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)
@@ -378,7 +404,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
 endif()
 if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
     add_subdirectory(src/core/vulkan)
-    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate)
+    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate sd_ggml_register_vulkan_rms_affine)
 endif()
 if(SD_BACKEND_REGISTRATIONS)
     string(REPLACE ";" ",\n    " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")

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

@@ -59,6 +59,7 @@ API and command-line option may change frequently.***
     - [Z-Image](./docs/z_image.md)
     - [Z-Image L2P](./docs/z_image_l2p.md)
     - [MiniT2I](./docs/minit2i.md)
+    - [P-H3-Image](./docs/p_h3_image.md)
     - [SenseNova U1.5](./docs/sensenova_u1.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
@@ -87,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
@@ -197,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,8 +3,11 @@
 
 #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);
 void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device);
 
 static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {

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

@@ -89,8 +89,13 @@ again after reordering.
 
 External inputs stay allocated through the final node. Native graph optimization
 and execution operate on the remaining contiguous ranges, so native fusion cannot
-consume nodes selected by an extension. Intermediate allocations are retained
-in this initial implementation; only their computation is removed.
+consume nodes selected by an extension. Fused intermediates are excluded from
+allocation through the backend's `skip_alloc` callback and
+`ggml_gallocr_set_skipped_tensors`. Their metadata and graph edges remain available,
+but execution must not access their data or buffers. Storage backing the final
+output is retained, including an internal owner used by an in-place final write.
+The skipped set is replaced for each allocation graph; changing which tensors are skipped
+invalidates cached allocation assignments even when tensor shapes are unchanged.
 
 Automatic extension fusion is enabled in scheduler execution with one pipeline
 copy and without evaluation callbacks. Multiple pipeline copies and evaluation
@@ -104,6 +109,16 @@ 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_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
+retain the producer's op marker. W4A4 uses row scales; W4A8 additionally exposes
+the codebook, channel scales, and FP8 relative scales as tensor dependencies.
+CPU, CUDA/HIP/MUSA, and Vulkan register kernels through the same extension API
+as INT8. No ggml submodule changes are required.
+
 Services are discovered through `ggml_backend_reg_get_proc_address`:
 
 | Service | Purpose |
@@ -115,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
@@ -127,14 +145,48 @@ Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; caller
 must use the representation appropriate to their backend. Vulkan extensions use
 the command recording services described below.
 
+CUDA INT8 Linear, RoPE, ConvRot, RMS affine and INT8 gated residual kernels omit output
+scratch when allocation guarantees that a fresh owning output cannot overlap
+live inputs. Existing outputs and views retain overlap protection. Internal
+GEMM and quantization workspace is still reserved when required; kernels also
+handle a null workspace when the entire step needs no scratch.
+
+Compute workspace measurement preserves the storage identities and offsets of
+external tensors, including views. Separate allocations must not share a synthetic
+buffer identity: that can invent overlap, overestimate scratch, and repeatedly
+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
@@ -192,6 +244,10 @@ ConvRot output fusion. These kernels use only the public extension API, borrowed
 CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or
 private conversion helpers. The existing CUDA fusion and CuTe disable environment
 variables are honored.
+Aligned contiguous FP16 K/V inputs are read directly without reserving conversion
+buffers. Fresh F32 outputs are written by the attention epilogue, retaining the
+original FP16 rounding before output scaling. Packed ConvRot and preallocated or
+view outputs retain the intermediate half buffer.
 
 `src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues:
 `GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy,
@@ -213,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
@@ -228,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
@@ -252,12 +308,21 @@ small-row policy, Stream-K scheduling, Ada column bands, and separate scaling
 order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row
 concatenation for custom GEMM nodes. Weight scales remain explicit
 inputs, including the gate/up scale permutation recorded on every SwiGLU launch.
+Two independent INT8 GEMMs sharing a quantized input can also fuse with split
+SwiGLU. A paired-column iterator reads the separate gate/up weight allocations,
+and the GEMM epilogue writes only the activated output. Both intermediate
+activations are omitted through the generic fusion allocation interface; weights
+are neither copied nor concatenated. This path requires bias-free branches,
+compatible layouts and alignments, and no external intermediate consumers.
+Only interleaved weight scales and any Stream-K workspace remain temporary.
 Plans reserve bounded Stream-K workspace, cache occupancy on first execution,
 and initialize barriers on every launch, so CUDA Graph replay observes updated
 inputs and scales. External consumers retain ordinary execution; overlapping
-outputs use workspace. The CUTLASS SwiGLU/ConvRot chain can write its packed
-output directly after GEMM, because quantization reads only its workspace
-intermediate. Shapes or alignments unsuitable for CUTLASS use cuBLAS.
+outputs use workspace. For fresh outputs, the CUTLASS SwiGLU epilogue writes
+an allocator-managed tensor that the subsequent ConvRot launch consumes. This
+keeps the same GPU launches while allowing its storage to be reused by other
+operations. The combined SwiGLU/ConvRot path retains its workspace intermediate
+for other cases. Shapes or alignments unsuitable for CUTLASS use cuBLAS.
 Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to
 preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
 settings are honored. No CUTLASS headers are needed for the cuBLAS path.
@@ -433,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.
 
@@ -452,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.
@@ -100,7 +101,68 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 
 LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
 
-The shared CUDA/HIP/MUSA convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution.
+The shared CUDA/HIP/MUSA convrot activation paths support H64 and H256. Vulkan also supports H64 activation quantization; its INT8 matrix multiplication path currently requires H256.
+
+## Packed INT4 convrot
+
+ComfyUI `convrot_w4a4` and `asym_w4a8_int8` safetensors load directly, including
+checkpoints that mix packed INT4, INT8 tensorwise, and floating-point layers.
+Packed weights retain their checkpoint byte layout; they are not expanded or
+converted to GGUF Q4 blocks at load time.
+
+- `convrot_w4a4`: signed 4-bit weights, low nibble first, with F32 per-output
+  `weight_scale`. Activations use the same INT8 convrot quantization as above;
+  the format name does not select a native INT4 activation kernel.
+- `asym_w4a8_int8`: 4-bit indices into `weight_codebook` (16 F32 values), F32
+  `weight_s_channel`, and F8_E4M3 `weight_s_rel` per 16 input features. The
+  metadata must specify `group_size: 16`. CPU evaluates the floating-point
+  codebook; GPU kernels requantize it to INT8 with `amax / 127`, introducing a
+  small additional rounding error.
+
+Both formats imply convrot rotation even without a `convrot` flag; missing
+`convrot_groupsize` defaults to 256.
+
+Both formats reuse the convrot activation cache when the input scale is `1` and
+produce F32 outputs with optional bias. CPU provides a portable implementation.
+CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS.
+With CUTLASS enabled, NVIDIA
+SM80 and newer devices use mixed-input S8 x S4 GEMM for eligible W4A4 shapes,
+reusing the INT8 scale epilogue without expanding the full weight tensor. W4A8
+uses CUTLASS integer MMA per 16-value group with FP32 scale restoration.
+Other CUDA shapes and builds retain the packed byte-dot path. Set
+`GGML_CUDA_W4_FORCE_DP4A=1` to compare against that path. Supported AMD targets
+use packed byte-dot instructions, with portable byte arithmetic for other targets.
+Vulkan requires integer dot product support and at least 13312 bytes of shared
+memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input
+width to be divisible by 64; H64 and H256 are supported. When supported by the
+compiler and device, Vulkan uses NV cooperative matrix integer GEMM. W4A4
+decodes packed weights directly into cooperative matrices for aligned input
+and output widths. Complete activation tiles and the final partial tile use
+separate shader specializations. Other W4A4 shapes use double-buffered weight
+and activation tiles (20480 bytes of shared memory, or 30720 for larger tiles).
+W4A8 stages both weights and activations in shared memory in K=64 tiles,
+reusing them for each 16-value group's integer
+MMA and FP32 scale restoration. FP8 scales use a shared lookup table, and
+expanding projections visit activation tiles in groups of weight panels for
+cache reuse. Adjacent groups share a K=32 activation fragment, with
+complementary zero padding in the weights to preserve each group's scale.
+This path uses 17476 bytes of shared memory, or 22596 for larger tiles.
+Both paths use 256-thread workgroups and keep the model's
+weights packed; neither expands a full layer in device memory. Other Vulkan
+devices retain the packed dot-product path; set `GGML_VK_W4_FORCE_DP4A=1` to compare
+against it. FP32 accumulation order can differ between these paths.
+Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA by the active
+compilation target's intrinsic support, preferring the GFX12 fragment layout
+over the GFX11 layout.
+There is no GPU architecture allowlist: newer targets supporting either
+intrinsic automatically use its matching layout. Selection happens per device
+image, including multi-architecture builds; targets supporting neither intrinsic
+retain the integer-dot fallback. W4A4 retains integer dot products on HIP.
+
+LoRA application remains at runtime, using the original unrotated activations
+and logical weight dimensions. Packed weights and their codebook/scale tensors
+are excluded from weight type conversion. `SD_USE_UPSTREAM_GGML` builds reject
+these formats because they require the application extension API.
 
 ## Example
 

docs/minimax_h3.md

@@ -99,10 +99,88 @@ The C API exposes the same inputs through `ref_images`, `ref_videos`, and
 `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
 frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
 
+## Latent hires second pass (two-stage sampling)
+
+MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers:
+
+- an interpolating latent upscaler (`Latent`, `Latent (nearest)`,
+  `Latent (bicubic)`, ...): after the first pass finishes, the packed AV latent
+  is resampled in latent space to `--hires-width` x `--hires-height` and
+  refined by a second sampling pass;
+- a model-backed upscaler (`--hires-upscaler <name>` + `--hires-upscalers-dir`):
+  a H3-shaped learned latent upscaler checkpoint runs between the passes,
+  like the LTX spatial upsampler. This is the sd.cpp equivalent of the
+  two-sampler workflows used for H3 in node-based frontends (e.g. the
+  community 3D-conv latent upscaler from
+  [LBH-123-AI/Minimax_h3_latent_Upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler)).
+
+Draft at a low megapixel count (e.g. 0.4 MP) and refine at a higher one
+(e.g. 0.8 MP) instead of paying for the full resolution in every step:
+
+```sh
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "a lovely cat" --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --video-frames 56 --hires --hires-upscaler "Latent (bicubic)" --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
+```
+
+or with the learned upscaler:
+
+```sh
+.\bin\Release\sd-cli.exe -M vid_gen ... --hires --hires-upscalers-dir ..\models\latent_upscale_models --hires-upscaler h3-latent-upscaler-3d --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
+```
+
+Notes for this mode:
+
+- The learned upscaler expects a second normalization of the diffusion latent,
+  matching its training convention. The runner applies and reverses this extra
+  transform internally; callers should pass the usual H3 diffusion latents.
+- The frame count is never changed by the second stage, so the packed audio
+  latent length stays valid and the generated soundtrack is preserved.
+- Width and height still have to be multiples of 32; the hires target is
+  aligned like any other resolution. The learned upscaler additionally wants
+  an even latent size on both axes, which 32-pixel targets provide.
+- `--hires-sigmas` can replace `--hires-denoising-strength` when the second
+  stage should resume from an explicit sigma schedule, which is how most
+  two-stage H3 workflows split the sampler.
+- Keyframe conditioning (`--init-img` / `--end-img`) is not available together
+  with hires yet: those latents are encoded at the first stage resolution and
+  would have to be re-encoded for the hires canvas. Use T2VA or Ref2VA.
+- The second stage runs at the larger latent, so it is the memory-hungry half;
+  combine it with quantization, `--vae-tiling` and `--offload-to-cpu` on
+  small cards.
+- The learned upscaler is scale-conditional: the target-to-source ratio is
+  fed into the network as an embedding, 1.2-2.0x is the trained domain, and
+  downscaling is rejected.
+
 ## Shape and runtime notes
 
 - Width and height are aligned upward to a multiple of 32.
 - Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5.
 - MiniMax-H3 runs at 24 fps; another requested value is overridden.
+- Attention head grouping is controlled by
+  `--model-args minimax_h3_attn_chunk_budget=2`, in GiB (default: `2`).
+  The budget estimates the size of a single FP32 Q/K tensor per group:
+  `head_dim * tokens * batch * 4 bytes * heads_per_group`. The group size adapts
+  to the sequence length, with a minimum of one head; inputs that fit use all
+  heads together. Set `0` to disable grouping, or a smaller budget such as
+  `0.25` to reduce peak activation memory for long sequences at the cost of
+  extra launches and concatenation. Each head still attends to the full sequence.
+  This is a grouping target, not a cap on total attention memory or VRAM usage;
+  other intermediates, backend workspace, and weights require additional memory.
+  `--max-vram` and `--offload-to-cpu` control weight residency, but the active
+  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/model_config.md

@@ -60,6 +60,8 @@ static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage
 ```
 
 Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes.
+Use `logical_ne0()` for Linear input widths; packed INT4 weights store half of
+the logical width in `ne[0]`. Keep raw `ne` values for storage layout calculations.
 Do not load or parse tensor data for config detection.
 
 Detection should respect `prefix`. For nested weights, construct full names from

docs/p_h3_image.md

@@ -0,0 +1,34 @@
+# How to Use
+
+P-H3-Image supports text-to-image and reference-image editing. It uses a MiniMax-H3 diffusion transformer, a single-frame VAE, and Qwen3-VL-4B with a text connector. Prompt enhancement is not included.
+
+## Download weights
+
+- Download P-H3-Image
+    - safetensors: https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/diffusion_models
+- Download P-H3-Image VAE
+    - safetensors: https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/vae
+- Download Qwen3-VL-4B with the text connector
+    - https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/text_encoders
+
+Alternatively, reuse Qwen3-VL-4B, including its vision tower, with the [separate text connector](https://huggingface.co/PrunaAI/P-H3-Image/blob/main/text_connector/diffusion_pytorch_model.safetensors). Save the connector as `p-h3-image_text_connector.safetensors` and pass it with `--embeddings-connectors`.
+
+Use the P-H3 image VAE; the MiniMax-H3 video VAE is not compatible. No audio VAE is needed.
+
+## Examples
+
+### Text-to-image
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\p-h3-image_int8_convrot.safetensors --llm ..\models\text_encoders\qwen3vl_4b_bf16.safetensors --embeddings-connectors ..\models\text_encoders\p-h3-image_text_connector.safetensors --vae ..\models\vae\p-h3-image_vae_bf16.safetensors --cfg-scale 1.0 --steps 8 --sampling-method euler --sigmas p-h3-image -p "a lovely cat holding a sign says 'p-h3-image.cpp'" -v --offload-to-cpu --rng cpu -H 2912 -W 2912
+```
+
+### Image editing
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\p-h3-image_int8_convrot.safetensors --llm ..\models\text_encoders\qwen3vl_4b_bf16.safetensors --embeddings-connectors ..\models\text_encoders\p-h3-image_text_connector.safetensors --vae ..\models\vae\p-h3-image_vae_bf16.safetensors --cfg-scale 1.0 --steps 8 --sampling-method euler --sigmas p-h3-image -r source.png -r reference.png -p "Edit <Picture 1>. Give the cat the scarf from <Picture 2>, preserving its face and the background." -v --offload-to-cpu --rng cpu -H 2912 -W 2912
+```
+
+Reference images are numbered in argument order as `<Picture 1>`, `<Picture 2>`, and so on. A plain edit instruction gets `Edit <Picture 1>. ` prepended automatically. Prompts starting with `Edit <Picture N>` or `subject_definitions:` are used directly.
+
+Use dimensions divisible by 32. The upstream default is 2912 × 2912; smaller canvases reduce memory use. Use CFG 1 and `--sigmas p-h3-image` for the upstream 8-step schedule. Keep the default VAE tiling size of 256 pixels.

docs/relay_vsr.md

@@ -0,0 +1,50 @@
+# RelayVSR image and video super-resolution
+
+[RelayVSR](https://github.com/kopperx/RelayVSR) upscales images and videos by
+2x or 4x with `-M vid_upscale`.
+
+## Weights
+
+[ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/RelayVSR) load
+directly. Select the matching pair for the requested scale:
+
+| `--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 \
+  --vae models/vae/relayvsr_flash_vae_x2_fp32.safetensors \
+  --input-frames input_frames --upscale-factor 2 --fps 24 \
+  --vsr-window-size 16 -s 42 -o output.avi
+```
+
+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`.
+
+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.
+
+## Settings and limits
+
+- `--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

@@ -3,6 +3,7 @@
 #include <time.h>
 #include <algorithm>
 #include <cctype>
+#include <climits>
 #include <filesystem>
 #include <functional>
 #include <iostream>
@@ -33,10 +34,14 @@ const char* previews_str[] = {
 std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
 
 struct SDCliParams {
-    SDMode mode             = IMG_GEN;
-    std::string output_path = "output.png";
-    int output_begin_idx    = -1;
-    int compression_quality = 90;
+    std::string input_frames;
+    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";
 
@@ -62,6 +67,7 @@ struct SDCliParams {
         ArgOptions options;
 
         options.string_options = {
+            {"", "--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",
@@ -90,6 +96,9 @@ struct SDCliParams {
         };
 
         options.int_options = {
+            {"", "--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",
@@ -190,7 +199,7 @@ struct SDCliParams {
         options.manual_options = {
             {"-M",
              "--mode",
-             "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
+             "run mode, one of [img_gen, adetailer, vid_gen, vid_upscale, upscale, convert, metadata], default: img_gen",
              on_mode_arg},
             {"",
              "--preview",
@@ -220,6 +229,10 @@ struct SDCliParams {
     }
 
     bool validate() {
+        if (mode == VID_UPSCALE && video_upscale_factor <= 0) {
+            LOG_ERROR("vid_upscale requires a positive --upscale-factor");
+            return false;
+        }
         if (mode != METADATA) {
             if (output_path.length() == 0) {
                 LOG_ERROR("error: the following arguments are required: output_path");
@@ -529,7 +542,9 @@ bool save_results(const SDCliParams& cli_params,
         return sucessful_reults != 0;
     }
 
-    if (cli_params.mode == VID_GEN && num_results > 1) {
+    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
@@ -664,6 +679,72 @@ int main(int argc, const char* argv[]) {
         return 0;
     }
 
+    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 (!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;
+        }
+        if (owned_frames.size() > size_t(INT_MAX))
+            return 1;
+        std::vector<sd_image_t> frames;
+        for (const auto& frame : owned_frames)
+            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.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);
+        if (!upscaler)
+            return 1;
+        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.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;
+        SDImageVec results;
+        results.adopt_video(output);
+        return save_results(cli_params, ctx_params, gen_params, results, int(results.size()), nullptr, video_info) ? 0 : 1;
+    }
+
     if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
         LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
         return 1;

examples/common/common.cpp

@@ -35,6 +35,7 @@ const char* const modes_str[] = {
     "convert",
     "upscale",
     "metadata",
+    "vid_upscale",
 };
 
 static sd_vae_format_t str_to_vae_format(const std::string& value) {
@@ -464,7 +465,7 @@ ArgOptions SDContextParams::get_options() {
          &uncond_diffusion_model_path},
         {"",
          "--embeddings-connectors",
-         "path to LTXAV embeddings connectors",
+         "path to embeddings connectors (LTXAV or P-H3-Image)",
          0,
          &embeddings_connectors_path},
         {"",
@@ -1135,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},
         {"",
@@ -1187,6 +1188,13 @@ ArgOptions SDGenerationParams::get_options() {
          "under --hires-upscalers-dir (default: Latent)",
          0,
          &hires_upscaler},
+        {"",
+         "--control-type",
+         "kind of the control image, one of [openpose, depth, scribble, canny, normal, segment, tile, repaint]. "
+         "Required by an SDXL Control Net Union model (tile and repaint only exist in its ProMax variant), ignored by any other Control Net. "
+         "Also accepted: hed, softedge, pidi, ted (scribble); lineart, anime_lineart, mlsd (canny); seg (segment); inpaint, outpaint (repaint)",
+         0,
+         &control_type},
         {"",
          "--extra-sample-args",
          "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
@@ -1530,6 +1538,14 @@ ArgOptions SDGenerationParams::get_options() {
         std::string sigmas_str = value;
         if (sigmas_str == "qwen-image-2.1-turbo") {
             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);
@@ -1777,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.",
+         "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.",
+         "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",
@@ -2247,6 +2263,7 @@ bool SDGenerationParams::from_json_str(
 
     load_if_exists("strength", strength);
     load_if_exists("control_strength", control_strength);
+    load_if_exists("control_type", control_type);
     load_if_exists("ip_adapter_strength", ip_adapter_strength);
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
@@ -2636,6 +2653,15 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
         sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
     }
 
+    resolved_control_type = SD_CONTROL_TYPE_NONE;
+    if (!control_type.empty()) {
+        resolved_control_type = str_to_sd_control_type(control_type.c_str());
+        if (resolved_control_type == SD_CONTROL_TYPE_COUNT) {
+            LOG_ERROR("invalid control type '%s' (openpose, depth, scribble, canny, normal, segment, tile, repaint)", control_type.c_str());
+            return false;
+        }
+    }
+
     hires_upscaler_model_path.clear();
     if (hires_enabled) {
         if (hires_upscaler.empty()) {
@@ -2875,6 +2901,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.qwen_image_layers   = qwen_image_layers;
     params.control_image       = control_image.get();
     params.control_strength    = control_strength;
+    params.control_type        = resolved_control_type;
     params.ip_adapter_image    = ip_adapter_image.get();
     params.ip_adapter_strength = ip_adapter_strength;
     params.pm_params           = pm_params;
@@ -3080,6 +3107,7 @@ std::string SDGenerationParams::to_string() const {
         << "  vace_strength: " << vace_strength << ",\n"
         << "  strength: " << strength << ",\n"
         << "  control_strength: " << control_strength << ",\n"
+        << "  control_type: \"" << control_type << "\",\n"
         << "  seed: " << seed << ",\n"
         << "  upscale_repeats: " << upscale_repeats << ",\n"
         << "  upscale_tile_size: " << upscale_tile_size << ",\n"

examples/common/common.h

@@ -22,7 +22,7 @@ std::string argv_to_utf8(int index, const char** argv);
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
 extern const char* const modes_str[];
-#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
+#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata, vid_upscale"
 
 enum SDMode {
     IMG_GEN,
@@ -31,6 +31,7 @@ enum SDMode {
     CONVERT,
     UPSCALE,
     METADATA,
+    VID_UPSCALE,
     MODE_COUNT
 };
 
@@ -227,6 +228,7 @@ struct SDGenerationParams {
     std::string end_image_path;
     std::string mask_image_path;
     std::string control_image_path;
+    std::string control_type;
     std::string ip_adapter_image_path;
     std::vector<std::string> ref_image_paths;
     std::vector<std::string> ref_video_paths;
@@ -291,6 +293,7 @@ struct SDGenerationParams {
     std::string prompt_with_lora;  // for metadata record only
     std::vector<sd_lora_t> lora_vec;
     sd_hires_upscaler_t resolved_hires_upscaler;
+    sd_control_type_t resolved_control_type = SD_CONTROL_TYPE_NONE;
 
     // Owned execution payload.
     SDImageOwner init_image;

examples/common/generation_metadata.cpp

@@ -1,5 +1,6 @@
 #include "generation_metadata.h"
 
+#include <cctype>
 #include <filesystem>
 #include <map>
 
@@ -70,6 +71,7 @@ std::string generation_info_to_webui(const std::string& text) {
     const auto& sampling = info.at("sample_params");
     std::string result   = info.value("prompt", "");
     if (info.contains("loras")) {
+        bool prompt_trimmed = false;
         for (const auto& lora : info.at("loras")) {
             // WebUI tags cannot restrict a LoRA to the high-noise model.
             if (lora.value("is_high_noise", false))
@@ -83,8 +85,15 @@ std::string generation_info_to_webui(const std::string& text) {
                 name.erase(extension);
             if (name.empty() || name.find_first_of(":<>\r\n") != std::string::npos)
                 continue;
-            if (!result.empty())
-                result += " ";
+            if (!result.empty()) {
+                if (!prompt_trimmed) {
+                    while (!result.empty() && std::isspace(static_cast<unsigned char>(result.back())))
+                        result.pop_back();
+                    prompt_trimmed = true;
+                }
+                if (!result.empty())
+                    result += " ";
+            }
             result += "<lora:" + name + ":" + lora.at("multiplier").dump() + ">";
         }
     }

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 |
@@ -936,6 +988,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH
 | `ref_image_args` | `string` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `hires` | `object` |
 | `hires.enabled` | `boolean` |
@@ -1116,6 +1169,7 @@ Example:
   "ref_image_args": "",
   "increase_ref_index": false,
   "control_strength": 0.9,
+  "control_type": "",
   "ip_adapter_strength": 1.0,
   "embed_image_metadata": true,
 
@@ -1241,6 +1295,7 @@ Top-level scalar fields:
 | `image_preprocess` | `string \| array<string>` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
 | `preview` | `string` |

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. |
@@ -71,6 +73,7 @@ For community edit adapters and other optional extensions, these are the base-mo
 | `minit2i` | 2.5 | 20 | 512 × 512 | [Source](https://huggingface.co/MiniT2I/MiniT2I/raw/main/pipeline.py) B/16 example and native model size; L/16 uses CFG 6. |
 | `moss-image-base` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base)  |
 | `moss-image-sft` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)  |
+| `p-h3-image` | 1 | 8 | 1024 × 1024 | [Runtime guide](../../../docs/p_h3_image.md) Fixed P-H3 sigma schedule; upstream 2912 × 2912 reduced to the catalog default. |
 | `nitrosd-realism` | 1 | 1 | 1024 × 1024 | [Source](https://huggingface.co/ChenDY/NitroFusion)  |
 | `ovis-image` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/ATH-MaaS/Ovis-Image-7B)  |
 | `pid` | 1 | 4 | 2048 × 2048 | [Source](https://huggingface.co/nvidia/PiD) Standalone PiD decoding at 2K; guidance is handled by the upstream generator. |

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"
+  ]
+}

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

@@ -0,0 +1,55 @@
+{
+  "schema_version": 1,
+  "id": "p-h3-image",
+  "name": "P-H3-Image",
+  "size_multiple": 32,
+  "description": "Supports text-to-image and reference-image editing with the fixed 8-step schedule. Use the P-H3 single-frame VAE and Qwen3-VL-4B with its vision tower. Supply the P-H3 text connector separately unless it is included in the text encoder.",
+  "model_url": "https://huggingface.co/PrunaAI/P-H3-Image",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "Qwen 3 VL 4B Text Encoder"
+    },
+    {
+      "id": "llm_vision",
+      "required": false,
+      "label": "Vision projector (required when not included in the text encoder)"
+    },
+    {
+      "id": "embeddings_connectors",
+      "required": false,
+      "label": "P-H3 Text Connector (required when not included in the text encoder)"
+    },
+    {
+      "id": "vae",
+      "required": true,
+      "label": "P-H3 Image VAE"
+    }
+  ],
+  "defaults": {
+    "sampling_method": "euler",
+    "cfg_scale": 1,
+    "steps": 8,
+    "width": 1024,
+    "height": 1024,
+    "custom_sigmas": [
+      1.0,
+      0.984753,
+      0.967905,
+      0.951688,
+      0.925372,
+      0.881923,
+      0.810083,
+      0.528718,
+      0.0
+    ]
+  },
+  "features": [
+    "ref_images"
+  ]
+}

examples/server/routes_sdcpp.cpp

@@ -154,6 +154,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"image_preprocess", defaults.image_preprocess},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
+        {"control_type", defaults.control_type},
         {"ip_adapter_strength", defaults.ip_adapter_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers, defaults.custom_sigmas)},
         {"hires", make_hires_json(defaults)},

include/stable-diffusion.h

@@ -400,6 +400,19 @@ enum sd_hires_upscaler_t {
     SD_HIRES_UPSCALER_COUNT,
 };
 
+enum sd_control_type_t {
+    SD_CONTROL_TYPE_NONE,
+    SD_CONTROL_TYPE_OPENPOSE,
+    SD_CONTROL_TYPE_DEPTH,
+    SD_CONTROL_TYPE_SCRIBBLE,  // thick lines: scribble, hed, softedge, ted
+    SD_CONTROL_TYPE_CANNY,     // thin lines: canny, lineart, anime lineart, mlsd
+    SD_CONTROL_TYPE_NORMAL,
+    SD_CONTROL_TYPE_SEGMENT,
+    SD_CONTROL_TYPE_TILE,
+    SD_CONTROL_TYPE_REPAINT,  // inpaint, outpaint
+    SD_CONTROL_TYPE_COUNT,
+};
+
 typedef struct {
     bool enabled;
     enum sd_hires_upscaler_t upscaler;
@@ -444,6 +457,7 @@ typedef struct {
     bool circular_x;
     bool circular_y;
     sd_image_preprocess_params_t image_preprocess;
+    enum sd_control_type_t control_type;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -486,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 {
@@ -496,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).
@@ -531,6 +559,10 @@ SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
 SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
 SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
+SD_API const char* sd_control_type_name(enum sd_control_type_t control_type);
+// Also accepts the names of the kinds a type stands for (hed, lineart, mlsd, ...);
+// returns SD_CONTROL_TYPE_COUNT when the name is unknown.
+SD_API enum sd_control_type_t str_to_sd_control_type(const char* str);
 
 SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
 SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
@@ -597,6 +629,42 @@ SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
 // Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model.
 SD_API int get_upscaler_model_scale(const char* model_path);
 
+typedef struct video_upscaler_ctx_t video_upscaler_ctx_t;
+
+typedef struct {
+    const char* diffusion_model_path;
+    const char* vae_path;
+    // Explicit output scale; RelayVSR requires 2 or 4 matching both checkpoints.
+    int upscale_factor;
+    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 {
+    const sd_image_t* frames;
+    int frame_count;
+    int fps;
+    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);
+SD_API void sd_video_upscale_params_init(sd_video_upscale_params_t* params);
+SD_API video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params);
+SD_API void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx);
+// Frames are RGB, have equal dimensions, and remain owned by the caller.
+// Free the result with free_sd_video_generation_result().
+SD_API bool upscale_video(video_upscaler_ctx_t* ctx,
+                          const sd_video_upscale_params_t* params,
+                          sd_video_generation_result_t** result_out);
+
 typedef struct adetailer_ctx_t adetailer_ctx_t;
 
 typedef struct {

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

@@ -17,6 +17,7 @@
 #include "model/te/llada_image_te.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/ming_image_te.hpp"
+#include "model/te/ph3_image_connector.hpp"
 #include "model/te/t5.hpp"
 #include "model_loader.h"
 #include "tokenizers/sensenova_u1_tokenizer.h"
@@ -1976,6 +1977,7 @@ struct LLMEmbedder : public Conditioner {
     std::shared_ptr<Tokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<T5Runner> byt5;
+    std::shared_ptr<PH3ImageTextConnectorRunner> ph3_connector;
 
     LLMEmbedder(ggml_backend_t backend,
                 const String2TensorStorage& tensor_storage_map      = {},
@@ -2022,12 +2024,18 @@ struct LLMEmbedder : public Conditioner {
         } else if (version == VERSION_MOSS_IMAGE) {
             arch = LLM::LLMArch::QWEN3_5;
         }
-        llm        = std::make_shared<LLM::LLMRunner>(arch,
-                                                      backend,
-                                                      tensor_storage_map,
-                                                      "text_encoders.llm",
-                                                      enable_vision,
-                                                      weight_manager);
+        llm = std::make_shared<LLM::LLMRunner>(arch,
+                                               backend,
+                                               tensor_storage_map,
+                                               "text_encoders.llm",
+                                               enable_vision,
+                                               weight_manager);
+        if (version == VERSION_P_H3_IMAGE) {
+            if (llm->config.num_layers != 36 || llm->config.hidden_size != 2560) {
+                throw std::runtime_error("P-H3-Image requires the Qwen3-VL-4B text encoder");
+            }
+            ph3_connector = std::make_shared<PH3ImageTextConnectorRunner>(backend, tensor_storage_map, weight_manager);
+        }
         int pad_id = 151643;
         if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
             pad_id = 11;
@@ -2063,6 +2071,9 @@ struct LLMEmbedder : public Conditioner {
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
+        if (ph3_connector) {
+            ph3_connector->get_param_tensors(tensors);
+        }
         if (byt5) {
             byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
         }
@@ -2074,6 +2085,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
+        if (ph3_connector) {
+            ph3_connector->set_max_graph_vram_bytes(max_vram_bytes);
+        }
         if (byt5) {
             byt5->set_max_graph_vram_bytes(max_vram_bytes);
         }
@@ -2081,6 +2095,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         llm->set_runtime_backends(backends);
+        if (ph3_connector) {
+            ph3_connector->set_runtime_backends(backends);
+        }
         if (byt5) {
             byt5->set_runtime_backends(backends);
         }
@@ -2113,6 +2130,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
+        if (ph3_connector) {
+            ph3_connector->set_flash_attention_enabled(enabled);
+        }
         if (byt5) {
             byt5->set_flash_attention_enabled(enabled);
         }
@@ -2120,12 +2140,18 @@ struct LLMEmbedder : public Conditioner {
 
     void set_scale_overrides(float linear_scale, float attn_scale) override {
         llm->set_scale_overrides(linear_scale, attn_scale);
+        if (ph3_connector) {
+            ph3_connector->set_scale_overrides(linear_scale, attn_scale);
+        }
         if (byt5) {
             byt5->set_scale_overrides(linear_scale, attn_scale);
         }
     }
 
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        if (ph3_connector) {
+            ph3_connector->set_weight_adapter(adapter);
+        }
         if (llm) {
             llm->set_weight_adapter(adapter);
         }
@@ -2135,6 +2161,9 @@ struct LLMEmbedder : public Conditioner {
     }
 
     void runner_end() override {
+        if (ph3_connector) {
+            ph3_connector->runner_end();
+        }
         if (llm) {
             llm->runner_end();
         }
@@ -2311,6 +2340,12 @@ struct LLMEmbedder : public Conditioner {
             LOG_ERROR("LLM prompt encoding failed");
             return {};
         }
+        if (ph3_connector) {
+            hidden_states = ph3_connector->compute(n_threads, hidden_states);
+            if (hidden_states.empty()) {
+                return {};
+            }
+        }
         hidden_states = apply_token_weights(std::move(hidden_states), weights);
         GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
 
@@ -2457,8 +2492,9 @@ struct LLMEmbedder : public Conditioner {
 
         if (sd_version_is_minimax_h3(version)) {
             prompt_template_encode_start_idx = 0;
-            out_layers                       = {50};
-            prompt_attn_range                = {0, 0};
+            // The final tap is normalized; LLMRunner exposes it at num_layers + 1.
+            out_layers        = version == VERSION_P_H3_IMAGE ? std::set<int>{13, 19, 25, 31, 37} : std::set<int>{50};
+            prompt_attn_range = {0, 0};
 
             if (llm->enable_vision) {
                 const int patch_size = llm->config.vision.patch_size;
@@ -2564,7 +2600,17 @@ struct LLMEmbedder : public Conditioner {
                     }
                 }
             }
-            prompt += conditioner_params.text;
+            auto text = conditioner_params.text;
+            if (version == VERSION_P_H3_IMAGE && !image_embeds.empty()) {
+                // Preserve explicit picture targets and structured briefs from the upstream upsampler.
+                static const std::regex verbatim(R"(^\s*(Edit <Picture \d+>|subject_definitions:))");
+                if (!std::regex_search(text, verbatim)) {
+                    text = "Edit <Picture 1>. " + text;
+                } else {
+                    text = trim(text);
+                }
+            }
+            prompt += text;
         } else if (sd_version_is_hunyuan_video(version)) {
             prompt_template_encode_start_idx = 98;
             out_layers                       = {26};
@@ -3402,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/compute_workspace.cpp

@@ -147,8 +147,8 @@ namespace sd {
                         append(type);
                     }
                 } else {
-                    const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
-                    append(alignment_offset);
+                    append(tensor->buffer);
+                    append(tensor->data);
                 }
                 const int32_t uses = ggml_graph_use_count(graph, tensor);
                 append(uses);
@@ -164,7 +164,7 @@ namespace sd {
             return {};
         }
         std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
-        std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
+        std::map<std::pair<ggml_backend_t, const void*>, ggml_backend_buffer_t> external_buffers;
         for (auto tensor : tensors) {
             auto copy      = ggml_dup_tensor(context, tensor);
             *copy          = *tensor;
@@ -179,7 +179,9 @@ namespace sd {
             }
             auto external = external_backend(source);
             if (external != nullptr && source->view_src == nullptr) {
-                auto& buffer = external_buffers[external];
+                // Preserve distinct external allocations so overlap checks do not invent aliases.
+                const void* storage = source->buffer ? static_cast<const void*>(source->buffer) : source;
+                auto& buffer        = external_buffers[{external, storage}];
                 if (buffer == nullptr) {
                     buffer = ggml_backend_alloc_buffer(external, 0);
                     GGML_ASSERT(buffer != nullptr);
@@ -191,6 +193,18 @@ namespace sd {
                 copy->extra = nullptr;
             }
         }
+        for (const auto& entry : copies) {
+            const auto* source = entry.first;
+            if (source->view_src && source->buffer) {
+                const auto* owner = source->view_src;
+                while (owner->view_src) {
+                    owner = owner->view_src;
+                }
+                if (external_backend(owner) != nullptr) {
+                    entry.second->buffer = copies.at(owner)->buffer;
+                }
+            }
+        }
         auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
         auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* {
             const auto& mapped = *static_cast<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);

src/core/cpu/registration.cpp

@@ -5,6 +5,7 @@ namespace sd_cpu {
 
     void register_extensions(ggml_backend_dev_t device) {
         register_i8(device);
+        register_w4(device);
         register_rope(device);
         register_modulate(device);
     }

src/core/cpu/registration.h

@@ -9,6 +9,7 @@ namespace sd_cpu {
     void register_rope(ggml_backend_dev_t device);
     void register_modulate(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
 }
 #endif
 

src/core/cpu/w4.cpp

@@ -0,0 +1,95 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cstring>
+
+namespace sd_cpu {
+    struct W4Execution {
+        SDW4Operation op;
+        float* output;
+        int64_t rows;
+    };
+
+    static void multiply_w4(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<W4Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.input->ne[0], n = op.weight->ne[1];
+        const auto* w       = static_cast<const uint8_t*>(op.weight->data);
+        const auto* x       = static_cast<const int8_t*>(op.input->data);
+        const auto* xs      = reinterpret_cast<const float*>(x + k * ((e.rows + 3) / 4 * 4));
+        const auto* ws      = static_cast<const float*>(op.weight_scale->data);
+        const auto* bias    = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        const auto* sc      = op.s_channel ? static_cast<const float*>(op.s_channel->data) : nullptr;
+        const auto* sr      = op.s_rel ? static_cast<const uint8_t*>(op.s_rel->data) : nullptr;
+        const auto to_float = ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float;
+        for (int64_t index = ith; index < n * e.rows; index += nth) {
+            const int64_t row = index / n, col = index % n;
+            const auto* packed = w + col * (k / 2);
+            float value;
+            if (op.kind == SDW4Kind::W4A4) {
+                uint32_t sum = 0;
+                for (int64_t i = 0; i < k; ++i) {
+                    const int nibble = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                    sum += uint32_t(((nibble ^ 8) - 8) * int32_t(x[row * k + i]));
+                }
+                int32_t signed_sum;
+                std::memcpy(&signed_sum, &sum, sizeof(sum));
+                value = float(signed_sum) * xs[row] * ws[col];
+            } else {
+                float sum = 0.f;
+                for (int64_t g = 0; g < k / 16; ++g) {
+                    float relative;
+                    to_float(sr + col * (k / 16) + g, &relative, 1);
+                    const float scale = sc[col] * relative;
+                    for (int j = 0; j < 16; ++j) {
+                        const int64_t i = g * 16 + j;
+                        const int code  = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                        sum += (ws[code] * scale) * float(x[row * k + i]);
+                    }
+                }
+                value = sum * xs[row];
+            }
+            e.output[index] = value + (bias ? bias[col] : 0.f);
+        }
+    }
+
+    static bool w4_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDW4Operation op;
+        return sd_ggml_w4_describe(node, &op);
+    }
+    static ggml_status w4_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status w4_execute(const ggml_ext_execution_context* ctx, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        W4Execution e = {};
+        if (!sd_ggml_w4_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->backend));
+        auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.rows      = ggml_nrows(dst);
+        e.output    = static_cast<float*>(ggml_backend_buffer_get_base(ctx->workspace));
+        auto status = run(ctx->backend, dst, multiply_w4, &e);
+        if (status == GGML_STATUS_SUCCESS)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    void register_w4(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.name                 = "sd.linear_w4_convrot.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = w4_supports;
+        d.prepare              = w4_prepare;
+        d.execute              = w4_execute;
+        d.destroy              = [](void*, void*) {};
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+}
+#endif

src/core/cuda/CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW)
 
 function(sd_cuda_sources VENDOR LANGUAGE)
     set(SOURCES)
-    foreach(OP rope int8 convrot affine modulate bf16 rms_affine registration)
+    foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration)
         set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
         string(TOUPPER "${VENDOR}" UPPER_VENDOR)
         file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
@@ -27,6 +27,12 @@ if(SD_CUDA)
     set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
     set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
     if(SD_CUDA_CUTLASS_PATH)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include")
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr)
+        if(MSVC)
+            set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+        endif()
         set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS)
         set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
         set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
@@ -49,6 +55,8 @@ if(SD_HIPBLAS)
         set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
     endif()
     set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
+    # W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12.
+    set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64")
     target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
     target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
 endif()

src/core/cuda/convrot.cu

@@ -204,7 +204,8 @@ namespace SD_CUDA_NAMESPACE {
             plan->device = cuda_device_info(ggml_backend_get_device(backend));
             if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
                 return GGML_STATUS_FAILED;
-            *workspace = ggml_nbytes(nodes[count - 1]);
+            const auto* dst = nodes[count - 1];
+            *workspace      = !dst->data && !dst->buffer && !dst->view_src ? 0 : ggml_nbytes(dst);
             return GGML_STATUS_SUCCESS;
         }
         static void destroy(void* state, void*) {

src/core/cuda/cute_attention.cu

@@ -150,7 +150,7 @@ namespace sd_fattn_cute {
         }
     }
 
-    template <int Rows>
+    template <int Rows, bool Masked = false>
     struct Softmax {
         decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
         template <bool First, typename Scores, typename Output>
@@ -180,8 +180,9 @@ namespace sd_fattn_cute {
                 }
 #pragma unroll
                 for (int n = 0; n < size<1>(scores); ++n) {
-                    const float value = exp2f(scores(m, n) * 1.4426950408889634f - next_max * 1.4426950408889634f);
-                    scores(m, n)      = value;
+                    const float safe_max = Masked && next_max == -INFINITY ? 0.f : next_max;
+                    const float value    = exp2f(scores(m, n) * 1.4426950408889634f - safe_max * 1.4426950408889634f);
+                    scores(m, n)         = value;
                     row_sum += value;
                 }
                 // Defer the four-lane sum reduction until the final output normalization.
@@ -196,7 +197,7 @@ namespace sd_fattn_cute {
                 float total = sum(m);
                 total += __shfl_xor_sync(0xffffffff, total, 2);
                 total += __shfl_xor_sync(0xffffffff, total, 1);
-                const float inverse = 1.f / total;
+                const float inverse = Masked && total == 0.f ? 0.f : 1.f / total;
 #pragma unroll
                 for (int n = 0; n < size<1>(output); ++n)
                     output(m, n) *= inverse;
@@ -236,6 +237,9 @@ namespace sd_fattn_cute {
         const cutlass::half_t *q, *k, *v;
         cutlass::half_t* o;
         int nq, nk, heads, kv_heads, dim;
+        float* output_f32  = nullptr;
+        float output_scale = 1.0f, output_bias = 0.0f;
+        bool scaled_output = false;
     };
 
     template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
@@ -332,18 +336,34 @@ namespace sd_fattn_cute {
         auto tOrO    = threadO.retile_S(rO);
         auto tOsO    = threadO.partition_D(sO);
         cute::copy(copyO, tOrO, tOsO);
-        auto mO = make_tensor(make_gmem_ptr(p.o + int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim),
-                              make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
-        auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
         typename T::OutputCopy output_copy;
         auto global_threadO = output_copy.get_thread_slice(tid);
         auto from_sO        = global_threadO.partition_S(sO);
-        auto to_gO          = global_threadO.partition_D(gO);
         auto output_id      = global_threadO.partition_D(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
         __syncthreads();
-        auto output_regs = make_tensor<Element>(shape(to_gO));
+        auto output_regs = make_tensor<Element>(shape(from_sO));
         cute::copy(output_copy, from_sO, output_regs);
-        copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+        const int64_t output_offset = int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim;
+        if (p.output_f32) {
+            auto float_copy = make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, float>{},
+                                              typename T::GlobalLayout{}, Layout<Shape<_1, _8>>{});
+            auto mO         = make_tensor(make_gmem_ptr(p.output_f32 + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+            auto gO         = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+            auto to_gO      = float_copy.get_thread_slice(tid).partition_D(gO);
+            auto values     = make_tensor<float>(shape(output_regs));
+#pragma unroll
+            for (int i = 0; i < size(values); ++i) {
+                // Preserve the half-rounded attention result before applying the F32 epilogue.
+                const float value = float(output_regs(i));
+                values(i)         = p.scaled_output ? __fmaf_rn(value, p.output_scale, p.output_bias) : value;
+            }
+            copy_rows<false, false, FullDim>(float_copy, values, to_gO, output_id, p.nq - m_block * BM, p.dim);
+        } else {
+            auto mO    = make_tensor(make_gmem_ptr(p.o + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+            auto gO    = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+            auto to_gO = global_threadO.partition_D(gO);
+            copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+        }
     }
 }
 
@@ -624,17 +644,34 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
 struct AttentionPlan {
     int index;
     size_t offsets[4];
+    bool direct_output;
 };
 
+static bool attention_direct_input(const ggml_tensor* src) {
+    size_t offset     = 0;
+    const auto* owner = src;
+    while (!owner->data && owner->view_src) {
+        offset += owner->view_offs;
+        owner = owner->view_src;
+    }
+    return src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && (uintptr_t(owner->data) + offset) % 16 == 0;
+}
+
 static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
     auto* plan  = new AttentionPlan{};
     *state      = plan;
     plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
                                                                                                       : 0;
     GGML_ASSERT(plan->index < count);
-    const auto* attn = nodes[plan->index];
-    size_t size      = 0;
+    const auto* attn    = nodes[plan->index];
+    const auto* dst     = nodes[count - 1];
+    plan->direct_output = dst->type == GGML_TYPE_F32 && !dst->data && !dst->buffer && !dst->view_src;
+    size_t size         = 0;
     for (int i = 0; i < 4; ++i) {
+        if ((i > 0 && i < 3 && attention_direct_input(attn->src[i])) || (i == 3 && plan->direct_output)) {
+            plan->offsets[i] = SIZE_MAX;
+            continue;
+        }
         plan->offsets[i]   = size;
         const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half);
         size += (bytes + 255) / 256 * 256;
@@ -680,7 +717,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     auto* base  = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
     half* work[4];
     for (int i = 0; i < 4; ++i)
-        work[i] = reinterpret_cast<half*>(base + plan.offsets[i]);
+        work[i] = plan.offsets[i] == SIZE_MAX ? nullptr : reinterpret_cast<half*>(base + plan.offsets[i]);
     const half* inputs[3] = {work[0], work[1], work[2]};
     for (int i = 0; i < 3; ++i) {
         const auto* src   = attn->src[i];
@@ -694,6 +731,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
         } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
             inputs[i] = static_cast<const half*>(src->data);
         } else {
+            if (!work[i])
+                return GGML_STATUS_FAILED;
             auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
             kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
                                                                           ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
@@ -701,12 +740,26 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
         if (cudaGetLastError() != cudaSuccess)
             return GGML_STATUS_FAILED;
     }
-    const auto* q           = attn->src[0];
-    const auto* k           = attn->src[1];
+    const auto* q = attn->src[0];
+    const auto* k = attn->src[1];
+    auto* dst     = nodes[count - 1];
+    float scale = 1.0f, bias = 0.0f;
+    for (int i = plan.index + 1; i < count; ++i) {
+        if (nodes[i]->op == GGML_OP_SCALE) {
+            scale = param_f32(nodes[i], 0);
+            bias  = param_f32(nodes[i], 1);
+        }
+    }
     sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
                                reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
                                reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
-    cudaError_t error       = cudaErrorInvalidValue;
+    if (plan.direct_output) {
+        p.output_f32    = static_cast<float*>(dst->data);
+        p.output_scale  = scale;
+        p.output_bias   = bias;
+        p.scaled_output = dst->op == GGML_OP_SCALE;
+    }
+    cudaError_t error = cudaErrorInvalidValue;
     switch (padded_dim(p.dim)) {
         case 64:
             error = attention_launch<64>(context, p, q->ne[3]);
@@ -732,14 +785,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     }
     if (error != cudaSuccess)
         return GGML_STATUS_FAILED;
-    float scale = 1.0f, bias = 0.0f;
-    for (int i = plan.index + 1; i < count; ++i) {
-        if (nodes[i]->op == GGML_OP_SCALE) {
-            scale = param_f32(nodes[i], 0);
-            bias  = param_f32(nodes[i], 1);
-        }
-    }
-    auto* dst = nodes[count - 1];
+    if (plan.direct_output)
+        return GGML_STATUS_SUCCESS;
     int convrot_group;
     if (sd_ggml_convrot_describe(dst, &convrot_group)) {
         const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4;
@@ -790,3 +837,5 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
     desc.priority  = 400;
     GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
 }
+
+#include "window_attention.cuh"

src/core/cuda/i8_gated_residual.cu

@@ -229,6 +229,7 @@ namespace sd_i8_gated {
         size_t kernel_workspace                                                     = 0;
         size_t output_offset                                                        = 0;
         int occupancy                                                               = 0;
+        bool output_scratch                                                         = true;
         ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
     };
     template <class Types>
@@ -397,24 +398,50 @@ namespace sd_i8_gated {
         return configure(describe(nodes, count), plan) ? count : 0;
 #endif
     }
-    static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
-        auto* plan   = new Plan;
-        *state       = plan;
-        plan->device = device_info(ggml_backend_get_device(backend));
-        if (!configure(describe(nodes, count), *plan))
-            return GGML_STATUS_FAILED;
-        *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
-        return GGML_STATUS_SUCCESS;
-    }
-    static void destroy(void* state, void*) {
-        delete static_cast<Plan*>(state);
+    static uintptr_t address(const ggml_tensor* tensor) {
+        size_t offset = 0;
+        while (tensor && !tensor->data && tensor->view_src) {
+            offset += tensor->view_offs;
+            tensor = tensor->view_src;
+        }
+        return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
     }
     static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
         if (!b)
             return false;
-        const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+        const auto x = address(a), y = address(b);
+        if (!x || !y)
+            return false;
         return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
     }
+    static bool needs_scratch(const ggml_tensor* dst, const Operation& op) {
+        if (address(dst) % 16 != 0)
+            return true;
+        for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
+            if (input == op.residual && !op.prefix && address(input) == address(dst))
+                continue;
+            if (overlaps(dst, input))
+                return true;
+        }
+        return false;
+    }
+    static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+        auto* plan    = new Plan;
+        *state        = plan;
+        plan->device  = device_info(ggml_backend_get_device(backend));
+        const auto op = describe(nodes, count);
+        if (!configure(op, *plan))
+            return GGML_STATUS_FAILED;
+        // Fusion dependencies keep external inputs live until this output is allocated.
+        const auto* dst         = nodes[count - 1];
+        const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+        plan->output_scratch    = !fresh_output && !(dst->data && !needs_scratch(dst, op));
+        *workspace              = plan->output_offset + (plan->output_scratch ? ggml_nbytes(dst) : 0);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void* state, void*) {
+        delete static_cast<Plan*>(state);
+    }
     static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
         auto& plan    = *static_cast<Plan*>(state);
         const auto op = describe(nodes, count);
@@ -424,14 +451,11 @@ namespace sd_i8_gated {
         ggml_ext_cuda_context context;
         if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
             return GGML_STATUS_FAILED;
-        auto stream   = static_cast<cudaStream_t>(context.stream);
-        auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
-        bool alias    = reinterpret_cast<uintptr_t>(dst->data) % 16 != 0;
-        for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
-            if (input == op.residual && !op.prefix && input->data == dst->data)
-                continue;
-            alias |= overlaps(dst, input);
-        }
+        auto stream      = static_cast<cudaStream_t>(context.stream);
+        auto* scratch    = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+        const bool alias = needs_scratch(dst, op);
+        if (alias && !plan.output_scratch)
+            return GGML_STATUS_FAILED;
         auto* output              = alias ? scratch + plan.output_offset : static_cast<char*>(dst->data);
         const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
         auto status               = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);

src/core/cuda/i8_gemm.cu

@@ -105,6 +105,39 @@ namespace sd_i8_cutlass {
         CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
     };
 
+    template <class Base>
+    struct SplitWeightIterator : Base {
+        using Layout      = typename Base::Layout;
+        using Pointer     = typename Base::Pointer;
+        using TensorCoord = typename Base::TensorCoord;
+        using AccessType  = typename Base::AccessType;
+        struct Params : Base::Params {
+            const int8_t* up = nullptr;
+            int64_t k        = 0;
+            Params()         = default;
+            CUTLASS_HOST_DEVICE Params(const Layout& layout)
+                : Base::Params(layout), k(layout.stride(0)) {}
+        };
+        uintptr_t gate, up;
+        size_t gate_bytes;
+        CUTLASS_DEVICE SplitWeightIterator(const Params& params, Pointer pointer, TensorCoord extent, int thread, TensorCoord offset, const int* indices = nullptr)
+            : Base(params, pointer, extent, thread, offset, indices),
+              gate(reinterpret_cast<uintptr_t>(pointer)),
+              up(reinterpret_cast<uintptr_t>(params.up)),
+              gate_bytes(size_t(params.k) * (extent.column() / 2)) {}
+
+        CUTLASS_DEVICE AccessType* get() const {
+            auto* address = Base::get();
+            if (!address)
+                return nullptr;
+            // Keep the packed SwiGLU iterator's column pairing, but redirect its
+            // virtual up half to a separate allocation. Params.k stays full-sized
+            // even when Stream-K restricts this iterator's K extent.
+            const uintptr_t offset = reinterpret_cast<uintptr_t>(address) - gate;
+            return offset < gate_bytes ? address : reinterpret_cast<AccessType*>(up + offset - gate_bytes);
+        }
+    };
+
     template <class Map>
     struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
         using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
@@ -137,34 +170,36 @@ namespace sd_i8_cutlass {
         }
     };
 
-    static __global__ void pair_scales(const float* src, float* dst, int n) {
+    static __global__ void pair_scales(const float* src, const float* up, float* dst, int n) {
         const int i = blockIdx.x * blockDim.x + threadIdx.x;
         if (i < n)
-            dst[i] = src[i / 2 + (i % 2) * (n / 2)];
+            dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)];
     }
 
-    template <class Config, bool StreamK>
+    template <class Config, bool StreamK, bool PackedW4 = false>
     struct DenseTypes {
-        using Tile     = typename Config::Tile;
-        using Warp     = typename Config::Warp;
-        using Swizzle  = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
-        using Map      = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
-        using Col      = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
-        using Row      = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
-        using Mul      = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
-        using ScaleX   = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
-        using ScaleW   = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
-        using Store    = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
-        using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
-        using Kernel   = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+        static constexpr bool kSplit = false;
+        using Weight                 = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
+        using Tile                   = typename Config::Tile;
+        using Warp                   = typename Config::Warp;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store                  = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Kernel                 = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
             int8_t,
             cutlass::layout::RowMajor,
             cutlass::ComplexTransform::kNone,
             16,
-            int8_t,
+            Weight,
             cutlass::layout::ColumnMajor,
             cutlass::ComplexTransform::kNone,
-            16,
+            PackedW4 ? 32 : 16,
             float,
             cutlass::layout::RowMajor,
             4,
@@ -178,48 +213,66 @@ namespace sd_i8_cutlass {
             Epilogue,
             Swizzle,
             4,
-            cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+            std::conditional_t<PackedW4, cutlass::arch::OpMultiplyAddMixedInputUpcast, cutlass::arch::OpMultiplyAddSaturate>>::GemmKernel;
         using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
-    template <class Config, bool StreamK>
+    template <class Config, bool StreamK, bool Split = false>
     struct SwigluTypes {
-        using Tile            = typename Config::Tile;
-        using Warp            = typename Config::Warp;
-        using Swizzle         = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
-        using Map             = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
-        using Col             = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
-        using Row             = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
-        using Mul             = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
-        using ScaleX          = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
-        using ScaleW          = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
-        using Store           = SwigluStore<Map>;
-        using Epilogue        = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
-        using GemmBase        = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
-        using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
-        using Kernel          = std::conditional_t<StreamK,
-                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
-                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
-        using Gemm            = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+        static constexpr bool kSplit = Split;
+        using Weight                 = int8_t;
+        using Tile                   = typename Config::Tile;
+        using Warp                   = typename Config::Warp;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store                  = SwigluStore<Map>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using GemmBase               = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+        using VisitorEpilogue        = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+        using BaseMma                = typename GemmBase::Mma;
+        using SplitMma               = cutlass::gemm::threadblock::MmaMultistage<
+            typename BaseMma::Shape,
+            typename BaseMma::IteratorA,
+            typename BaseMma::SmemIteratorA,
+            BaseMma::kCacheOpA,
+            SplitWeightIterator<typename BaseMma::IteratorB>,
+            typename BaseMma::SmemIteratorB,
+            BaseMma::kCacheOpB,
+            typename BaseMma::ElementC,
+            typename BaseMma::LayoutC,
+            typename BaseMma::Policy,
+            BaseMma::Detail::kStages>;
+        using Mma    = std::conditional_t<Split, SplitMma, BaseMma>;
+        using Kernel = std::conditional_t<StreamK,
+                                          cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<Mma, VisitorEpilogue, Swizzle>,
+                                          cutlass::gemm::kernel::GemmWithEpilogueVisitor<Mma, VisitorEpilogue, Swizzle>>;
+        using Gemm   = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
     struct Plan {
         DeviceInfo device;
         int64_t rows, n, k;
         bool swiglu;
+        bool packed_w4;
+        bool split;
         size_t kernel_bytes = 0, scratch_offset = 0;
-        int occupancy                                                                                                       = 0;
-        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t) = nullptr;
+        int occupancy                                                                                                                                    = 0;
+        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr;
     };
     template <class Types>
     static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
         typename Types::ScaleW::Arguments scales     = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
         typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
-        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, w, nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
+        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast<const typename Types::Weight*>(w), nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
     }
     template <class Types>
-    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales) {
         using Kernel = typename Types::Kernel;
         if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
-             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out)) %
+             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales)) %
             16)
             return GGML_STATUS_FAILED;
         auto* paired = static_cast<float*>(scratch);
@@ -236,14 +289,19 @@ namespace sd_i8_cutlass {
                 return GGML_STATUS_FAILED;
         }
         typename Kernel::Params params(args, p.device.sms, p.occupancy);
+        if constexpr (Types::kSplit) {
+            if (!up || !up_scales)
+                return GGML_STATUS_FAILED;
+            params.params_B.up = up;
+        }
         if (params.get_workspace_size() > p.kernel_bytes)
             return GGML_STATUS_FAILED;
         if (p.swiglu) {
-            pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, paired, int(p.n));
+            pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, up_scales, paired, int(p.n));
             if (cudaGetLastError() != cudaSuccess)
                 return GGML_STATUS_FAILED;
         }
-        if (params.init_workspace(static_cast<char*>(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess)
+        if (params.init_workspace(scratch ? static_cast<char*>(scratch) + p.scratch_offset : nullptr, stream) != cutlass::Status::kSuccess)
             return GGML_STATUS_FAILED;
         cutlass::arch::synclog_setup();
         cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
@@ -272,6 +330,10 @@ namespace sd_i8_cutlass {
     }
     template <class Config>
     static bool configure_tile(Plan& p) {
+        if (p.packed_w4)
+            return configure_kernel<DenseTypes<Config, false, true>>(p);
+        if (p.split)
+            return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true, true>>(p) : configure_kernel<SwigluTypes<Config, false, true>>(p);
         if (p.swiglu)
             return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
         return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
@@ -300,20 +362,22 @@ namespace sd_i8_cutlass {
     }
 }
 struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
-SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace) {
+static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4, bool split) {
 #ifdef SD_CUDA_FORCE_CUBLAS
     return nullptr;
 #else
     const auto info = sd_i8_cutlass::device_info(device);
-    if (info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
+    if ((packed_w4 && (swiglu || k > INT_MAX / 1024)) || info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
         rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
         return nullptr;
-    auto* p   = new SDI8CudaGemmPlan;
-    p->device = info;
-    p->rows   = rows;
-    p->n      = n;
-    p->k      = k;
-    p->swiglu = swiglu;
+    auto* p      = new SDI8CudaGemmPlan;
+    p->device    = info;
+    p->rows      = rows;
+    p->n         = n;
+    p->k         = k;
+    p->swiglu    = swiglu;
+    p->packed_w4 = packed_w4;
+    p->split     = split;
     if (!sd_i8_cutlass::configure(*p)) {
         delete p;
         return nullptr;
@@ -322,8 +386,19 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
     return p;
 #endif
 }
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) {
+    return prepare_gemm(device, rows, n, k, swiglu, workspace, packed_w4, false);
+}
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace) {
+    if (n > INT_MAX / 2 || k > INT_MAX / 16384)
+        return nullptr;
+    return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
+}
 ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
-    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream);
+    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr);
+}
+ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) {
+    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales);
 }
 void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
     delete plan;

src/core/cuda/i8_gemm.h

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

src/core/cuda/int8.cu

@@ -386,12 +386,55 @@ namespace SD_CUDA_NAMESPACE {
             return info.int8_blas &&
                    ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
         }
+        template <bool Quantize>
+        static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) {
+            const int64_t group = blockIdx.x, groups = k / 64, row = group / groups;
+            const int tid = threadIdx.x;
+            __shared__ float values[64];
+            __shared__ float maxima[16];
+#pragma unroll
+            for (int j = 0; j < 4; ++j)
+                values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f);
+            __syncthreads();
+#pragma unroll
+            for (int stride = 1; stride < 64; stride *= 4) {
+                const int i   = (tid / stride) * 4 * stride + tid % stride;
+                const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                values[i]              = convrot_h4_cuda(a, b, c, d, 0);
+                values[i + stride]     = convrot_h4_cuda(a, b, c, d, 1);
+                values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2);
+                values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3);
+                __syncthreads();
+            }
+            if constexpr (Quantize) {
+                const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row];
+#pragma unroll
+                for (int j = 0; j < 4; ++j) {
+                    const int value                = __float2int_rn(values[tid + j * 16] * inverse);
+                    dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value)));
+                }
+            } else {
+                float amax = 0.f;
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                    amax = fmaxf(amax, fabsf(values[tid + j * 16]));
+                maxima[tid] = amax;
+                __syncthreads();
+                for (int stride = 8; stride; stride >>= 1) {
+                    if (tid < stride)
+                        maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                    __syncthreads();
+                }
+                if (!tid)
+                    partial[group] = maxima[0];
+            }
+        }
         static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
             int group;
-            if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
+            if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0]))
                 return false;
             const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
-            return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0;
+            return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0;
         }
         enum Kind { MAT,
                     QUANTIZE,
@@ -402,6 +445,13 @@ namespace SD_CUDA_NAMESPACE {
             static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
             return disabled;
         }
+#ifdef SD_CUDA_I8_CUTLASS
+        static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) {
+            return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 &&
+                   aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
+                   (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16)));
+        }
+#endif
         static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
             auto* g = const_cast<ggml_cgraph*>(graph);
             if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
@@ -421,15 +471,30 @@ namespace SD_CUDA_NAMESPACE {
                 return 0;
             if (first + 2 < ggml_graph_n_nodes(g)) {
                 const auto* quant = ggml_graph_node(g, first + 2);
-                if (quant->src[0] == next && convrot_supports(device, quant, nullptr) &&
-                    size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
+                int group;
+                if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) &&
+                    size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) {
+#ifdef SD_CUDA_I8_CUTLASS
+                    if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) {
+                        size_t workspace          = 0;
+                        auto* gemm                = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], true, &workspace);
+                        const bool fused_epilogue = gemm != nullptr;
+                        sd_cuda_i8_gemm_free(gemm);
+                        // ConvRot is a separate launch either way. Keep its input in the
+                        // tensor arena so other operations can reuse this large allocation.
+                        if (fused_epilogue)
+                            return 2;
+                    }
+#endif
                     return 3;
+                }
             }
             return 2;
         }
         struct Plan {
             ggml_ext_gpu_device_info device;
             Kind kind;
+            int group;
             int64_t k, n, rows, padded;
             size_t quant, scales, accum, partial, output, gemm_offset;
             SDI8CudaGemmPlan* gemm = nullptr;
@@ -456,16 +521,17 @@ namespace SD_CUDA_NAMESPACE {
             p->n           = mat ? nodes[0]->ne[0] : 0;
             p->rows        = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]);
             p->padded      = (p->rows + 3) / 4 * 4;
-            size_t total   = 0;
+            p->group       = mat ? op.group_size : 0;
+            if (!mat)
+                sd_ggml_convrot_describe(dst, &p->group);
+            size_t total = 0;
             if (mat) {
                 p->quant  = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0);
                 p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0);
 #ifdef SD_CUDA_I8_CUTLASS
                 size_t gemm_bytes = 0;
                 // Retain cuBLAS when I32 accumulation can overflow.
-                if (op.weight_scale && !op.bias && p->k <= INT_MAX / 16384 &&
-                    aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
-                    (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) &&
+                if (cutlass_compatible(op, dst) &&
                     (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
                     if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
                         p->gemm        = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
@@ -480,9 +546,14 @@ namespace SD_CUDA_NAMESPACE {
                 p->accum                = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
             }
             const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
-            p->partial                = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
-            p->output                 = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
-            *workspace                = total;
+            p->partial                = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
+                                                                                                                                       : 0);
+            // Custom/concat outputs and packed GLU cannot reuse their input storage.
+            const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src &&
+                                      (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT ||
+                                       (dst->op == GGML_OP_GLU && !dst->src[1]));
+            p->output               = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
+            *workspace              = total;
             return GGML_STATUS_SUCCESS;
         }
         static void destroy(void* state, void*) {
@@ -499,6 +570,99 @@ namespace SD_CUDA_NAMESPACE {
             const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
             return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
         }
+#ifdef SD_CUDA_I8_CUTLASS
+        static bool split_describe(ggml_backend_dev_t device, ggml_tensor* const* nodes, SDI8Operation& gate, SDI8Operation& up) {
+            const auto* glu = nodes[2];
+            if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || glu->op_params[1] != 0 ||
+                !glu->src[0] || !glu->src[1] || glu->src[0] == glu->src[1] || glu->view_src ||
+                glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || !aligned(glu, 16) ||
+                !((nodes[0] == glu->src[0] && nodes[1] == glu->src[1]) || (nodes[1] == glu->src[0] && nodes[0] == glu->src[1])))
+                return false;
+            for (const auto* mat : {glu->src[0], glu->src[1]})
+                if (mat->view_src || !supports(device, mat, nullptr) || !ggml_are_same_shape(mat, glu))
+                    return false;
+            if (!sd_ggml_i8_describe(glu->src[0], &gate) || !sd_ggml_i8_describe(glu->src[1], &up) ||
+                gate.input != up.input || gate.input_scale != up.input_scale || gate.packed != up.packed ||
+                gate.input->type != GGML_TYPE_I8 || gate.bias || up.bias || !gate.weight_scale || !up.weight_scale ||
+                !ggml_are_same_shape(gate.weight, up.weight))
+                return false;
+            for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+                if (!aligned(t, 16))
+                    return false;
+            return true;
+        }
+        static int split_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (fusion_disabled() || first + 2 >= ggml_graph_n_nodes(g))
+                return 0;
+            ggml_tensor* nodes[] = {ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)};
+            SDI8Operation gate, up;
+            if (!split_describe(device, nodes, gate, up))
+                return 0;
+            size_t workspace     = 0;
+            auto* gemm           = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(nodes[2]), nodes[2]->ne[0], gate.weight->ne[0], &workspace);
+            const bool supported = gemm != nullptr;
+            sd_cuda_i8_gemm_free(gemm);
+            return supported ? 3 : 0;
+        }
+        struct SplitPlan {
+            SDI8CudaGemmPlan* gemm = nullptr;
+            size_t output          = SIZE_MAX;
+            int device;
+        };
+        static ggml_status split_prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            auto* p     = new SplitPlan;
+            *state      = p;
+            auto device = ggml_backend_get_device(backend);
+            p->device   = cuda_device_info(device).device;
+            SDI8Operation gate, up;
+            if (count != 3 || !split_describe(device, nodes, gate, up))
+                return GGML_STATUS_FAILED;
+            auto* dst    = nodes[2];
+            size_t bytes = 0;
+            p->gemm      = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(dst), dst->ne[0], gate.weight->ne[0], &bytes);
+            if (!p->gemm)
+                return GGML_STATUS_FAILED;
+            size_t total = (bytes + 255) / 256 * 256;
+            if (dst->data || dst->buffer || dst->view_src)
+                p->output = reserve(total, ggml_nbytes(dst));
+            *workspace = total;
+            return GGML_STATUS_SUCCESS;
+        }
+        static void split_destroy(void* state, void*) {
+            auto* p = static_cast<SplitPlan*>(state);
+            if (p)
+                sd_cuda_i8_gemm_free(p->gemm);
+            delete p;
+        }
+        static ggml_status split_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) {
+            auto& p   = *static_cast<SplitPlan*>(state);
+            auto* dst = nodes[2];
+            SDI8Operation gate, up;
+            if (!sd_ggml_i8_describe(dst->src[0], &gate) || !sd_ggml_i8_describe(dst->src[1], &up))
+                return GGML_STATUS_FAILED;
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != p.device)
+                return GGML_STATUS_FAILED;
+            auto stream   = static_cast<CUDAStream>(context.stream);
+            auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+            bool alias    = false;
+            for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+                alias |= overlaps(dst, t);
+            if (alias && p.output == SIZE_MAX)
+                return GGML_STATUS_FAILED;
+            auto* output         = alias ? reinterpret_cast<float*>(scratch + p.output) : static_cast<float*>(dst->data);
+            const auto* input    = static_cast<const int8_t*>(gate.input->data);
+            const int64_t padded = (ggml_nrows(dst) + 3) / 4 * 4;
+            const auto* scales   = gate.packed ? reinterpret_cast<const float*>(input + gate.weight->ne[0] * padded) : static_cast<const float*>(gate.input_scale->data);
+            auto status          = sd_cuda_i8_gemm_execute_split(p.gemm, input, static_cast<const int8_t*>(gate.weight->data), static_cast<const int8_t*>(up.weight->data),
+                                                                 scales, static_cast<const float*>(gate.weight_scale->data), static_cast<const float*>(up.weight_scale->data),
+                                                                 output, scratch, stream);
+            if (status != GGML_STATUS_SUCCESS || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success))
+                return GGML_STATUS_FAILED;
+            return GGML_STATUS_SUCCESS;
+        }
+#endif
         static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
             if (p.padded > p.rows) {
                 auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
@@ -506,6 +670,17 @@ namespace SD_CUDA_NAMESPACE {
                     return error;
             }
             const size_t shared = size_t(p.k) * sizeof(float);
+            if (p.group == 1) {
+                quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(src, dst, scales, p.k, p.rows);
+                return cuda_last_error();
+            }
+            if (p.group == 64) {
+                const auto groups = p.k / 64;
+                convrot_h64<false><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+                convrot_h64<true><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                return cuda_last_error();
+            }
             if (shared + 128 <= p.device.shared_memory) {
                 static std::mutex mutex;
                 static std::unordered_map<int, bool> configured;
@@ -582,9 +757,10 @@ namespace SD_CUDA_NAMESPACE {
             ggml_ext_gpu_context context;
             if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
                 return GGML_STATUS_FAILED;
-            auto stream   = static_cast<CUDAStream>(context.stream);
-            auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
-            auto* dst     = nodes[count - 1];
+            auto stream     = static_cast<CUDAStream>(context.stream);
+            auto* scratch   = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+            auto scratch_at = [scratch](size_t offset) { return scratch ? scratch + offset : nullptr; };
+            auto* dst       = nodes[count - 1];
             SDI8Operation op;
             const bool mat = sd_ggml_i8_describe(nodes[0], &op);
             bool alias     = false;
@@ -599,20 +775,22 @@ namespace SD_CUDA_NAMESPACE {
             // CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes.
             if (p.gemm && p.kind == SWIGLU_QUANTIZE)
                 alias = false;
-            void* output    = alias ? scratch + p.output : dst->data;
+            if (alias && p.output == SIZE_MAX)
+                return GGML_STATUS_FAILED;
+            void* output    = alias ? scratch_at(p.output) : dst->data;
             CUDAError error = cuda_success;
             if (!mat) {
                 auto* quant  = static_cast<int8_t*>(output);
                 auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
-                error        = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch + p.partial));
+                error        = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch_at(p.partial)));
             } else {
-                auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
+                auto* quant = reinterpret_cast<int8_t*>(scratch_at(p.quant));
                 const float* scales;
                 if (op.input->type == GGML_TYPE_F32) {
-                    auto* temporary_scales = reinterpret_cast<float*>(scratch + p.scales);
+                    auto* temporary_scales = reinterpret_cast<float*>(scratch_at(p.scales));
                     scales                 = temporary_scales;
                     if (op.group_size)
-                        error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
+                        error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch_at(p.partial)));
                     else {
                         if (p.padded > p.rows)
                             error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
@@ -636,12 +814,12 @@ namespace SD_CUDA_NAMESPACE {
                 const auto* bias        = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
                 const size_t prefix     = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
                 auto* gemm_output       = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
-                auto* accum             = reinterpret_cast<int32_t*>(scratch + p.accum);
+                auto* accum             = reinterpret_cast<int32_t*>(scratch_at(p.accum));
                 const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
 #ifdef SD_CUDA_I8_CUTLASS
                 if (p.gemm) {
                     auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
-                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch + p.gemm_offset, stream);
+                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream);
                     if (status != GGML_STATUS_SUCCESS)
                         return status;
                 } else
@@ -734,6 +912,17 @@ namespace SD_CUDA_NAMESPACE {
         desc.custom_op = sd_ggml_convrot_op();
         desc.supports  = sd_i8_cuda::convrot_supports;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#ifdef SD_CUDA_I8_CUTLASS
+        desc.name      = "sd.linear_i8.swiglu_split.cuda.v1";
+        desc.custom_op = sd_ggml_i8_op();
+        desc.priority  = 260;
+        desc.supports  = nullptr;
+        desc.match     = sd_i8_cuda::split_match;
+        desc.prepare   = sd_i8_cuda::split_prepare;
+        desc.execute   = sd_i8_cuda::split_execute;
+        desc.destroy   = sd_i8_cuda::split_destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#endif
     }
 
 }

src/core/cuda/registration.cu

@@ -4,11 +4,13 @@
 void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
 void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
 void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device);
 #endif
 
 namespace SD_CUDA_NAMESPACE {
     void register_rope(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
     void register_convrot(ggml_backend_dev_t device);
     void register_affine(ggml_backend_dev_t device);
     void register_modulate(ggml_backend_dev_t device);
@@ -25,6 +27,7 @@ namespace SD_CUDA_NAMESPACE {
         register_affine(device);
         register_modulate(device);
         register_i8(device);
+        register_w4(device);
         register_convrot(device);
         register_rope(device);
         register_bf16(device);
@@ -37,6 +40,7 @@ namespace SD_CUDA_NAMESPACE {
         sd_ggml_register_cuda_i8_projection(device);
         sd_ggml_register_cuda_i8_gated_residual(device);
         sd_ggml_register_cuda_attention(device);
+        sd_ggml_register_cuda_window_attention(device);
 #endif
     }
 }

src/core/cuda/rms_affine.cu

@@ -2,6 +2,7 @@
 #include <cstdlib>
 #include <cstring>
 #include "bf16.h"
+#include "core/ggml_extensions.h"
 #include "rms_affine.h"
 
 namespace SD_CUDA_NAMESPACE {
@@ -18,6 +19,28 @@ namespace SD_CUDA_NAMESPACE {
             return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t);
         }
 
+        static int modulate_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (disabled() || first + 2 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* norm     = ggml_graph_node(g, first);
+            const auto* affine   = ggml_graph_node(g, first + 1);
+            const auto* modulate = ggml_graph_node(g, first + 2);
+            if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm ||
+                !sd_ggml_modulate_describe(modulate) || modulate->src[0] != affine)
+                return 0;
+            const auto* x = norm->src[0];
+            if (!f32(x) || !f32(norm) || !f32(affine) || x->ne[0] < 32 || x->ne[0] > 8192 ||
+                ggml_nrows(x) > INT_MAX || !ggml_are_same_shape(x, norm) || !ggml_are_same_shape(x, affine))
+                return 0;
+            for (const auto* parameter : {affine->src[1], modulate->src[1], modulate->src[2]}) {
+                if (parameter && (!f32(parameter) || parameter == norm || parameter == affine ||
+                                  parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]))
+                    return 0;
+            }
+            return affine->src[1] ? 3 : 0;
+        }
+
         static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
             auto* g = const_cast<ggml_cgraph*>(graph);
             if (disabled() || first + 3 >= ggml_graph_n_nodes(g))
@@ -44,13 +67,17 @@ namespace SD_CUDA_NAMESPACE {
         static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
             if (disabled())
                 return;
-            for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) {
+            for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
                 const auto* mul = ggml_graph_node(graph, i);
-                const auto* add = ggml_graph_node(graph, i + 1);
-                if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
-                    continue;
+                if (!sd_ggml_modulate_describe(mul)) {
+                    if (i + 1 >= ggml_graph_n_nodes(graph))
+                        continue;
+                    const auto* add = ggml_graph_node(graph, i + 1);
+                    if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+                        continue;
+                }
                 const auto* affine = mul->src[0];
-                if (affine->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM)
+                if (!affine || affine->op != GGML_OP_MUL || !affine->src[0] || affine->src[0]->op != GGML_OP_RMS_NORM)
                     continue;
                 for (int j = i - 2; j >= 0 && j >= i - 32; --j) {
                     if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) {
@@ -68,64 +95,156 @@ namespace SD_CUDA_NAMESPACE {
             return value;
         }
 
-        template <int Block, bool Modulate, typename T = float>
-        static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) {
+        enum class Kind { Residual,
+                          SelfResidual,
+                          Modulate };
+
+        struct Operation {
+            Kind kind = Kind::Residual;
+            const ggml_tensor* input;
+            const ggml_tensor* weight;
+            const ggml_tensor* gate;
+            const ggml_tensor* residual = nullptr;
+            const ggml_tensor* shift    = nullptr;
+            ggml_tensor* output;
+            float eps;
+        };
+
+        static Operation describe(ggml_tensor* const* nodes, int count) {
+            Operation op;
+            op.input  = nodes[0]->src[0];
+            op.weight = nodes[1]->src[1];
+            op.gate   = nodes[2]->src[1];
+            op.output = nodes[count - 1];
+            std::memcpy(&op.eps, nodes[0]->op_params, sizeof(op.eps));
+            if (count == 3) {
+                op.kind  = Kind::Modulate;
+                op.shift = op.output->src[2];
+            } else {
+                op.residual = op.output->src[0] == nodes[2] ? op.output->src[1] : op.output->src[0];
+                if (op.residual == nodes[1]) {
+                    op.kind     = Kind::SelfResidual;
+                    op.residual = nullptr;
+                }
+            }
+            return op;
+        }
+
+        template <int Block, int Groups, Kind Op, typename T>
+        static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, const float* shift, T* dst, int width, float eps) {
             const int col0     = threadIdx.x;
             const int64_t base = int64_t(blockIdx.x) * width;
-            float values[8], sum = 0.f;
-            for (int j = 0, col = col0; col < width; ++j, col += Block) {
-                values[j] = x[base + col];
-                sum       = __fmaf_rn(values[j], values[j], sum);
-            }
-            sum = warp_sum(sum);
+            float values[Groups][8];
             __shared__ float sums[32];
             const int lane = col0 % 32;
-            if (lane == 0)
-                sums[col0 / 32] = sum;
+#pragma unroll
+            for (int group = 0; group < Groups; ++group) {
+                float sum = 0.f;
+                for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+                    values[group][j] = x[base + col];
+                    sum              = __fmaf_rn(values[group][j], values[group][j], sum);
+                }
+                sum = warp_sum(sum);
+                if (lane == 0)
+                    sums[(col0 + group * Block) / 32] = sum;
+            }
+            // Groups preserve the native reduction tree when using fewer physical threads.
             __syncthreads();
-            sum               = warp_sum(lane < Block / 32 ? sums[lane] : 0.f);
+            const float sum   = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f);
             const float scale = rsqrtf(sum / width + eps);
-            for (int j = 0, col = col0; col < width; ++j, col += Block) {
-                const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]);
-                const float r     = Modulate ? value : residual[base + col];
-                dst[base + col]   = T(__fadd_rn(__fmul_rn(value, gate[col]), r));
+#pragma unroll
+            for (int group = 0; group < Groups; ++group) {
+                for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+                    float value = __fmul_rn(__fmul_rn(scale, values[group][j]), weight[col]);
+                    if constexpr (Op == Kind::Modulate) {
+                        value = __fadd_rn(value, __fmul_rn(value, gate[col]));
+                        if (shift)
+                            value = __fadd_rn(value, shift[col]);
+                    } else {
+                        const float r = Op == Kind::SelfResidual ? value : residual[base + col];
+                        value         = __fadd_rn(__fmul_rn(value, gate[col]), r);
+                    }
+                    dst[base + col] = T(value);
+                }
             }
         }
 
+        static uintptr_t address(const ggml_tensor* tensor) {
+            size_t offset = 0;
+            while (tensor && !tensor->data && tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
+        }
         static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
-            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            if (!b)
+                return false;
+            const auto x = address(a), y = address(b);
+            if (!x || !y)
+                return false;
             return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
         }
+        static bool needs_scratch(const Operation& op) {
+            const auto* dst = op.output;
+            return (overlaps(dst, op.input) && address(dst) != address(op.input)) ||
+                   (overlaps(dst, op.residual) && address(dst) != address(op.residual)) ||
+                   overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
+        }
 
-        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
-            *workspace = ggml_nbytes(nodes[3]);
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+            const auto op    = describe(nodes, count);
+            const auto* dst  = op.output;
+            const bool fresh = !dst->data && !dst->buffer && !dst->view_src;
+            // Fresh outputs may reuse the residual elementwise; external inputs stay live.
+            *workspace = fresh || (dst->data && !needs_scratch(op)) ? 0 : ggml_nbytes(dst);
             return GGML_STATUS_SUCCESS;
         }
         static void destroy(void*, void*) {}
 
-        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+        template <typename T, Kind Op>
+        static CUDAError launch(CUDAStream stream, const Operation& op, T* output) {
+            auto compute = kernel<256, 1, Op, T>;
+            int threads  = 256;
+            if (op.input->ne[0] >= 1024) {
+                if constexpr (Op == Kind::Modulate) {
+                    compute = kernel<512, 2, Op, T>;
+                    threads = 512;
+                } else {
+                    compute = kernel<1024, 1, Op, T>;
+                    threads = 1024;
+                }
+            }
+            compute<<<unsigned(ggml_nrows(op.input)), threads, 0, stream>>>(
+                static_cast<const float*>(op.input->data), static_cast<const float*>(op.weight->data),
+                static_cast<const float*>(op.gate->data), op.residual ? static_cast<const float*>(op.residual->data) : nullptr,
+                op.shift ? static_cast<const float*>(op.shift->data) : nullptr, output, int(op.input->ne[0]), op.eps);
+            return cuda_last_error();
+        }
+
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
             ggml_ext_gpu_context context;
             if (!cuda_context(execution->backend, &context))
                 return GGML_STATUS_FAILED;
-            const auto* x         = nodes[0]->src[0];
-            const auto* w         = nodes[1]->src[1];
-            const auto* gate      = nodes[2]->src[1];
-            auto* dst             = nodes[3];
-            const auto* residual  = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0];
-            const bool modulation = residual == nodes[1];
-            const bool alias      = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) ||
-                               (!modulation && overlaps(dst, residual) && dst->data != residual->data);
-            auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
-            float eps;
-            std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
-            auto launch       = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>)
-                                                : (modulation ? kernel<1024, true> : kernel<1024, false>);
+            const auto op    = describe(nodes, count);
+            auto* dst        = op.output;
+            const bool alias = needs_scratch(op);
+            if (alias && !execution->workspace)
+                return GGML_STATUS_FAILED;
+            auto* output      = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
             const auto stream = static_cast<CUDAStream>(context.stream);
-            const int threads = x->ne[0] < 1024 ? 256 : 1024;
-            launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
-                static_cast<const float*>(x->data), static_cast<const float*>(w->data), static_cast<const float*>(gate->data),
-                modulation ? nullptr : static_cast<const float*>(residual->data), output, int(x->ne[0]), eps);
-            auto error = cuda_last_error();
+            CUDAError error   = cuda_invalid_value;
+            switch (op.kind) {
+                case Kind::Residual:
+                    error = launch<float, Kind::Residual>(stream, op, output);
+                    break;
+                case Kind::SelfResidual:
+                    error = launch<float, Kind::SelfResidual>(stream, op, output);
+                    break;
+                case Kind::Modulate:
+                    error = launch<float, Kind::Modulate>(stream, op, output);
+                    break;
+            }
             if (error == cuda_success && alias)
                 error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
             return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
@@ -137,19 +256,12 @@ namespace SD_CUDA_NAMESPACE {
     }
 
     ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) {
-        const auto* x         = nodes[0]->src[0];
-        const auto* residual  = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0];
-        const bool modulation = residual == nodes[1];
-        auto launch           = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>)
-                                                : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>);
-        float eps;
-        std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
-        const int threads = x->ne[0] < 1024 ? 256 : 1024;
-        launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
-            static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),
-            static_cast<const float*>(nodes[2]->src[1]->data), modulation ? nullptr : static_cast<const float*>(residual->data),
-            static_cast<CUDABFloat16*>(output), int(x->ne[0]), eps);
-        return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        const auto op    = sd_rms_affine::describe(nodes, 4);
+        auto* dst        = static_cast<CUDABFloat16*>(output);
+        const auto error = op.kind == sd_rms_affine::Kind::SelfResidual
+                               ? sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::SelfResidual>(stream, op, dst)
+                               : sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::Residual>(stream, op, dst);
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
     }
 
     void register_rms_affine(ggml_backend_dev_t device) {
@@ -164,5 +276,11 @@ namespace SD_CUDA_NAMESPACE {
         desc.destroy              = sd_rms_affine::destroy;
         desc.execute              = sd_rms_affine::execute;
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.rms_affine_modulate." SD_CUDA_NAME ".v1";
+        // Prefer the existing fusion that also produces ConvRot input when available.
+        desc.priority = 200;
+        desc.match    = sd_rms_affine::modulate_match;
+        desc.optimize = nullptr;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
     }
 }

src/core/cuda/rope.cu

@@ -169,7 +169,7 @@ namespace SD_CUDA_NAMESPACE {
             return false;
         const auto* x      = norm->src[0];
         const auto* weight = mul->src[1];
-        return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+        return x && weight && weight != norm && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
                x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
                ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
                weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
@@ -232,6 +232,13 @@ namespace SD_CUDA_NAMESPACE {
                 ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
                 return 5;
         }
+        if (first + 3 < ggml_graph_n_nodes(g)) {
+            const auto* rope = ggml_graph_node(g, first + 2);
+            const auto* copy = ggml_graph_node(g, first + 3);
+            if (copy->op == GGML_OP_CPY && copy->src[0] == rope && copy->type == GGML_TYPE_F16 &&
+                ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 4;
+        }
         return 3;
     }
 
@@ -270,7 +277,10 @@ namespace SD_CUDA_NAMESPACE {
     }
 
     static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
-        *workspace = ggml_nbytes(nodes[count - 1]);
+        const auto* dst = nodes[count - 1];
+        // Owning casts are fresh outputs; copies into existing storage retain view_src.
+        const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+        *workspace              = fresh_output ? 0 : ggml_nbytes(dst);
         return GGML_STATUS_SUCCESS;
     }
 
@@ -293,7 +303,7 @@ namespace SD_CUDA_NAMESPACE {
         void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
         CUDAError error;
         if (count > 1) {
-            error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+            error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
         } else {
             RopeApplyStrides strides;
             const auto* pe = rope->src[1];

src/core/cuda/w4.cu

@@ -0,0 +1,939 @@
+#include <climits>
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+#ifdef SD_CUDA_W4_CUTLASS
+#include <cstdlib>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+#include "w4_mma.h"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_w4_cuda {
+        static __device__ __forceinline__ int shuffle_sum(int value, int lane, int width) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            // ROCm LLVM 24 can crash lowering __shfl's wave_shuffle on RDNA2.
+            // These 8/16-lane groups stay within one ds_bpermute half-wave.
+            const int index = (lane & (width - 1)) + (__lane_id() & ~(width - 1));
+            return __builtin_amdgcn_ds_bpermute(index * 4, value);
+#else
+            return cuda_shuffle(value, lane, width);
+#endif
+        }
+
+        // Signed byte dots preserve the INT4 grid after the XOR-8 bias correction.
+        static __device__ __forceinline__ int dot4(int a, int b, int c) {
+#if defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
+            return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
+#elif defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__))
+            return __builtin_amdgcn_sdot4(a, b, c, false);
+#elif defined(SD_CUDA_VENDOR_MUSA) || (!defined(SD_CUDA_VENDOR_HIP) && __CUDA_ARCH__ >= 610)
+            return __dp4a(a, b, c);
+#else
+            const auto* av = reinterpret_cast<const signed char*>(&a);
+            const auto* bv = reinterpret_cast<const signed char*>(&b);
+            for (int i = 0; i < 4; ++i)
+                c += int(av[i]) * int(bv[i]);
+            return c;
+#endif
+        }
+        static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
+            const int exponent = (v >> 3) & 15, mantissa = v & 7;
+            float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9)
+                                                                                                        : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7);
+            return v & 128 ? -result : result;
+        }
+        // A CTA shares decoded weights across 32 activation rows. Each double-buffered
+        // stage covers 64 K values; W4A8 restores the codebook scale per 16-value group.
+
+        static __device__ __forceinline__ uint32_t perm_bytes(uint32_t lo, uint32_t hi, uint32_t sel) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            return __builtin_amdgcn_perm(hi, lo, sel);
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+            return __byte_perm(lo, hi, sel);
+#else
+            uint32_t r = 0;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const uint32_t b    = (sel >> (4 * i)) & 0xF;
+                const uint32_t byte = (b < 4 ? lo : hi) >> (8 * (b & 3)) & 0xFF;
+                r |= byte << (8 * i);
+            }
+            return r;
+#endif
+        }
+
+#if defined(SD_CUDA_VENDOR_HIP)
+#define PERM_SEL_EVEN 0x06040200u
+#define PERM_SEL_ODD 0x07050301u
+#else
+#define PERM_SEL_EVEN 0x6420u
+#define PERM_SEL_ODD 0x7531u
+#endif
+
+        static __device__ __forceinline__ void interleave_act_pair(
+            const int lo,
+            const int hi,
+            int* __restrict__ dst) {
+            dst[0] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_EVEN);
+            dst[1] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_ODD);
+        }
+
+#define MUL_MAT_W4_BLOCK 16
+#define MUL_MAT_W4_ROWS (2 * MUL_MAT_W4_BLOCK)
+#define MUL_MAT_W4_GPS 4
+#define MUL_MAT_W4_SLOTS (4 * MUL_MAT_W4_GPS)
+#define MUL_MAT_W4_OUT (MUL_MAT_W4_BLOCK * 4)
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4_BLOCK* MUL_MAT_W4_BLOCK) mul_mat_w4a4_convrot_dp4a(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            int* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t packed_k) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4_SLOTS * MUL_MAT_W4_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4_ROWS * MUL_MAT_W4_SLOTS];
+
+            const uint32_t w_words = packed_k / 4;
+            const uint32_t x_words = (packed_k * 2) / 4;
+            const uint32_t groups  = packed_k / 8;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            const uint32_t fl = lx & 7;
+            const uint32_t fr = (lx >> 3) ? r1 : r0;
+            const uint32_t tr = (lx >> 3) * MUL_MAT_W4_BLOCK + ly;
+            int my_sum        = 0;
+
+            {
+                const uint32_t ocol = o_base + dc;
+                int2 wv             = make_int2(0, 0);
+                if (ocol < n && dw * 2 < w_words) {
+                    wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + dw];
+                }
+                w_dec[0][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+
+                const uint32_t gi = fl >> 1;
+                const uint32_t xb = gi * 4 + (fl & 1) * 2;
+                const int a       = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                const int b       = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                interleave_act_pair(a, b, &x_tile[0][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+            }
+            __syncthreads();
+
+            int acc0a = 0, acc0b = 0;
+            int acc1a = 0, acc1b = 0;
+            int acc2a = 0, acc2b = 0;
+            int acc3a = 0, acc3b = 0;
+            int acc4a = 0, acc4b = 0;
+            int acc5a = 0, acc5b = 0;
+            int acc6a = 0, acc6b = 0;
+            int acc7a = 0, acc7b = 0;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4_GPS;
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        const uint32_t widx = ngs * 2 + dw * 2;
+                        int2 wv             = make_int2(0, 0);
+                        if (ocol < n && widx < w_words) {
+                            wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + ngs + dw];
+                        }
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        const uint32_t gi                                         = ngs + (fl >> 1);
+                        const uint32_t xb                                         = gi * 4 + (fl & 1) * 2;
+                        const int a                                               = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                        const int b                                               = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                        interleave_act_pair(a, b, &x_tile[nbuf][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                        my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+                    }
+                }
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+#pragma unroll
+                    for (uint32_t jp = 0; jp < 2; ++jp) {
+                        const int ue0     = jp == 0 ? u0[0] : u0[2];
+                        const int uo0     = jp == 0 ? u0[1] : u0[3];
+                        const int ue1     = jp == 0 ? u1[0] : u1[2];
+                        const int uo1     = jp == 0 ? u1[1] : u1[3];
+                        const uint32_t we = (gl * 4 + jp * 2) * MUL_MAT_W4_OUT;
+                        const uint32_t wo = we + MUL_MAT_W4_OUT;
+
+                        acc0a = dot4(w_dec[buf][we + lx], ue0, acc0a);
+                        acc1a = dot4(w_dec[buf][we + 16 + lx], ue0, acc1a);
+                        acc2a = dot4(w_dec[buf][we + 32 + lx], ue0, acc2a);
+                        acc3a = dot4(w_dec[buf][we + 48 + lx], ue0, acc3a);
+                        acc0b = dot4(w_dec[buf][wo + lx], uo0, acc0b);
+                        acc1b = dot4(w_dec[buf][wo + 16 + lx], uo0, acc1b);
+                        acc2b = dot4(w_dec[buf][wo + 32 + lx], uo0, acc2b);
+                        acc3b = dot4(w_dec[buf][wo + 48 + lx], uo0, acc3b);
+                        acc4a = dot4(w_dec[buf][we + lx], ue1, acc4a);
+                        acc5a = dot4(w_dec[buf][we + 16 + lx], ue1, acc5a);
+                        acc6a = dot4(w_dec[buf][we + 32 + lx], ue1, acc6a);
+                        acc7a = dot4(w_dec[buf][we + 48 + lx], ue1, acc7a);
+                        acc4b = dot4(w_dec[buf][wo + lx], uo1, acc4b);
+                        acc5b = dot4(w_dec[buf][wo + 16 + lx], uo1, acc5b);
+                        acc6b = dot4(w_dec[buf][wo + 32 + lx], uo1, acc6b);
+                        acc7b = dot4(w_dec[buf][wo + 48 + lx], uo1, acc7b);
+                    }
+                }
+                __syncthreads();
+            }
+
+            int v = my_sum;
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 4) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 2) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 1) % 8, 8);
+            const int sum0 = shuffle_sum(v, 0, 16);
+            const int sum1 = shuffle_sum(v, 8, 16);
+
+            const int accv[2][4] = {
+                {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
+                {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
+            };
+            const int sumv[2] = {sum0, sum1};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            y[(size_t)rt * n + oj] = accv[t][j] - 8 * sumv[t];
+                        }
+                    }
+                }
+            }
+        }
+
+#define MUL_MAT_W4A8_BLOCK 16
+#define MUL_MAT_W4A8_ROWS (2 * MUL_MAT_W4A8_BLOCK)
+#define MUL_MAT_W4A8_GPS 4
+#define MUL_MAT_W4A8_SLOTS (4 * MUL_MAT_W4A8_GPS)
+#define MUL_MAT_W4A8_OUT (MUL_MAT_W4A8_BLOCK * 4)
+
+#if defined(SD_CUDA_VENDOR_HIP) && defined(__HIP_DEVICE_COMPILE__)
+// __has_builtin checks the active device target's features, not just whether
+// the compiler recognizes the name. Prefer the non-replicated fragment layout.
+#if __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12)
+#define SD_W4_WMMA_GFX12
+#elif __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64)
+#define SD_W4_WMMA_GFX11
+#endif
+#endif
+
+// Keep each intrinsic paired with its fragment layout, including on newer GPUs.
+#if defined(SD_W4_WMMA_GFX11)
+        typedef int v4si32_w4a8 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t q      = lane >> 4;
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]           = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8 zero = {};
+
+            int w0c = 0, w1c = 0, xc0 = 0, xc1 = 0, xc2 = 0, xc3 = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn0 = 0, xn1 = 0, xn2 = 0, xn3 = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc0           = ra < rows ? xr[0] : 0;
+                xc1           = ra < rows ? xr[1] : 0;
+                xc2           = ra < rows ? xr[2] : 0;
+                xc3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words + 4;
+                xn0           = ra < rows ? xr[0] : 0;
+                xn1           = ra < rows ? xr[1] : 0;
+                xn2           = ra < rows ? xr[2] : 0;
+                xn3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                v4si32_w4a8 bfrag = {};
+                float s           = 0.0f;
+                if (o_valid) {
+                    bfrag = {
+                        (lut[w0c & 0xFF] & 0xFFFF) | ((lut[(w0c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w0c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w0c >> 24) & 0xFF] & 0xFFFF) << 16),
+                        (lut[w1c & 0xFF] & 0xFFFF) | ((lut[(w1c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w1c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w1c >> 24) & 0xFF] & 0xFFFF) << 16),
+                    };
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8 afrag = {xc0, xc1, xc2, xc3};
+                const v4si32_w4a8 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64(true, afrag, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc0 = xn0;
+                xc1 = xn1;
+                xc2 = xn2;
+                xc3 = xn3;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words + (g + 2) * 4;
+                    xn0           = ra < rows ? xr[0] : 0;
+                    xn1           = ra < rows ? xr[1] : 0;
+                    xn2           = ra < rows ? xr[2] : 0;
+                    xn3           = ra < rows ? xr[3] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * d + q;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+#if defined(SD_W4_WMMA_GFX12)
+        typedef int v4si32_w4a8_g12 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma_gfx12(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t h      = lane >> 4;
+            const uint32_t kh     = ((h & 1) << 1) | (h >> 1);
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]               = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8_g12 zero = {};
+
+            int w0c = 0, w1c = 0, xc = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc            = ra < rows ? xr[kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words;
+                xn            = ra < rows ? xr[4 + kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                int bfrag = 0;
+                float s   = 0.0f;
+                if (o_valid) {
+                    const int bwd = (kh & 2) ? w1c : w0c;
+                    const int bsh = (kh & 1) * 16;
+                    bfrag         = (lut[(bwd >> bsh) & 0xFF] & 0xFFFF) |
+                                    ((lut[(bwd >> (bsh + 8)) & 0xFF] & 0xFFFF) << 16);
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8_g12 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12(true, xc, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc  = xn;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words;
+                    xn            = ra < rows ? xr[(g + 2) * 4 + kh] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * kh + d;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4A8_BLOCK* MUL_MAT_W4A8_BLOCK) mul_mat_w4a8_convrot(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k) {
+#if defined(SD_W4_WMMA_GFX11) || defined(SD_W4_WMMA_GFX12)
+            // Retain the 32-row launch grid shared with DP4A. Each WMMA wave
+            // computes a 16-row tile; the barrier protects the reused LUTs.
+            for (uint32_t tile = 0; tile < 2; ++tile) {
+#if defined(SD_W4_WMMA_GFX12)
+                mul_mat_w4a8_convrot_wmma_gfx12(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                                blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#else
+                mul_mat_w4a8_convrot_wmma(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                          blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#endif
+                __syncthreads();
+            }
+#else
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4A8_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4A8_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4A8_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4A8_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4A8_SLOTS * MUL_MAT_W4A8_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4A8_ROWS * MUL_MAT_W4A8_SLOTS];
+            __shared__ float s_dec[2][MUL_MAT_W4A8_GPS * MUL_MAT_W4A8_OUT];
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (ly == 0 && lx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : __float2int_rn(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+            }
+            __syncthreads();
+
+            lut[idx] = (cb8[idx & 0xF] & 0xFF) | ((cb8[idx >> 4] & 0xFF) << 8);
+            __syncthreads();
+
+#pragma unroll
+            for (uint32_t j = 0; j < 2; ++j) {
+                const uint32_t pw                              = dw * 2 + j;
+                const uint32_t ocol                            = o_base + dc;
+                const int wv                                   = (ocol < n && pw < w_words) ? wi[(size_t)ocol * w_words + pw] : 0;
+                w_dec[0][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                w_dec[0][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+            }
+            {
+                const uint32_t ocol = o_base + dc;
+                float s             = 0.0f;
+                if (ocol < n && dw < groups) {
+                    s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + dw]);
+                }
+                s_dec[0][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+            }
+            if (lx < 16) {
+                const uint32_t aidx = lx;
+                x_tile[0][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                x_tile[0][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+            }
+            __syncthreads();
+
+            float acc0 = 0.0f;
+            float acc1 = 0.0f;
+            float acc2 = 0.0f;
+            float acc3 = 0.0f;
+            float acc4 = 0.0f;
+            float acc5 = 0.0f;
+            float acc6 = 0.0f;
+            float acc7 = 0.0f;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4A8_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4A8_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4A8_GPS;
+#pragma unroll
+                    for (uint32_t j = 0; j < 2; ++j) {
+                        const uint32_t pw                                 = dw * 2 + j;
+                        const uint32_t ocol                               = o_base + dc;
+                        const uint32_t widx                               = ngs * 2 + pw;
+                        const int wv                                      = (ocol < n && widx < w_words) ? wi[(size_t)ocol * w_words + widx] : 0;
+                        w_dec[nbuf][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                        w_dec[nbuf][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+                    }
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        float s             = 0.0f;
+                        if (ocol < n && ngs + dw < groups) {
+                            s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + ngs + dw]);
+                        }
+                        s_dec[nbuf][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+                    }
+                    if (lx < 16) {
+                        const uint32_t aidx = ngs * 4 + lx;
+                        x_tile[nbuf][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                        x_tile[nbuf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+                    }
+                }
+
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4A8_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    int accg0 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u0[0], 0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u0[1], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u0[2], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u0[3], accg0);
+                    int accg1 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u0[0], 0);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u0[1], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u0[2], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u0[3], accg1);
+                    int accg2 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u0[0], 0);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u0[1], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u0[2], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u0[3], accg2);
+                    int accg3 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u0[0], 0);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u0[1], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u0[2], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u0[3], accg3);
+                    int accg4 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u1[0], 0);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u1[1], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u1[2], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u1[3], accg4);
+                    int accg5 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u1[0], 0);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u1[1], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u1[2], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u1[3], accg5);
+                    int accg6 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u1[0], 0);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u1[1], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u1[2], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u1[3], accg6);
+                    int accg7 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u1[0], 0);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u1[1], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u1[2], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u1[3], accg7);
+                    acc0      = fmaf((float)accg0, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc0);
+                    acc1      = fmaf((float)accg1, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc1);
+                    acc2      = fmaf((float)accg2, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc2);
+                    acc3      = fmaf((float)accg3, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc3);
+                    acc4      = fmaf((float)accg4, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc4);
+                    acc5      = fmaf((float)accg5, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc5);
+                    acc6      = fmaf((float)accg6, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc6);
+                    acc7      = fmaf((float)accg7, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc7);
+                }
+                __syncthreads();
+            }
+
+            const float accv[8] = {acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            float value = accv[t * 4 + j] * act_scales[rt];
+                            if (bias != nullptr) {
+                                value += bias[oj];
+                            }
+                            y[(size_t)rt * n + oj] = value;
+                        }
+                    }
+                }
+            }
+#endif
+        }
+
+        static __global__ void dequantize(const int32_t* accum, const float* xs, const float* ws, const float* bias, float* dst, int64_t n, int64_t total) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i < total) {
+                float value = float(accum[i]) * xs[i / n] * ws[i % n];
+                dst[i]      = value + (bias ? bias[i % n] : 0.f);
+            }
+        }
+        static bool aligned(const ggml_tensor* t, size_t alignment) {
+            if (!t)
+                return true;
+            size_t offset = 0;
+            while (!t->data && t->view_src) {
+                offset += t->view_offs;
+                t = t->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
+        }
+        static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(node, &op) || cuda_device_info(device).device < 0)
+                return false;
+            const int64_t k = op.input->ne[0], n = node->ne[0], rows = ggml_nrows(node);
+            if (k % 64 || k > INT_MAX || n > INT_MAX - 63 || rows > 65535 * 32 ||
+                ggml_nelements(node) > int64_t(INT_MAX) * 256 || !aligned(op.weight, 8) || !aligned(node, 4))
+                return false;
+            for (auto* t : node->src)
+                if (!aligned(t, 4))
+                    return false;
+            return true;
+        }
+        struct Plan {
+            bool mma           = false;
+            size_t gemm_offset = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+            SDI8CudaGemmPlan* gemm = nullptr;
+            ~Plan() { sd_cuda_i8_gemm_free(gemm); }
+#endif
+        };
+        static bool use_mma(ggml_backend_t backend, const SDW4Operation& op) {
+#ifdef SD_CUDA_W4_CUTLASS
+            if (!aligned(op.input, 16) || !aligned(op.weight, 16))
+                return false;
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_W4_FORCE_DP4A"); return v && std::atoi(v); }();
+            if (disabled)
+                return false;
+            static std::mutex mutex;
+            static std::unordered_map<int, bool> devices;
+            const auto device = cuda_device_info(ggml_backend_get_device(backend)).device;
+            std::lock_guard<std::mutex> lock(mutex);
+            auto found = devices.find(device);
+            if (found != devices.end())
+                return found->second;
+            cudaDeviceProp prop  = {};
+            const bool available = cudaGetDeviceProperties(&prop, device) == cudaSuccess && prop.major >= 8;
+            devices.emplace(device, available);
+            return available;
+#else
+            return false;
+#endif
+        }
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int, void** state, size_t* bytes, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(nodes[0], &op))
+                return GGML_STATUS_FAILED;
+            const auto* dst = nodes[0];
+            bool may_alias  = dst->data || dst->view_src;
+            if (dst->data && dst->buffer && !dst->view_src) {
+                may_alias = false;
+                for (const auto* input : dst->src) {
+                    if (!input)
+                        continue;
+                    size_t offset     = 0;
+                    const size_t size = ggml_nbytes(input);
+                    while (!input->data && input->view_src) {
+                        offset += input->view_offs;
+                        input = input->view_src;
+                    }
+                    if (!input->data)
+                        continue;
+                    const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(input->data) + offset;
+                    may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst));
+                }
+            }
+            auto* plan        = new Plan;
+            plan->mma         = use_mma(backend, op);
+            *state            = plan;
+            size_t gemm_bytes = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+            if (plan->mma && op.kind == SDW4Kind::W4A4) {
+                if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16))
+                    plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true);
+                plan->mma = plan->gemm != nullptr;
+                if (plan->gemm && gemm_bytes == 0)
+                    gemm_bytes = 256;
+            }
+#endif
+            const size_t tensor_bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 && !plan->mma ? 1 : 0) + (may_alias ? 1 : 0));
+            plan->gemm_offset         = (tensor_bytes + 255) / 256 * 256;
+            *bytes                    = plan->gemm_offset + gemm_bytes;
+            return GGML_STATUS_SUCCESS;
+        }
+        static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+            if (!src)
+                return false;
+            const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+            return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+        }
+        static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void* state, void*) {
+            auto* dst = nodes[0];
+            SDW4Operation op;
+            ggml_ext_gpu_context context;
+            if (!sd_ggml_w4_describe(dst, &op) || !cuda_context(e->backend, &context))
+                return GGML_STATUS_FAILED;
+            const auto stream = static_cast<CUDAStream>(context.stream);
+            const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0];
+            auto* scratch = e->workspace ? static_cast<float*>(ggml_backend_buffer_get_base(e->workspace)) : nullptr;
+            bool alias    = false;
+            for (auto* input : dst->src)
+                alias |= overlaps(dst, input);
+            auto* output     = alias ? scratch : static_cast<float*>(dst->data);
+            const auto* w    = static_cast<const char*>(op.weight->data);
+            const auto* x    = static_cast<const char*>(op.input->data);
+            const auto* xs   = reinterpret_cast<const float*>(x + k * ((rows + 3) / 4 * 4));
+            const auto* ws   = static_cast<const float*>(op.weight_scale->data);
+            const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+            const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16);
+#ifdef SD_CUDA_W4_CUTLASS
+            auto* plan = static_cast<Plan*>(state);
+            if (plan->mma) {
+                if (plan->gemm) {
+                    auto* workspace = scratch ? reinterpret_cast<char*>(scratch) + plan->gemm_offset : nullptr;
+                    if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast<const int8_t*>(x), reinterpret_cast<const int8_t*>(w), xs, ws, output, workspace, stream) != GGML_STATUS_SUCCESS)
+                        return GGML_STATUS_FAILED;
+                } else {
+                    const dim3 tiles(unsigned((n + 63) / 64), unsigned((rows + 127) / 128));
+                    sd_w4_mma::w4a8<128, 64, 128><<<tiles, 256, 0, stream>>>(reinterpret_cast<const uint8_t*>(w), reinterpret_cast<const int8_t*>(x), xs, ws,
+                                                                             static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data), bias, output, int(rows), int(n), int(k));
+                }
+            } else
+#endif
+                if (op.kind == SDW4Kind::W4A4) {
+                auto* accum = reinterpret_cast<int32_t*>(scratch + (alias ? ggml_nelements(dst) : 0));
+                mul_mat_w4a4_convrot_dp4a<<<grid, block, 0, stream>>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2));
+                if (cuda_last_error() != cuda_success)
+                    return GGML_STATUS_FAILED;
+                dequantize<<<unsigned((n * rows + 255) / 256), 256, 0, stream>>>(accum, xs, ws, bias, output, n, n * rows);
+            } else {
+                mul_mat_w4a8_convrot<<<grid, block, 0, stream>>>(w, x, ws,
+                                                                 static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data),
+                                                                 xs, bias, output, uint32_t(n), uint32_t(rows), uint32_t(k));
+            }
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+    void register_w4(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.name                 = "sd.linear_w4_convrot." SD_CUDA_NAME ".v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = sd_w4_cuda::supports;
+        d.prepare              = sd_w4_cuda::prepare;
+        d.execute              = sd_w4_cuda::execute;
+        d.destroy              = [](void* state, void*) { delete static_cast<sd_w4_cuda::Plan*>(state); };
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/cuda/w4_mma.h

@@ -0,0 +1,138 @@
+#ifndef __SD_CORE_CUDA_W4_MMA_H__
+#define __SD_CORE_CUDA_W4_MMA_H__
+
+#include <cuda_runtime.h>
+#include <cutlass/arch/mma_sm80.h>
+#include <cstdint>
+
+namespace sd_w4_mma {
+    static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
+        const int exponent = (v >> 3) & 15, mantissa = v & 7;
+        const float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9)
+                                                                                                          : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7);
+        return v & 128 ? -result : result;
+    }
+
+    template <int M, int N, int K>
+    static __global__ void __launch_bounds__(256) w4a8(
+        const uint8_t* __restrict__ w,
+        const int8_t* __restrict__ x,
+        const float* __restrict__ xs,
+        const float* __restrict__ codebook,
+        const float* __restrict__ channel,
+        const uint8_t* __restrict__ relative,
+        const float* __restrict__ bias,
+        float* __restrict__ dst,
+        int rows,
+        int n,
+        int k) {
+        const int row_base = blockIdx.y * M, col_base = blockIdx.x * N;
+        const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32;
+        __shared__ int cb[16];
+        __shared__ float cbscale;
+        if (tid == 0) {
+            float amax = 0.f;
+#pragma unroll
+            for (int i = 0; i < 16; ++i)
+                amax = fmaxf(amax, fabsf(codebook[i]));
+            cbscale = amax / 127.f;
+#pragma unroll
+            for (int i = 0; i < 16; ++i) {
+                const int value = cbscale == 0.f ? 0 : __float2int_rn(codebook[i] / cbscale);
+                cb[i]           = max(-127, min(127, value));
+            }
+        }
+        __syncthreads();
+#if __CUDA_ARCH__ >= 800
+        constexpr int stride = K / 4 + 4;
+        __shared__ __align__(16) uint32_t a[M][stride], b[N][stride];
+        __shared__ uint32_t lut[256];
+        __shared__ float scales[K / 16][N];
+        lut[tid] = (cb[tid & 15] & 255) | ((cb[tid >> 4] & 255) << 8);
+        __syncthreads();
+        constexpr int warp_cols = N / 32;
+        constexpr int row_tiles = M / (16 * (8 / warp_cols));
+        const int wr = (warp / warp_cols) * 16 * row_tiles, wc = (warp % warp_cols) * 32;
+        float acc[row_tiles][4][4] = {};
+        for (int base = 0; base < k; base += K) {
+            for (int i = tid; i < M * K / 16; i += 256) {
+                const int r = i / (K / 16), t = (i % (K / 16)) * 16;
+                const int4 value                       = row_base + r < rows && base + t < k ? *reinterpret_cast<const int4*>(x + int64_t(row_base + r) * k + base + t) : make_int4(0, 0, 0, 0);
+                *reinterpret_cast<int4*>(&a[r][t / 4]) = value;
+            }
+            for (int i = tid; i < N * K / 8; i += 256) {
+                const int c = i / (K / 8), t = i % (K / 8);
+                const uint32_t value = col_base + c < n && base + t * 8 < k ? *reinterpret_cast<const uint32_t*>(w + int64_t(col_base + c) * (k / 2) + base / 2 + t * 4) : 0;
+                b[c][2 * t]          = lut[value & 255] | (lut[(value >> 8) & 255] << 16);
+                b[c][2 * t + 1]      = lut[(value >> 16) & 255] | (lut[value >> 24] << 16);
+            }
+            for (int i = tid; i < N * K / 16; i += 256) {
+                const int c = i % N, g = i / N;
+                scales[g][c] = col_base + c < n && base + g * 16 < k ? cbscale * (channel[col_base + c] * fp8_e4m3(relative[int64_t(col_base + c) * (k / 16) + base / 16 + g])) : 0.f;
+            }
+            __syncthreads();
+#pragma unroll
+            for (int g = 0; g < K / 16; ++g) {
+#pragma unroll
+                for (int r = 0; r < row_tiles; ++r) {
+                    const int ar = wr + r * 16 + lane / 4, ac = g * 4 + lane % 4;
+                    const uint32_t a0 = a[ar][ac], a1 = a[ar + 8][ac];
+#pragma unroll
+                    for (int j = 0; j < 4; ++j) {
+                        const uint32_t b0 = b[wc + j * 8 + lane / 4][ac];
+                        using Mma         = cutlass::arch::Mma<cutlass::gemm::GemmShape<16, 8, 16>, 32,
+                                                               int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
+                                                               int, cutlass::layout::RowMajor, cutlass::arch::OpMultiplyAddSaturate>;
+                        Mma::FragmentA af;
+                        Mma::FragmentB bf;
+                        Mma::FragmentC d;
+                        reinterpret_cast<uint32_t*>(&af)[0] = a0;
+                        reinterpret_cast<uint32_t*>(&af)[1] = a1;
+                        reinterpret_cast<uint32_t*>(&bf)[0] = b0;
+                        d.clear();
+                        // Each K=16 MMA has its own scale; combining groups before
+                        // dequantization would change the W4A8 representation.
+                        Mma{}(d, af, bf, d);
+                        const float s0 = scales[g][wc + j * 8 + (lane % 4) * 2], s1 = scales[g][wc + j * 8 + (lane % 4) * 2 + 1];
+#pragma unroll
+                        for (int v = 0; v < 4; ++v)
+                            acc[r][j][v] = fmaf(float(d[v]), v % 2 ? s1 : s0, acc[r][j][v]);
+                    }
+                }
+            }
+            __syncthreads();
+        }
+#pragma unroll
+        for (int r = 0; r < row_tiles; ++r) {
+#pragma unroll
+            for (int j = 0; j < 4; ++j) {
+#pragma unroll
+                for (int v = 0; v < 4; ++v) {
+                    const int row = row_base + wr + r * 16 + lane / 4 + (v / 2) * 8;
+                    const int col = col_base + wc + j * 8 + (lane % 4) * 2 + v % 2;
+                    if (row < rows && col < n)
+                        dst[int64_t(row) * n + col] = __fmul_rn(acc[r][j][v], xs[row]) + (bias ? bias[col] : 0.f);
+                }
+            }
+        }
+#else
+        for (int i = tid; i < M * N; i += 256) {
+            const int row = row_base + i / N, col = col_base + i % N;
+            if (row < rows && col < n) {
+                float sum = 0.f;
+                for (int g = 0; g < k / 16; ++g) {
+                    int d = 0;
+                    for (int j = 0; j < 16; ++j) {
+                        const int t = g * 16 + j;
+                        d += cb[(w[int64_t(col) * (k / 2) + t / 2] >> (4 * (t % 2))) & 15] * int(x[int64_t(row) * k + t]);
+                    }
+                    sum = fmaf(float(d), cbscale * (channel[col] * fp8_e4m3(relative[int64_t(col) * (k / 16) + g])), sum);
+                }
+                dst[int64_t(row) * n + col] = __fmul_rn(sum, xs[row]) + (bias ? bias[col] : 0.f);
+            }
+        }
+#endif
+    }
+
+}
+#endif

src/core/cuda/window_attention.cuh

@@ -0,0 +1,284 @@
+/*
+ * Adapted from FlashAttention v2.8.3:
+ * https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
+ * Sources: kernel_traits.h, utils.h, softmax.h, flash_fwd_kernel.h.
+ * Copyright (c) 2023-2024, Tri Dao.
+ *
+ * BSD 3-Clause License
+ *
+ * Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice, this
+ *   list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of the copyright holder nor the names of its
+ *   contributors may be used to endorse or promote products derived from
+ *   this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#ifndef __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__
+#define __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__
+
+namespace sd_fattn_cute {
+    struct WindowRegion {
+        int left, top, width, height;
+    };
+
+    static __device__ WindowRegion window_region(const SDWindowAttentionParams& p, int plane, int x, int y) {
+        int w = p.window_width[plane], h = p.window_height[plane];
+        if (p.preserve_boundary) {
+            w = min(w, p.width);
+            h = min(h, p.height);
+            return {min(max(x - w / 2, 0), p.width - w), min(max(y - h / 2, 0), p.height - h), w, h};
+        }
+        int left = max(x - w / 2, 0), top = max(y - h / 2, 0);
+        return {left, top, min(p.width, x + w / 2) - left, min(p.height, y + h / 2) - top};
+    }
+
+    template <typename Dst, typename Id, typename Index>
+    __device__ __forceinline__ void window_copy(const cutlass::half_t* source, Dst& dst, const Id& id, int dim, Index index) {
+#pragma unroll
+        for (int m = 0; m < size<1>(dst); ++m) {
+            const int row = index(get<0>(id(0, m, 0)));
+#pragma unroll
+            for (int k = 0; k < size<2>(dst); ++k) {
+                if (row >= 0) {
+                    const auto* src        = source + int64_t(row) * dim + get<1>(id(0, m, k));
+                    const unsigned address = unsigned(__cvta_generic_to_shared(&dst(0, m, k)));
+                    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
+                                 :
+                                 : "r"(address), "l"(src));
+                } else {
+                    cute::clear(dst(_, m, k));
+                }
+            }
+        }
+    }
+
+    template <int D>
+    __global__ __launch_bounds__(128) void window_kernel(const cutlass::half_t* q, const cutlass::half_t* k, const cutlass::half_t* v, float* output, SDWindowAttentionParams p, int kv_heads) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+        constexpr int BM = 64, BN = 64, TW = 8;
+        using T       = Traits<BM, BN, D>;
+        using Element = typename T::Element;
+        extern __shared__ char storage[];
+        const int tid = threadIdx.x, head = blockIdx.z;
+        const int qx = blockIdx.x * TW, qy = blockIdx.y * (BM / TW);
+        const int tokens = p.width * p.height;
+        q += int64_t(head) * tokens * D;
+        const int64_t kv_offset = int64_t(head / (p.heads / kv_heads)) * tokens * p.planes * D;
+        k += kv_offset;
+        v += kv_offset;
+        auto sQ      = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
+        auto sK      = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
+        auto sV      = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
+        auto sVt     = make_tensor(sV.data(), typename T::VtLayout{});
+        auto sVplain = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
+        typename T::GlobalCopy global_copy;
+        auto thread_copy = global_copy.get_thread_slice(tid);
+        auto tQsQ        = thread_copy.partition_D(sQ);
+        auto tKsK        = thread_copy.partition_D(sK);
+        auto tVsV        = thread_copy.partition_D(sV);
+        auto q_id        = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+        auto kv_id       = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
+        auto query_index = [&](int row) {
+            const int x = qx + row % TW, y = qy + row / TW;
+            return x < p.width && y < p.height ? y * p.width + x : -1;
+        };
+        window_copy(q, tQsQ, q_id, D, query_index);
+        cute::cp_async_fence();
+
+        typename T::Mma mma;
+        auto thread_mma = mma.get_thread_slice(tid);
+        auto rQ         = thread_mma.partition_fragment_A(sQ);
+        auto rK         = thread_mma.partition_fragment_B(sK);
+        auto rV         = thread_mma.partition_fragment_B(sVplain);
+        auto acc_o      = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
+        auto copyQ      = make_tiled_copy_A(typename T::CopyAtom{}, mma);
+        auto copyK      = make_tiled_copy_B(typename T::CopyAtom{}, mma);
+        auto copyV      = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
+        auto threadQ    = copyQ.get_thread_slice(tid);
+        auto threadK    = copyK.get_thread_slice(tid);
+        auto threadV    = copyV.get_thread_slice(tid);
+        auto tSsQ       = threadQ.partition_S(sQ);
+        auto tSsK       = threadK.partition_S(sK);
+        auto tOsV       = threadV.partition_S(sVt);
+        auto score_id   = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
+        clear(acc_o);
+        Softmax<2 * size<1>(acc_o), true> softmax;
+        const float scale = rsqrtf(float(D));
+        for (int plane = 0; plane < p.planes; ++plane) {
+            const auto first = window_region(p, plane, qx, qy);
+            const auto last  = window_region(p, plane, min(qx + TW - 1, p.width - 1), min(qy + BM / TW - 1, p.height - 1));
+            const int kw     = last.left + last.width - first.left;
+            const int kh     = last.top + last.height - first.top;
+            for (int start = 0; start < kw * kh; start += BN) {
+                auto key_index = [&](int row) {
+                    row += start;
+                    return row < kw * kh ? plane * tokens + (first.top + row / kw) * p.width + first.left + row % kw : -1;
+                };
+                window_copy(k, tKsK, kv_id, D, key_index);
+                window_copy(v, tVsV, kv_id, D, key_index);
+                cute::cp_async_fence();
+                cute::cp_async_wait<0>();
+                __syncthreads();
+                auto acc_s = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
+                clear(acc_s);
+                gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+#pragma unroll
+                for (int i = 0; i < size(acc_s); ++i) {
+                    const int qr = get<0>(score_id(i)), kr = start + get<1>(score_id(i));
+                    const int x = qx + qr % TW, y = qy + qr / TW;
+                    const auto region = window_region(p, plane, x, y);
+                    const int kx = first.left + kr % kw, ky = first.top + kr / kw;
+                    const bool valid = x < p.width && y < p.height && kr < kw * kh &&
+                                       kx >= region.left && kx < region.left + region.width &&
+                                       ky >= region.top && ky < region.top + region.height;
+                    acc_s(i)         = valid ? acc_s(i) * scale : -INFINITY;
+                }
+                if (plane == 0 && start == 0)
+                    softmax.template update<true>(acc_s, acc_o);
+                else
+                    softmax.template update<false>(acc_s, acc_o);
+                auto rP   = to_half(acc_s);
+                auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
+                gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
+                __syncthreads();
+            }
+        }
+        softmax.normalize(acc_o);
+        auto output_id = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+#pragma unroll
+        for (int i = 0; i < size(acc_o); ++i) {
+            const int token = query_index(get<0>(output_id(i)));
+            if (token >= 0)
+                output[(int64_t(token) * p.heads + head) * D + get<1>(output_id(i))] = acc_o(i);
+        }
+#else
+        // This translation unit also targets older GPUs; runtime dispatch excludes this kernel there.
+        asm volatile("trap;");
+#endif
+    }
+}
+
+static bool window_attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+    static const bool disabled = [] {
+        const char* value = std::getenv("GGML_CUDA_DISABLE_WINDOW_ATTENTION");
+        return value && std::atoi(value) != 0;
+    }();
+    SDWindowAttentionParams params;
+    const auto* info = cuda_device(device);
+    return !disabled && info && info->properties.major >= 8 && sd_ggml_window_attn_describe(node, &params);
+}
+
+static ggml_status window_attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+    *workspace = 0;
+    for (int i = 0; i < 3; ++i)
+        *workspace += (size_t(ggml_nelements(nodes[0]->src[i])) * sizeof(half) + 255) / 256 * 256;
+    return GGML_STATUS_SUCCESS;
+}
+
+static void window_attention_destroy(void*, void*) {}
+
+template <int D>
+static cudaError_t window_attention_launch(const ggml_ext_cuda_context& context, const half* const* inputs, float* output, SDWindowAttentionParams params, int kv_heads) {
+    auto kernel                = sd_fattn_cute::window_kernel<D>;
+    constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
+    static std::mutex mutex;
+    static std::unordered_map<int, bool> initialized;
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        if (!initialized[context.device]) {
+            auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
+            if (error != cudaSuccess)
+                return error;
+            initialized[context.device] = true;
+        }
+    }
+    kernel<<<dim3((params.width + 7) / 8, (params.height + 7) / 8, params.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(
+        reinterpret_cast<const cutlass::half_t*>(inputs[0]), reinterpret_cast<const cutlass::half_t*>(inputs[1]),
+        reinterpret_cast<const cutlass::half_t*>(inputs[2]), output, params, kv_heads);
+    return cudaGetLastError();
+}
+
+static ggml_status window_attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+    auto device = ggml_backend_get_device(execution->backend);
+    auto reg    = ggml_backend_dev_backend_reg(device);
+    auto query  = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    SDWindowAttentionParams params;
+    auto* node = nodes[0];
+    if (!sd_ggml_window_attn_describe(node, &params))
+        return GGML_STATUS_FAILED;
+    auto stream   = static_cast<cudaStream_t>(context.stream);
+    auto* base    = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+    const int dim = int(node->ne[0] / params.heads);
+    const half* inputs[3];
+    for (int i = 0; i < 3; ++i) {
+        const auto* src     = node->src[i];
+        auto* packed        = reinterpret_cast<half*>(base);
+        const int64_t heads = src->ne[0] / dim;
+        copy_to_half<float><<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), packed,
+                                                                                   ggml_nelements(src), dim, src->ne[1], heads, src->nb[1], dim * sizeof(float), ggml_nbytes(src), 1.f);
+        if (cudaGetLastError() != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        inputs[i] = packed;
+        base += (size_t(ggml_nelements(src)) * sizeof(half) + 255) / 256 * 256;
+    }
+    const int kv_heads = int(node->src[1]->ne[0] / dim);
+    auto* output       = static_cast<float*>(node->data);
+    const auto error   = dim == 64 ? window_attention_launch<64>(context, inputs, output, params, kv_heads)
+                                   : window_attention_launch<128>(context, inputs, output, params, kv_heads);
+    return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device) {
+    const auto* info = cuda_device(device);
+    if (!info || info->properties.major < 8)
+        return;
+    int previous;
+    if (cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info->ordinal) != cudaSuccess)
+        return;
+    cudaFuncAttributes attributes64{}, attributes128{};
+    const auto error64  = cudaFuncGetAttributes(&attributes64, sd_fattn_cute::window_kernel<64>);
+    const auto error128 = cudaFuncGetAttributes(&attributes128, sd_fattn_cute::window_kernel<128>);
+    const auto restored = cudaSetDevice(previous);
+    // JIT can produce an SM80+ binary from older PTX containing only the trap.
+    if (restored != cudaSuccess || error64 != cudaSuccess || error128 != cudaSuccess ||
+        std::min(attributes64.binaryVersion, attributes64.ptxVersion) < 80 ||
+        std::min(attributes128.binaryVersion, attributes128.ptxVersion) < 80)
+        return;
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.window_attention.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_window_attn_op();
+    desc.priority             = 100;
+    desc.supports             = window_attention_supports;
+    desc.prepare              = window_attention_prepare;
+    desc.destroy              = window_attention_destroy;
+    desc.execute              = window_attention_execute;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
+
+#endif

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,36 +269,70 @@ 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];
+
+static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+    if (!attributes || count < 3)
+        return false;
+    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) ||
+        count != (codebook ? 5 : 3) + attr.has_bias)
+        return false;
+    for (int i = 0; i < count; ++i)
+        if (!inputs[i] || !ggml_is_contiguous(inputs[i]))
+            return false;
+    const auto* w = inputs[0];
+    const auto* x = inputs[1];
+    if (w->type != GGML_TYPE_I8 || w->ne[0] <= 0 || w->ne[0] > i8_size_limit / 2 ||
+        w->ne[1] <= 0 || w->ne[2] != 1 || w->ne[3] != 1)
+        return false;
+    const int64_t k = 2 * w->ne[0], n = w->ne[1];
+    int64_t rows = 1;
+    for (auto dim : attr.shape) {
+        if (dim <= 0 || dim > i8_size_limit / rows)
+            return false;
+        rows *= dim;
     }
-    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);
+    int64_t height;
+    if (!packed_height(k, rows, &height) || n > i8_size_limit / 4 / rows ||
+        x->type != GGML_TYPE_I8 || x->ne[0] != k || x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1 ||
+        !scale_matches(inputs[2], codebook ? 16 : n) ||
+        (attr.has_bias && !scale_matches(inputs[count - 1], n)))
+        return false;
+    if (codebook && (k % 16 || !scale_matches(inputs[3], n) || inputs[4]->type != GGML_TYPE_F8_E4M3 ||
+                     inputs[4]->ne[0] != k / 16 || inputs[4]->ne[1] != n || inputs[4]->ne[2] != 1 || inputs[4]->ne[3] != 1))
+        return false;
+    out->type  = GGML_TYPE_F32;
+    out->ne[0] = n;
+    std::memcpy(out->ne + 1, attr.shape, sizeof(attr.shape));
+    return true;
 }
-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);
+
+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(SDW4Attributes), w4_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return 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;
+    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);
+    ggml_ext_tensor_desc out;
+    if (!w4_infer(node->src, count, &attr, &out) || node->type != out.type ||
+        std::memcmp(node->ne, out.ne, sizeof(out.ne)))
+        return false;
+    *op = {node->src[0], node->src[1], node->src[2], codebook ? node->src[3] : nullptr,
+           codebook ? node->src[4] : nullptr, attr.has_bias ? node->src[count - 1] : nullptr, attr.kind};
+    return true;
 }
 
 static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
@@ -348,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;
@@ -400,10 +428,56 @@ 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);
+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;
+    std::memcpy(&p, attributes, sizeof(p));
+    if (p.width <= 0 || p.height <= 0 || int64_t(p.width) * p.height > INT32_MAX / 4 ||
+        p.heads <= 0 || p.planes < 1 || p.planes > 4 || (p.preserve_boundary != 0 && p.preserve_boundary != 1))
+        return false;
+    for (int i = 0; i < p.planes; ++i)
+        if (p.window_height[i] <= 0 || p.window_width[i] <= 0 ||
+            (!p.preserve_boundary && (p.window_height[i] % 2 || p.window_width[i] % 2)))
+            return false;
+    for (int i = 0; i < 3; ++i) {
+        const auto* t = inputs[i];
+        if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->ne[2] != 1 || t->ne[3] != 1 ||
+            t->ne[1] != int64_t(p.width) * p.height * (i ? p.planes : 1))
+            return false;
+    }
+    const auto* q = inputs[0];
+    const auto* k = inputs[1];
+    if (q->ne[0] % p.heads || (q->ne[0] / p.heads != 64 && q->ne[0] / p.heads != 128) ||
+        k->ne[0] <= 0 || k->ne[0] > q->ne[0] || q->ne[0] % k->ne[0] ||
+        k->ne[0] % (q->ne[0] / p.heads) || !ggml_are_same_shape(k, inputs[2]))
+        return false;
+    *out = {GGML_TYPE_F32, {q->ne[0], q->ne[1], 1, 1}};
+    return true;
+}
+
+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), sd_ggml_window_attn_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params) {
+    if (!node || ggml_ext_op_get(node) != sd_ggml_window_attn_op() || node->src[3] || !ggml_is_contiguous(node))
+        return false;
+    const auto* attributes = ggml_ext_op_attributes(node);
+    ggml_ext_tensor_desc out;
+    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;
 }
 
 void sd_ggml_register_extensions(ggml_backend_t backend) {
@@ -429,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,12 +38,38 @@ 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;
+    const ggml_tensor* weight_scale;
+    const ggml_tensor* s_channel;
+    const ggml_tensor* s_rel;
+    const ggml_tensor* bias;
+    SDW4Kind kind;
+};
+
+ggml_ext_op_t sd_ggml_w4_op();
+
+bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation);
 
 enum ggml_sage_attn_mode {
     GGML_SAGE_ATTN_AUTO = 0,
@@ -51,20 +83,33 @@ 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;
     float tau;
 };
 
+struct SDWindowAttentionParams {
+    int32_t width, height, heads, planes;
+    int16_t window_height[4], window_width[4];
+    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);
+
 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());
@@ -211,6 +195,27 @@ ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
     return nullptr;
 }
 
+static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
+    // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
+    if (sd_backend_is(backend, "Vulkan")) {
+        return true;
+    }
+    auto buft        = ggml_backend_get_default_buffer_type(backend);
+    size_t alignment = ggml_backend_buft_get_alignment(buft);
+    size_t max_size  = ggml_backend_buft_get_max_size(buft);
+    max_size -= max_size % alignment;
+    if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
+        return false;
+    }
+    // Cross-backend copies must also fit the destination's allocation limit.
+    for (auto src : node->src) {
+        if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
+            return false;
+        }
+    }
+    return true;
+}
+
 void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
     if (sched == nullptr || gf == nullptr) {
         return;
@@ -235,7 +240,12 @@ void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph*
             node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
             continue;
         }
-        if (ggml_backend_supports_op(current, node)) {
+        if (!backend_can_allocate_op(current, node)) {
+            auto cpu = workspace_.cpu_backend();
+            if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
+                ggml_backend_sched_set_tensor_backend(sched, node, cpu);
+            }
+        } else if (ggml_backend_supports_op(current, node)) {
             ggml_backend_sched_set_tensor_backend(sched, node, current);
         }
     }

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;
             }
@@ -206,7 +265,7 @@ namespace sd {
         for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
             ggml_backend_t backend = backend_by_segment[seg_idx];
             const auto& segment    = plan.segments[seg_idx];
-            auto assign_node = [&](int node_index) {
+            auto assign_node       = [&](int node_index) {
                 if (node_index < 0 || node_index >= n_nodes) {
                     return;
                 }

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/core/tensor.hpp

@@ -222,14 +222,28 @@ namespace sd {
                 return result;
             }
 
-            for (int64_t flat = 0; flat < result.numel(); ++flat) {
-                std::vector<int64_t> out_coord = tensor_unravel_index(flat, out_shape);
-                std::vector<int64_t> src_coord(static_cast<size_t>(dim()), 0);
-                for (size_t i = 0; i < dims.size(); ++i) {
-                    src_coord[dims[i]] = out_coord[i];
+            std::vector<int64_t> source_strides(dims.size(), 1), strides(dims.size());
+            for (size_t i = 1; i < dims.size(); ++i)
+                source_strides[i] = source_strides[i - 1] * shape_[i - 1];
+            for (size_t i = 0; i < dims.size(); ++i)
+                strides[i] = source_strides[dims[i]];
+            parallel_for(0, result.numel(), 65536, [&](int64_t begin, int64_t end) {
+                auto coord     = tensor_unravel_index(begin, out_shape);
+                int64_t source = 0;
+                for (size_t i = 0; i < dims.size(); ++i)
+                    source += coord[i] * strides[i];
+                for (int64_t flat = begin; flat < end; ++flat) {
+                    result.data()[flat] = data_[source];
+                    for (size_t i = 0; i < dims.size(); ++i) {
+                        ++coord[i];
+                        source += strides[i];
+                        if (coord[i] < out_shape[i])
+                            break;
+                        source -= coord[i] * strides[i];
+                        coord[i] = 0;
+                    }
                 }
-                result[flat] = index(src_coord);
-            }
+            });
 
             return result;
         }
@@ -1037,11 +1051,13 @@ namespace sd {
             auto x_contributors = make_interpolate_contributors(input_width, output_width, mode, antialias);
             auto y_contributors = make_interpolate_contributors(input_height, output_height, mode, antialias);
 
-            for (int64_t plane = 0; plane < plane_count; ++plane) {
-                const int64_t input_plane_offset  = plane * input_plane;
-                const int64_t output_plane_offset = plane * output_plane;
-                for (int64_t y = 0; y < output_height; ++y) {
-                    const auto& y_axis = y_contributors[static_cast<size_t>(y)];
+            parallel_for(0, plane_count * output_height, std::max<int64_t>(1, 16384 / output_width), [&](int64_t begin, int64_t end) {
+                for (int64_t row = begin; row < end; ++row) {
+                    const int64_t plane               = row / output_height;
+                    const int64_t y                   = row % output_height;
+                    const int64_t input_plane_offset  = plane * input_plane;
+                    const int64_t output_plane_offset = plane * output_plane;
+                    const auto& y_axis                = y_contributors[static_cast<size_t>(y)];
                     for (int64_t x = 0; x < output_width; ++x) {
                         const auto& x_axis = x_contributors[static_cast<size_t>(x)];
                         double value       = 0.0;
@@ -1055,7 +1071,7 @@ namespace sd {
                         output.data()[output_plane_offset + y * output_width + x] = static_cast<T>(value);
                     }
                 }
-            }
+            });
 
             return output;
         }

src/core/util.cpp

@@ -807,13 +807,11 @@ sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
                                                  static_cast<int64_t>(image.height),
                                                  static_cast<int64_t>(image.channel),
                                                  1});
-    for (uint32_t iw = 0; iw < image.width; ++iw) {
-        for (uint32_t ih = 0; ih < image.height; ++ih) {
-            for (uint32_t ic = 0; ic < image.channel; ++ic) {
-                tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale);
-            }
-        }
-    }
+    const int64_t pixels     = int64_t(image.width) * image.height;
+    sd::tensor_for_each(tensor.numel(), [&](int64_t i) {
+        float value      = image.data[(i % pixels) * image.channel + i / pixels];
+        tensor.data()[i] = scale ? value / 255.f : value;
+    });
     if (target_width >= 0 && target_height >= 0 &&
         (tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) {
         tensor = sd::ops::interpolate(tensor,

src/core/vulkan/CMakeLists.txt

@@ -49,9 +49,16 @@ endfunction()
 # and is disabled for KHR cooperative matrix shaders.
 sd_vulkan_shader(rope_apply rope_apply ON)
 sd_vulkan_shader(modulate modulate ON -O)
+sd_vulkan_shader(rms_modulate rms_modulate ON -O)
 sd_vulkan_shader(rms_rope_apply rms_rope_apply ON)
 sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16)
 sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(quantize_i8_convrot_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64)
+sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(mul_mat_w4a4_convrot_cm2 mul_mat_w4a4_convrot_cm2 ${SD_VK_CM2} -O)
+sd_vulkan_shader(mul_mat_w4a4_convrot_decode_cm2 mul_mat_w4a4_convrot_decode_cm2 ${SD_VK_CM2} -O)
+sd_vulkan_shader(mul_mat_w4a8_convrot_cm2 mul_mat_w4a8_convrot_cm2 ${SD_VK_CM2} -O)
 sd_vulkan_shader(quantize_i8_convrot_subgroup quantize_i8_convrot_subgroup ${SD_VK_DOT} -O)
 sd_vulkan_shader(quantize_i8_convrot_register quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSINGLE_PASS=1)
 sd_vulkan_shader(swiglu_quantize_i8_convrot quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1)
@@ -66,5 +73,5 @@ sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_C
 sd_vulkan_shader(mul_mat_i8_residual_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DRESIDUAL=1)
 add_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS})
 add_dependencies(${SD_LIB} sd-vulkan-shaders)
-target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp")
+target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/rms_affine.cpp")
 target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}")

src/core/vulkan/int8.cpp

@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <array>
 #include <climits>
+#include <cstdlib>
 #include <cstring>
 #include <vector>
 #include "core/ggml_extensions.h"
@@ -11,7 +12,13 @@
 #include "mul_mat_i8_tensorwise.h"
 #include "mul_mat_i8_tensorwise_cm1.h"
 #include "mul_mat_i8_tensorwise_cm2.h"
+#include "mul_mat_w4a4_convrot.h"
+#include "mul_mat_w4a4_convrot_cm2.h"
+#include "mul_mat_w4a4_convrot_decode_cm2.h"
+#include "mul_mat_w4a8_convrot.h"
+#include "mul_mat_w4a8_convrot_cm2.h"
 #include "quantize_i8_convrot.h"
+#include "quantize_i8_convrot_h64.h"
 #include "quantize_i8_convrot_register.h"
 #include "quantize_i8_convrot_subgroup.h"
 #include "rms_modulate_convrot.h"
@@ -20,11 +27,17 @@
 #include "swiglu_quantize_i8_convrot_16384.h"
 
 namespace sd_vulkan_i8 {
+    // Keep these budgets in sync with the cooperative matrix shader tile layouts.
+    static constexpr uint32_t w4a4_shared       = 2 * (64 + 64) * (64 / 4 + 4) * 4;
+    static constexpr uint32_t w4a4_large_shared = 2 * (64 + 128) * (64 / 4 + 4) * 4;
+    static constexpr uint32_t w4a8_shared       = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 64 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
+    static constexpr uint32_t w4a8_large_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 128 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
     struct Device {
         const ggml_ext_vulkan_api* api;
         ggml_ext_vulkan_device_info info;
         ggml_ext_vulkan_compute_info compute;
         bool subgroup, cm2, cm2_small;
+        bool w4a4_cm2, w4a8_cm2;
         uint32_t cm_m, cm_n, cm_k;
     };
     static bool device_info(ggml_backend_dev_t dev, Device& d) {
@@ -52,6 +65,10 @@ namespace sd_vulkan_i8 {
                 d.cm2 = true;
                 d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
             }
+            if (prop.flexible && prop.workgroup_scope && prop.invocations == 256 && d.info.max_group_size[0] >= 256 && 64 % prop.m == 0 && 64 % prop.n == 0) {
+                d.w4a4_cm2 |= d.compute.shared_memory >= w4a4_shared && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_cm2) > 4;
+                d.w4a8_cm2 |= d.compute.shared_memory >= w4a8_shared && 32 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot_cm2) > 4;
+            }
         }
         return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
     }
@@ -111,6 +128,23 @@ namespace sd_vulkan_i8 {
                (cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
                d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
     }
+    static bool w4_supports(ggml_backend_dev_t dev, const ggml_tensor* dst, void*) {
+        Device d;
+        SDW4Operation op;
+        if (!device_info(dev, d) || !sd_ggml_w4_describe(dst, &op) || !bounded(d, dst) ||
+            sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot) <= 4 || sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot) <= 4 ||
+            d.compute.shared_memory < (op.kind == SDW4Kind::W4A8 ? 16452u : 13312u) || d.info.max_push_constants < 16)
+            return false;
+        for (auto* t : dst->src)
+            if (t && (!bounded(d, t) || ggml_nbytes(t) > UINT32_MAX))
+                return false;
+        const uint64_t k = op.input->ne[0], n = dst->ne[0], rows = ggml_nrows(dst);
+        return k % 64 == 0 && k <= UINT32_MAX && n <= UINT32_MAX && rows <= UINT32_MAX &&
+               ggml_nelements(dst) <= UINT32_MAX &&
+               (k * ((rows + 3) / 4 * 4)) % d.info.buffer_alignment == 0 &&
+               (n + 63) / 64 <= d.info.max_group_count[0] && (rows + 31) / 32 <= d.info.max_group_count[1] &&
+               d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
+    }
     static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
         if (i + 2 >= count_nodes(graph) || !device.subgroup) {
             return false;
@@ -534,9 +568,12 @@ namespace sd_vulkan_i8 {
             status                     = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
         } else {
             ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
-            auto pipeline                      = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
-            bool subgroup                      = d.subgroup && k >= 1024;
-            uint32_t width                     = k;
+            int group;
+            if (!sd_ggml_convrot_describe(output, &group))
+                return GGML_STATUS_FAILED;
+            auto pipeline  = group == 64 ? SD_VK_PIPELINE(quantize_i8_convrot_h64, 3, QuantParameters) : SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
+            bool subgroup  = group == 256 && d.subgroup && k >= 1024;
+            uint32_t width = k;
             if (subgroup) {
                 pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
                 if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
@@ -653,6 +690,78 @@ namespace sd_vulkan_i8 {
         }
         return finish(e, d, output, dst, alias, status);
     }
+    static ggml_status multiply_w4(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* output) {
+        SDW4Operation op;
+        if (!sd_ggml_w4_describe(output, &op))
+            return GGML_STATUS_FAILED;
+        const uint32_t k = uint32_t(op.input->ne[0]), n = uint32_t(output->ne[0]), rows = uint32_t(ggml_nrows(output));
+        const size_t scale_offset                     = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
+        const ggml_ext_vulkan_binding xs              = {op.input, nullptr, scale_offset, size_t(rows) * 4};
+        std::vector<ggml_ext_vulkan_binding> bindings = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}};
+        if (op.kind == SDW4Kind::W4A8) {
+            bindings.push_back({op.s_channel});
+            bindings.push_back({op.s_rel});
+        }
+        bindings.push_back(xs);
+        bindings.push_back({op.bias ? op.bias : op.weight_scale});
+        ggml_ext_vulkan_binding dst;
+        bool alias;
+        if (!output_binding(e, d, output, bindings, dst, alias))
+            return GGML_STATUS_FAILED;
+        bindings.push_back(dst);
+        struct Parameters {
+            uint32_t n, rows, k, bias;
+        };
+        const Parameters pc = {n, rows, k, op.bias ? 1u : 0u};
+        auto pipeline       = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
+        uint32_t tile_cols = 64, tile_rows = 32;
+        const bool large_w4a4      = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared;
+        const bool large_w4a8      = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared;
+        const bool decode_w4a4     = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 &&
+                                     aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4;
+        const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64};
+        const char* force_dp4a     = std::getenv("GGML_VK_W4_FORCE_DP4A");
+        if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) &&
+            (op.kind == SDW4Kind::W4A4 ? d.w4a4_cm2 : d.w4a8_cm2 && aligned(op.weight, 8))) {
+            pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot_cm2, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot_cm2, 8, Parameters);
+            if (decode_w4a4) {
+                struct DecodeParameters {
+                    uint32_t n, rows, k, bias, row_offset;
+                };
+                DecodeParameters decoded    = {n, rows, k, pc.bias, 0};
+                uint32_t decode_constants[] = {64, 128, 128, 1};
+                pipeline                    = SD_VK_PIPELINE(mul_mat_w4a4_convrot_decode_cm2, 6, DecodeParameters);
+                pipeline.constants          = decode_constants;
+                pipeline.constant_count     = 4;
+                uint32_t groups[]           = {n / 64, rows / 128, 1};
+                auto status                 = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
+                if (status == GGML_STATUS_SUCCESS && rows % 128) {
+                    // Specialize the tail separately so bounds checks do not inhibit full-tile loads.
+                    decoded.row_offset  = rows / 128 * 128;
+                    decode_constants[1] = rows % 128 <= 64 ? 64 : 128;
+                    decode_constants[3] = 0;
+                    groups[1]           = 1;
+                    status              = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
+                }
+                return finish(e, d, output, dst, alias, status);
+            }
+            pipeline.constants      = constants;
+            pipeline.constant_count = 3;
+            tile_cols               = constants[0];
+            tile_rows               = constants[1];
+        }
+        const uint32_t groups[] = {(n + tile_cols - 1) / tile_cols, (rows + tile_rows - 1) / tile_rows, 1};
+        auto status             = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
+        return finish(e, d, output, dst, alias, status);
+    }
+    static bool quant64_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
+        Device d;
+        int group;
+        return device_info(dev, d) && sd_ggml_convrot_describe(n, &group) && group == 64 &&
+               bounded(d, n) && bounded(d, n->src[0]) && n->ne[0] <= UINT32_MAX &&
+               ggml_nrows(n->src[0]) <= UINT32_MAX && d.info.max_group_invocations >= 16 && d.info.max_group_size[0] >= 16 &&
+               (uint64_t(n->ne[0]) * ((ggml_nrows(n->src[0]) + 3) / 4 * 4)) % d.info.buffer_alignment == 0;
+    }
     static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
         Device d;
         return device_info(dev, d) && can_quantize(d, n);
@@ -747,6 +856,8 @@ namespace sd_vulkan_i8 {
         Device d;
         if (!device_info(ggml_backend_get_device(e->backend), d))
             return GGML_STATUS_FAILED;
+        if (ggml_ext_op_get(nodes[0]) == sd_ggml_w4_op())
+            return multiply_w4(e, d, nodes[0]);
         return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
     }
 #undef SD_VK_PIPELINE
@@ -768,6 +879,13 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
     d.custom_op            = sd_ggml_convrot_op();
     d.supports             = quant_supports;
     GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.name     = "sd.quantize_i8_convrot.h64.vulkan.v1";
+    d.supports = quant64_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.name      = "sd.linear_w4_convrot.vulkan.v1";
+    d.custom_op = sd_ggml_w4_op();
+    d.supports  = w4_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
     d.name      = "sd.linear_i8.vulkan.v1";
     d.op        = GGML_OP_CUSTOM;
     d.custom_op = sd_ggml_i8_op();

src/core/vulkan/rms_affine.cpp

@@ -0,0 +1,192 @@
+#include <algorithm>
+#include <climits>
+#include <cstring>
+#include "core/ggml_extensions.h"
+#include "ggml-ext-vulkan.h"
+#include "rms_modulate.h"
+
+namespace sd_vulkan_rms_affine {
+    struct Parameters {
+        uint32_t ne[4], stride[4], offsets[5];
+        float eps;
+        uint32_t ne10;
+    };
+    static_assert(sizeof(Parameters) == 60, "RMS affine push constant layout");
+
+    static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) {
+        auto query = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(
+            ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1"));
+        return query ? query() : nullptr;
+    }
+    static bool bounded(const ggml_tensor* t, const ggml_ext_vulkan_device_info& info) {
+        if (!t || t->type != GGML_TYPE_F32 || t->nb[0] != 4 || ggml_nbytes(t) > info.max_buffer_range - (info.buffer_alignment - 1))
+            return false;
+        size_t offset     = 0;
+        const auto* owner = t;
+        while (!owner->data && owner->view_src) {
+            offset += owner->view_offs;
+            owner = owner->view_src;
+        }
+        if ((reinterpret_cast<uintptr_t>(owner->data) + offset) % 4)
+            return false;
+        for (int d = 0; d < 4; ++d)
+            if (t->ne[d] > UINT32_MAX || t->nb[d] % 4 || t->nb[d] / 4 > UINT32_MAX)
+                return false;
+        return true;
+    }
+    static const ggml_tensor* other(const ggml_tensor* node, const ggml_tensor* input) {
+        return node->src[0] == input ? node->src[1] : node->src[1] == input ? node->src[0]
+                                                                            : nullptr;
+    }
+    struct Operation {
+        const ggml_tensor* x;
+        const ggml_tensor* weight;
+        const ggml_tensor* scale;
+        const ggml_tensor* shift;
+        const ggml_tensor* output;
+    };
+    static Operation describe(ggml_tensor* const* nodes, int count) {
+        Operation op = {nodes[0]->src[0], other(nodes[1], nodes[0]), nullptr, nullptr, nodes[count - 1]};
+        if (count == 3) {
+            op.scale = nodes[2]->src[1];
+            op.shift = nodes[2]->src[2];
+        } else if (count == 5) {
+            op.scale = other(nodes[2], nodes[1]);
+            op.shift = other(nodes[4], nodes[3]);
+        }
+        return op;
+    }
+    static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+        const auto* service = api(device);
+        ggml_ext_vulkan_device_info info;
+        auto* g = const_cast<ggml_cgraph*>(graph);
+        if (!service || !service->device_info(device, &info) || info.fusion_disabled || first + 1 >= ggml_graph_n_nodes(g))
+            return 0;
+        ggml_tensor* nodes[5] = {};
+        const int available   = std::min(5, ggml_graph_n_nodes(g) - first);
+        for (int i = 0; i < available; ++i)
+            nodes[i] = ggml_graph_node(g, first + i);
+        if (nodes[0]->op != GGML_OP_RMS_NORM || nodes[1]->op != GGML_OP_MUL)
+            return 0;
+        const auto* x = nodes[0]->src[0];
+        if (!bounded(x, info) || x->ne[0] != 3072 || available < 3 || !ggml_is_contiguous(x) ||
+            x->ne[1] <= 1 || x->ne[1] > info.max_group_count[0] || x->ne[2] != 1 || x->ne[3] != 1 ||
+            info.max_group_size[0] < 512 || info.max_group_invocations < 512)
+            return 0;
+        int count;
+        if (sd_ggml_modulate_describe(nodes[2]) && nodes[2]->src[0] == nodes[1]) {
+            count = 3;
+        } else if (available == 5 && nodes[2]->op == GGML_OP_MUL && nodes[3]->op == GGML_OP_ADD &&
+                   other(nodes[3], nodes[1]) == nodes[2] && nodes[4]->op == GGML_OP_ADD) {
+            count = 5;
+        } else {
+            return 0;
+        }
+        const auto op = describe(nodes, count);
+        if (!op.weight || !op.scale || (count == 5 && !op.shift))
+            return 0;
+        for (const auto* parameter : {op.weight, op.scale, op.shift}) {
+            if (!parameter)
+                continue;
+            if (!bounded(parameter, info) || !ggml_is_contiguous(parameter) || parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0])
+                return 0;
+            for (int i = 0; i < count; ++i)
+                if (parameter == nodes[i])
+                    return 0;
+        }
+        for (int i = 0; i < count; ++i)
+            if (!bounded(nodes[i], info) || !ggml_are_same_shape(nodes[i], x) || !ggml_is_contiguous(nodes[i]))
+                return 0;
+        return count;
+    }
+    static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+        const auto* service = api(device);
+        ggml_ext_vulkan_device_info info;
+        if (!service || !service->device_info(device, &info) || info.fusion_disabled)
+            return;
+        for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
+            const auto* node = ggml_graph_node(graph, i);
+            if (!sd_ggml_modulate_describe(node) && node->op != GGML_OP_MUL)
+                continue;
+            const auto* weighted = node->src[0];
+            if (!weighted || weighted->op != GGML_OP_MUL || !weighted->src[0] ||
+                weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 3072)
+                continue;
+            for (int j = i - 2; j >= std::max(0, i - 32); --j) {
+                if (ggml_graph_node(graph, j) == weighted->src[0] && ggml_graph_node(graph, j + 1) == weighted) {
+                    if (j + 2 < i)
+                        ggml_ext_graph_delay(graph, j, 2, i);
+                    break;
+                }
+            }
+        }
+    }
+    static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[count - 1]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) {
+        auto device         = ggml_backend_get_device(e->backend);
+        const auto* service = api(device);
+        ggml_ext_vulkan_device_info info;
+        if (!service || !service->device_info(device, &info))
+            return GGML_STATUS_FAILED;
+        const auto op                      = describe(nodes, count);
+        ggml_ext_vulkan_binding bindings[] = {{op.x}, {op.weight}, {op.output}, {op.scale}, {op.shift ? op.shift : op.scale}};
+        ggml_ext_vulkan_buffer_info buffers[5];
+        const int n = 5;
+        for (int i = 0; i < n; ++i)
+            if (!service->buffer_info(e->backend, bindings + i, buffers + i))
+                return GGML_STATUS_FAILED;
+        bool alias = false;
+        for (int i = 0; i < n; ++i)
+            if (i != 2)
+                alias |= buffers[i].storage == buffers[2].storage && buffers[i].offset < buffers[2].offset + buffers[2].size &&
+                         buffers[2].offset < buffers[i].offset + buffers[i].size;
+        const auto destination = bindings[2];
+        if (alias) {
+            bindings[2] = {nullptr, e->workspace, 0, ggml_nbytes(op.output)};
+            if (!service->buffer_info(e->backend, bindings + 2, buffers + 2))
+                return GGML_STATUS_FAILED;
+        }
+        Parameters p = {};
+        p.ne10       = uint32_t(op.weight->ne[0]);
+        for (int i = 0; i < 4; ++i) {
+            p.ne[i]     = uint32_t(op.x->ne[i]);
+            p.stride[i] = uint32_t(op.x->nb[i] / 4);
+        }
+        for (int i = 0; i < n; ++i)
+            p.offsets[i] = buffers[i].shader_offset / 4;
+        std::memcpy(&p.eps, nodes[0]->op_params, sizeof(p.eps));
+        const uint32_t has_shift          = op.shift != nullptr;
+        ggml_ext_vulkan_pipeline pipeline = {};
+        pipeline.name                     = "sd.rms_modulate.vulkan";
+        pipeline.spirv                    = sd_vulkan_shaders::rms_modulate;
+        pipeline.spirv_size               = sizeof(sd_vulkan_shaders::rms_modulate);
+        pipeline.binding_count            = n;
+        pipeline.push_constant_size       = sizeof(p);
+        pipeline.constants                = &has_shift;
+        pipeline.constant_count           = 1;
+        uint32_t groups[]                 = {p.ne[1], std::min(p.ne[2], info.max_group_count[1]), std::min(p.ne[3], info.max_group_count[2])};
+        auto status                       = service->dispatch(e->backend, &pipeline, bindings, &p, groups);
+        if (status == GGML_STATUS_SUCCESS && alias)
+            status = service->copy(e->backend, bindings + 2, &destination);
+        return status;
+    }
+}
+void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device) {
+    if (!sd_vulkan_rms_affine::api(device))
+        return;
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.rms_affine.vulkan.v1";
+    desc.op                   = GGML_OP_RMS_NORM;
+    desc.priority             = 150;
+    desc.match                = sd_vulkan_rms_affine::match;
+    desc.optimize             = sd_vulkan_rms_affine::optimize;
+    desc.prepare              = sd_vulkan_rms_affine::prepare;
+    desc.execute              = sd_vulkan_rms_affine::execute;
+    desc.destroy              = sd_vulkan_rms_affine::destroy;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+}

src/core/vulkan/rope.cpp

@@ -72,7 +72,7 @@ namespace sd_vulkan_rope {
             return false;
         const auto* x      = norm->src[0];
         const auto* weight = mul->src[1];
-        return x && weight && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 &&
+        return x && weight && weight != norm && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 &&
                norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_are_same_shape(x, norm) && ggml_are_same_shape(norm, mul) &&
                weight->type == GGML_TYPE_F32 && weight->ne[0] == 128 && ggml_nelements(weight) == 128 &&
                ggml_is_contiguous(weight) && bounded(weight, info) && bounded(x, info) && (x->ne[1] + 3) / 4 <= info.max_group_count[0];
@@ -81,32 +81,35 @@ namespace sd_vulkan_rope {
         ggml_ext_vulkan_device_info info;
         api(device)->device_info(device, &info);
         // Workspace writeback uses Vulkan's four-byte copy granularity.
-        return info.float16 && scale->op == GGML_OP_SCALE && scale->src[0] == rope && scale->type == GGML_TYPE_F32 &&
-               param(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale && copy->type == GGML_TYPE_F16 &&
-               ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(scale, copy);
+        return info.float16 && (!scale || (scale->op == GGML_OP_SCALE && scale->src[0] == rope && scale->type == GGML_TYPE_F32 && param(scale, 1) == 0.f)) &&
+               copy->op == GGML_OP_CPY && copy->src[0] == (scale ? scale : rope) && copy->type == GGML_TYPE_F16 &&
+               ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(rope, copy);
     }
     static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
         auto* g = const_cast<ggml_cgraph*>(graph);
         if (first + 2 >= ggml_graph_n_nodes(g) || !fused_supports(device, ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
             return 0;
-        return first + 4 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), ggml_graph_node(g, first + 3), ggml_graph_node(g, first + 4)) ? 5 : 3;
+        if (first + 4 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), ggml_graph_node(g, first + 3), ggml_graph_node(g, first + 4)))
+            return 5;
+        return first + 3 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), nullptr, ggml_graph_node(g, first + 3)) ? 4 : 3;
     }
     static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
         for (int last = 2; last < ggml_graph_n_nodes(graph); ++last) {
             const auto* output = ggml_graph_node(graph, last);
-            const auto* scale  = output->op == GGML_OP_CPY ? output->src[0] : nullptr;
-            const auto* rope   = scale && scale->op == GGML_OP_SCALE ? scale->src[0] : output;
+            const auto* source = output->op == GGML_OP_CPY ? output->src[0] : nullptr;
+            const auto* scale  = source && source->op == GGML_OP_SCALE ? source : nullptr;
+            const auto* rope   = source ? (scale ? scale->src[0] : source) : output;
             if (!is_rope(rope))
                 continue;
             const auto* mul  = rope->src[0];
             const auto* norm = mul->src[0];
             if (!norm || !fused_supports(device, norm, mul, rope))
                 continue;
-            const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? 5
+            const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? (scale ? 5 : 4)
                                                                                               : 0;
             if (!count)
                 continue;
-            const ggml_tensor* nodes[] = {norm, mul, rope, scale, output};
+            const ggml_tensor* nodes[] = {norm, mul, rope, scale ? scale : output, output};
             int indices[5], next = count - 1;
             for (int i = last; i >= std::max(0, last - 64) && next >= 0; --i)
                 if (ggml_graph_node(graph, i) == nodes[next])
@@ -170,6 +173,7 @@ namespace sd_vulkan_rope {
         pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type));
         std::memcpy(&pc.interleaved, ggml_ext_op_attributes(rope), 4);
         uint32_t groups[3]                = {1, 1, 1};
+        const uint32_t apply_output_scale = count == 5;
         ggml_ext_vulkan_pipeline pipeline = {};
         pipeline.name                     = "sd.rope_apply.vulkan";
         pipeline.spirv                    = sd_vulkan_shaders::rope_apply;
@@ -181,9 +185,13 @@ namespace sd_vulkan_rope {
             groups[0]               = std::min(n_groups, info.max_group_count[0]);
             groups[1]               = (n_groups + groups[0] - 1) / groups[0];
         } else {
-            pipeline.name                   = count == 5 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan";
-            pipeline.spirv                  = count == 5 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply;
-            pipeline.spirv_size             = count == 5 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply);
+            pipeline.name       = count > 3 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan";
+            pipeline.spirv      = count > 3 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply;
+            pipeline.spirv_size = count > 3 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply);
+            if (count > 3) {
+                pipeline.constants      = &apply_output_scale;
+                pipeline.constant_count = 1;
+            }
             pipeline.push_constant_size     = sizeof(FusedParameters);
             pipeline.full_subgroups         = true;
             pipeline.required_subgroup_size = info.subgroup_size_control ? 32 : 0;

src/core/vulkan/shaders/modulate.comp

@@ -1,16 +1,25 @@
 #version 450
 
 layout(local_size_x = 256) in;
-layout(binding = 0) readonly buffer X { float x[]; };
-layout(binding = 1) readonly buffer S { float scale[]; };
-layout(binding = 2) readonly buffer H { float shift[]; };
-layout(binding = 3) writeonly buffer D { float dst[]; };
+layout(binding = 0) readonly buffer X {
+    float x[];
+};
+layout(binding = 1) readonly buffer S {
+    float scale[];
+};
+layout(binding = 2) readonly buffer H {
+    float shift[];
+};
+layout(binding = 3) writeonly buffer D {
+    float dst[];
+};
 layout(push_constant) uniform Parameters {
     uint ne[4], scale_ne[4], shift_ne[4];
     uint x_stride[4], scale_stride[4], shift_stride[4];
     uint offsets[4];
     uint count, has_shift;
-} p;
+}
+p;
 
 void main() {
     uint group = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x;

src/core/vulkan/shaders/modulate_convrot.comp

@@ -2,8 +2,8 @@
 
 #extension GL_EXT_control_flow_attributes : require
 
-layout(constant_id = 0) const uint K      = 4096;
-layout(constant_id = 1) const bool AFFINE = false;
+layout(constant_id = 0) const uint K         = 4096;
+layout(constant_id = 1) const bool AFFINE    = false;
 layout(constant_id = 2) const bool HAS_SHIFT = true;
 layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
 layout(binding = 0) readonly buffer X {

src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp

@@ -0,0 +1,189 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_integer_dot_product : require
+
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    float data_d[];
+};
+
+shared int w_dec[2][16 * 64];
+shared int x_tile[2][32 * 16];
+shared int sum_p[2][16 * 8];
+
+void interleave_pair(int a, int b, uint buf, uint li) {
+    const uint a0       = uint(a);
+    const uint b0       = uint(b);
+    x_tile[buf][li + 0] = int((a0 & 0x000000FFu) | ((a0 >> 8) & 0x0000FF00u) | ((b0 << 16) & 0x00FF0000u) | ((b0 << 8) & 0xFF000000u));
+    x_tile[buf][li + 1] = int(((a0 >> 8) & 0x000000FFu) | ((a0 >> 16) & 0x0000FF00u) | ((b0 << 8) & 0x00FF0000u) | (b0 & 0xFF000000u));
+}
+
+void main() {
+    const uint out_base = gl_WorkGroupID.x * 64u;
+    const uint lx       = gl_LocalInvocationID.x;
+    const uint ly       = gl_LocalInvocationID.y;
+    const uint idx      = gl_LocalInvocationIndex;
+    const uint w_words  = p.k / 8u;
+    const uint a_words  = p.k / 4u;
+    const uint groups   = p.k / 16u;
+
+    const uint dc   = idx & 63u;
+    const uint dw   = idx >> 6u;
+    const uint ocol = out_base + dc;
+
+    const uint fl = lx & 7u;
+    const uint ft = lx >> 3u;
+    const uint fr = ft == 0u ? gl_WorkGroupID.y * 32u + ly : gl_WorkGroupID.y * 32u + ly + 16u;
+    const uint tr = ft * 16u + ly;
+    int my_sum    = 0;
+    int acc0a     = 0;
+    int acc0b     = 0;
+    int acc1a     = 0;
+    int acc1b     = 0;
+    int acc2a     = 0;
+    int acc2b     = 0;
+    int acc3a     = 0;
+    int acc3b     = 0;
+    int acc4a     = 0;
+    int acc4b     = 0;
+    int acc5a     = 0;
+    int acc5b     = 0;
+    int acc6a     = 0;
+    int acc6b     = 0;
+    int acc7a     = 0;
+    int acc7b     = 0;
+
+    {
+        [[unroll]] for (uint j = 0; j < 2; ++j) {
+            const uint pw                    = dw * 2 + j;
+            const uint widx                  = pw;
+            const uint wv                    = ocol < p.n && widx < w_words
+                                                   ? uint(data_w[ocol * w_words + widx])
+                                                   : 0u;
+            w_dec[0][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
+            w_dec[0][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
+        }
+        const uint gi = fl >> 1u;
+        const uint xb = gi * 4u + (fl & 1u) * 2u;
+        const int a   = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
+        const int b   = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
+        interleave_pair(a, b, 0, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
+        my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
+    }
+    barrier();
+
+    for (uint gs = 0; gs < groups; gs += 4) {
+        const uint buf = (gs >> 2) & 1u;
+
+        if (gs + 4 < groups) {
+            const uint nbuf = buf ^ 1u;
+            const uint ngs  = gs + 4;
+            [[unroll]] for (uint j = 0; j < 2; ++j) {
+                const uint pw                       = dw * 2 + j;
+                const uint widx                     = ngs * 2 + pw;
+                const uint wv                       = ocol < p.n && widx < w_words
+                                                          ? uint(data_w[ocol * w_words + widx])
+                                                          : 0u;
+                w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
+                w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
+            }
+            const uint gi = ngs + (fl >> 1u);
+            const uint xb = gi * 4u + (fl & 1u) * 2u;
+            const int a   = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
+            const int b   = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
+            interleave_pair(a, b, nbuf, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
+            my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
+        }
+
+        [[unroll]] for (uint gl = 0; gl < 4; ++gl) {
+            const int u0e  = x_tile[buf][ly * 16u + gl * 4u + 0u];
+            const int u0o  = x_tile[buf][ly * 16u + gl * 4u + 1u];
+            const int u0e2 = x_tile[buf][ly * 16u + gl * 4u + 2u];
+            const int u0o2 = x_tile[buf][ly * 16u + gl * 4u + 3u];
+            const int u1e  = x_tile[buf][(16u + ly) * 16u + gl * 4u + 0u];
+            const int u1o  = x_tile[buf][(16u + ly) * 16u + gl * 4u + 1u];
+            const int u1e2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 2u];
+            const int u1o2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 3u];
+            [[unroll]] for (uint jp = 0; jp < 2; ++jp) {
+                const int ue0 = jp == 0u ? u0e : u0e2;
+                const int uo0 = jp == 0u ? u0o : u0o2;
+                const int ue1 = jp == 0u ? u1e : u1e2;
+                const int uo1 = jp == 0u ? u1o : u1o2;
+                const uint we = (gl * 4u + jp * 2u) * 64u;
+                const uint wo = we + 64u;
+                acc0a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue0));
+                acc1a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue0));
+                acc2a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue0));
+                acc3a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue0));
+                acc0b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo0));
+                acc1b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo0));
+                acc2b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo0));
+                acc3b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo0));
+                acc4a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue1));
+                acc5a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue1));
+                acc6a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue1));
+                acc7a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue1));
+                acc4b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo1));
+                acc5b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo1));
+                acc6b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo1));
+                acc7b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo1));
+            }
+        }
+        barrier();
+    }
+
+    sum_p[ft][ly * 8u + fl] = my_sum;
+    barrier();
+    int sum0 = 0;
+    int sum1 = 0;
+    [[unroll]] for (uint i = 0; i < 8; ++i) {
+        sum0 += sum_p[0][ly * 8u + i];
+        sum1 += sum_p[1][ly * 8u + i];
+    }
+
+    const int accv[2][4] = {
+        {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
+        {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
+    };
+    const int sumv[2] = {sum0, sum1};
+    [[unroll]] for (uint t = 0; t < 2u; ++t) {
+        const uint rt = gl_WorkGroupID.y * 32u + ly + t * 16u;
+        if (rt < p.rows) {
+            [[unroll]] for (uint j = 0; j < 4; ++j) {
+                const uint oj = out_base + j * 16u + lx;
+                if (oj < p.n) {
+                    float value = float(accv[t][j] - 8 * sumv[t]) * data_as[rt] * data_ws[oj];
+                    if (p.has_bias != 0) {
+                        value += data_b[oj];
+                    }
+                    data_d[rt * p.n + oj] = value;
+                }
+            }
+        }
+    }
+}

src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp

@@ -0,0 +1,88 @@
+#version 460
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_KHR_cooperative_matrix : require
+#extension GL_NV_cooperative_matrix2 : require
+#extension GL_KHR_memory_scope_semantics : require
+layout(local_size_x = 256) in;
+layout(constant_id = 0) const uint BM = 64;
+layout(constant_id = 1) const uint BN = 64;
+layout(constant_id = 2) const uint BK = 64;
+const uint STRIDE                     = BK / 4 + 4;
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+layout(binding = 0, std430) readonly buffer W {
+    uint data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    uvec4 data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    int32_t data_d[];
+};
+shared uint weights[2 * BM * STRIDE];
+shared uint activations[2 * BN * STRIDE];
+void load_tile(uint base, uint buf, uint row_base, uint col_base) {
+    const uint tid = gl_LocalInvocationIndex;
+    for (uint i = tid; i < BN * BK / 16; i += 256) {
+        const uint r = i / (BK / 16), t = i % (BK / 16);
+        const uvec4 value       = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0);
+        const uint offset       = (buf * BN + r) * STRIDE + t * 4;
+        activations[offset]     = value.x;
+        activations[offset + 1] = value.y;
+        activations[offset + 2] = value.z;
+        activations[offset + 3] = value.w;
+    }
+    for (uint i = tid; i < BM * BK / 8; i += 256) {
+        const uint c = i / (BK / 8), t = i % (BK / 8);
+        const uint v                                 = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 8) + base / 8 + t] : 0u;
+        const uint lo                                = (v & 15u) | ((v << 4) & 0xf00u) | ((v << 8) & 0xf0000u) | ((v << 12) & 0xf000000u);
+        const uint hi                                = ((v >> 16) & 15u) | ((v >> 12) & 0xf00u) | ((v >> 8) & 0xf0000u) | ((v >> 4) & 0xf000000u);
+        weights[(buf * BM + c) * STRIDE + 2 * t]     = (lo & 0x07070707u) | ((lo & 0x08080808u) * 31u);
+        weights[(buf * BM + c) * STRIDE + 2 * t + 1] = (hi & 0x07070707u) | ((hi & 0x08080808u) * 31u);
+    }
+}
+int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
+    if (row + c >= p.rows || col + r >= p.n)
+        return 0;
+    precise float value = float(v) * data_as[row + c] * data_ws[col + r];
+    if (p.has_bias != 0)
+        value += data_b[col + r];
+    return floatBitsToInt(value);
+}
+void main() {
+    const uint row = gl_WorkGroupID.y * BN, col = gl_WorkGroupID.x * BM;
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    load_tile(0, 0, row, col);
+    barrier();
+    for (uint base = 0; base < p.k; base += BK) {
+        const uint buf = (base / BK) & 1u;
+        if (base + BK < p.k)
+            load_tile(base + BK, buf ^ 1u, row, col);
+        coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+        coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+        coopMatLoad(w, weights, buf * BM * STRIDE, STRIDE, gl_CooperativeMatrixLayoutRowMajor);
+        coopMatLoad(a, activations, buf * BN * STRIDE, STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
+        sum = coopMatMulAdd(w, a, sum);
+        barrier();
+    }
+    coopMatPerElementNV(sum, sum, dequantize, row, col);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    dl                                                            = setTensorLayoutDimensionNV(dl, p.rows, p.n);
+    tensorViewNV<2, false, 1, 0> transpose                        = createTensorViewNV(2, false, 1, 0);
+    coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
+}

src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp

@@ -0,0 +1,92 @@
+#version 460
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_buffer_reference : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_KHR_cooperative_matrix : require
+#extension GL_NV_cooperative_matrix2 : require
+#extension GL_KHR_memory_scope_semantics : require
+layout(local_size_x = 256) in;
+layout(constant_id = 0) const uint BM   = 64;
+layout(constant_id = 1) const uint BN   = 64;
+layout(constant_id = 2) const uint BK   = 64;
+layout(constant_id = 3) const bool FULL = true;
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+    uint row_offset;
+}
+p;
+layout(binding = 0, std430) readonly buffer W {
+    uvec4 data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int8_t data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    int32_t data_d[];
+};
+layout(buffer_reference, std430, buffer_reference_align = 16) buffer PackedBlock {
+    uvec4 value;
+};
+int8_t decode_weight(const in PackedBlock block, const in uint block_coords[2], const in uint coord[2]) {
+    return int8_t(bitfieldExtract(int(block.value[coord[1] / 8]), int((coord[1] % 8) * 4), 4));
+}
+int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
+    if (row + c >= p.rows || col + r >= p.n)
+        return 0;
+    precise float value = float(v) * data_as[row + c] * data_ws[col + r];
+    if (p.has_bias != 0)
+        value += data_b[col + r];
+    return floatBitsToInt(value);
+}
+// The host selects complete weight tiles and a 256-aligned input width.
+void main() {
+    const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y;
+    const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x;
+    const uint panel = min(32u, cols), first_col = (linear / (panel * rows)) * panel;
+    const uint group_cols = min(panel, cols - first_col), within = linear - first_col * rows;
+    const uint row = p.row_offset + (within / group_cols) * BN, col = (first_col + within % group_cols) * BM;
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    tensorLayoutNV<2> wl                                                     = createTensorLayoutNV(2);
+    wl                                                                       = setTensorLayoutBlockSizeNV(wl, 1, 32);
+    wl                                                                       = setTensorLayoutDimensionNV(wl, p.n, p.k & ~255u);
+    tensorLayoutNV<2> al                                                     = createTensorLayoutNV(2);
+    al                                                                       = setTensorLayoutDimensionNV(al, p.rows, p.k & ~255u);
+    tensorViewNV<2, false, 1, 0> at                                          = createTensorViewNV(2, false, 1, 0);
+    if (FULL) {
+        for (uint base = 0; base < p.k; base += BK) {
+            coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+            coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+            coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight);
+            coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(al, row, BN, base, BK), at);
+            sum = coopMatMulAdd(w, a, sum);
+        }
+    } else {
+        tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tail = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+        tail                                                            = setTensorLayoutDimensionNV(tail, p.rows, p.k & ~255u);
+        for (uint base = 0; base < p.k; base += BK) {
+            coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+            coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+            coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight);
+            coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(tail, row, BN, base, BK), at);
+            sum = coopMatMulAdd(w, a, sum);
+        }
+    }
+    coopMatPerElementNV(sum, sum, dequantize, row, col);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    dl                                                            = setTensorLayoutDimensionNV(dl, p.rows, p.n);
+    tensorViewNV<2, false, 1, 0> transpose                        = createTensorViewNV(2, false, 1, 0);
+    coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
+}

src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp

@@ -0,0 +1,244 @@
+#version 460
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_integer_dot_product : require
+
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int data_a[];
+};
+layout(binding = 2, std430) readonly buffer CB {
+    float data_cb[];
+};
+layout(binding = 3, std430) readonly buffer SC {
+    float data_sc[];
+};
+layout(binding = 4, std430) readonly buffer SR {
+    uint data_sr[];
+};
+layout(binding = 5, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 6, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 7, std430) writeonly buffer D {
+    float data_d[];
+};
+
+shared int w_dec[2][16 * 64];
+shared int x_tile[2][32 * 16];
+shared float s_dec[2][4 * 64];
+shared int lut[256];
+shared float s_lut[256];
+shared int cb8s[16];
+shared float s_cb;
+
+float f8_e4m3_to_f32(uint v) {
+    const uint sign = (v >> 7) & 1u;
+    const uint exp  = (v >> 3) & 0xFu;
+    const uint man  = v & 0x7u;
+    float result;
+    if (exp == 0xFu && man == 0x7u) {
+        result = uintBitsToFloat(0x7FC00000u);
+    } else if (exp == 0u) {
+        result = ldexp(float(man), -9);
+    } else {
+        result = ldexp(1.0 + float(man) / 8.0, int(exp) - 7);
+    }
+    return sign != 0u ? -result : result;
+}
+
+void main() {
+    const uint out_base = gl_WorkGroupID.x * 64u;
+    const uint lx       = gl_LocalInvocationID.x;
+    const uint ly       = gl_LocalInvocationID.y;
+
+    const uint row0     = gl_WorkGroupID.y * 32u + ly;
+    const uint row1     = row0 + 16u;
+    const uint idx      = gl_LocalInvocationIndex;
+    const uint w_words  = p.k / 8u;
+    const uint a_words  = p.k / 4u;
+    const uint groups   = p.k / 16u;
+    const uint sr_words = groups / 4u;
+
+    const uint dc   = idx & 63u;
+    const uint dw   = idx >> 6u;
+    const uint ocol = out_base + dc;
+    float sum0      = 0.0;
+    float sum1      = 0.0;
+    float sum2      = 0.0;
+    float sum3      = 0.0;
+    float sum4      = 0.0;
+    float sum5      = 0.0;
+    float sum6      = 0.0;
+    float sum7      = 0.0;
+
+    if (idx == 0u) {
+        float amax = 0.0;
+        for (uint i = 0; i < 16; ++i) {
+            amax = max(amax, abs(data_cb[i]));
+        }
+        s_cb = amax / 127.0;
+        for (uint i = 0; i < 16; ++i) {
+            const int v = (s_cb == 0.0 ? 0 : int(roundEven(data_cb[i] / s_cb)));
+            cb8s[i]     = v < -127 ? -127 : (v > 127 ? 127 : v);
+        }
+    }
+    barrier();
+
+    lut[idx]   = (cb8s[idx & 0xFu] & 0xFF) | (cb8s[idx >> 4] << 8);
+    s_lut[idx] = f8_e4m3_to_f32(idx);
+    barrier();
+
+    {
+        [[unroll]] for (uint j = 0; j < 2; ++j) {
+            const uint pw   = dw * 2 + j;
+            const uint widx = pw;
+            const uint wv   = ocol < p.n && widx < w_words
+                                  ? uint(data_w[ocol * w_words + widx])
+                                  : 0u;
+
+            w_dec[0][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
+            w_dec[0][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
+        }
+        {
+            const uint g = dw;
+            float s      = 0.0;
+            if (ocol < p.n && g < groups) {
+                const uint sr_word = data_sr[ocol * sr_words + g / 4u];
+                s                  = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
+            }
+            s_dec[0][dw * 64 + dc] = s_cb * s;
+        }
+        if (lx < 16u) {
+            const uint aidx                 = lx;
+            x_tile[0][ly * 16 + lx]         = row0 < p.rows && aidx < a_words
+                                                  ? data_a[row0 * a_words + aidx]
+                                                  : 0;
+            x_tile[0][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
+                                                  ? data_a[row1 * a_words + aidx]
+                                                  : 0;
+        }
+    }
+    barrier();
+
+    for (uint gs = 0; gs < groups; gs += 4) {
+        const uint buf = (gs >> 2) & 1u;
+
+        if (gs + 4 < groups) {
+            const uint nbuf = buf ^ 1u;
+            const uint ngs  = gs + 4;
+            [[unroll]] for (uint j = 0; j < 2; ++j) {
+                const uint pw                       = dw * 2 + j;
+                const uint widx                     = ngs * 2 + pw;
+                const uint wv                       = ocol < p.n && widx < w_words
+                                                          ? uint(data_w[ocol * w_words + widx])
+                                                          : 0u;
+                w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
+                w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
+            }
+            {
+                const uint g = ngs + dw;
+                float s      = 0.0;
+                if (ocol < p.n && g < groups) {
+                    const uint sr_word = data_sr[ocol * sr_words + g / 4u];
+                    s                  = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
+                }
+                s_dec[nbuf][dw * 64 + dc] = s_cb * s;
+            }
+            if (lx < 16u) {
+                const uint aidx                    = ngs * 4 + lx;
+                x_tile[nbuf][ly * 16 + lx]         = row0 < p.rows && aidx < a_words
+                                                         ? data_a[row0 * a_words + aidx]
+                                                         : 0;
+                x_tile[nbuf][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
+                                                         ? data_a[row1 * a_words + aidx]
+                                                         : 0;
+            }
+        }
+
+        float partial0 = 0.0;
+        float partial1 = 0.0;
+        float partial2 = 0.0;
+        float partial3 = 0.0;
+        float partial4 = 0.0;
+        float partial5 = 0.0;
+        float partial6 = 0.0;
+        float partial7 = 0.0;
+        {
+            [[unroll]] for (uint gl = 0; gl < 4; ++gl) {
+                int accg0 = 0;
+                int accg1 = 0;
+                int accg2 = 0;
+                int accg3 = 0;
+                int accg4 = 0;
+                int accg5 = 0;
+                int accg6 = 0;
+                int accg7 = 0;
+                [[unroll]] for (uint q = 0; q < 4; ++q) {
+                    const int a0 = x_tile[buf][ly * 16 + gl * 4 + q];
+                    const int a1 = x_tile[buf][(16u + ly) * 16 + gl * 4 + q];
+                    accg0 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a0));
+                    accg1 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a0));
+                    accg2 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a0));
+                    accg3 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a0));
+                    accg4 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a1));
+                    accg5 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a1));
+                    accg6 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a1));
+                    accg7 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a1));
+                }
+                const float s0 = s_dec[buf][gl * 64 + lx];
+                const float s1 = s_dec[buf][gl * 64 + 16u + lx];
+                const float s2 = s_dec[buf][gl * 64 + 32u + lx];
+                const float s3 = s_dec[buf][gl * 64 + 48u + lx];
+                partial0 += float(accg0) * s0;
+                partial1 += float(accg1) * s1;
+                partial2 += float(accg2) * s2;
+                partial3 += float(accg3) * s3;
+                partial4 += float(accg4) * s0;
+                partial5 += float(accg5) * s1;
+                partial6 += float(accg6) * s2;
+                partial7 += float(accg7) * s3;
+            }
+        }
+        sum0 += partial0;
+        sum1 += partial1;
+        sum2 += partial2;
+        sum3 += partial3;
+        sum4 += partial4;
+        sum5 += partial5;
+        sum6 += partial6;
+        sum7 += partial7;
+        barrier();
+    }
+
+    const float sumv[8] = {sum0, sum1, sum2, sum3, sum4, sum5, sum6, sum7};
+    [[unroll]] for (uint t = 0; t < 2u; ++t) {
+        const uint rt = t == 0u ? row0 : row1;
+        if (rt < p.rows) {
+            [[unroll]] for (uint j = 0; j < 4; ++j) {
+                const uint oj = out_base + j * 16u + lx;
+                if (oj < p.n) {
+                    float value = sumv[t * 4u + j] * data_as[rt];
+                    if (p.has_bias != 0) {
+                        value += data_b[oj];
+                    }
+                    data_d[rt * p.n + oj] = value;
+                }
+            }
+        }
+    }
+}

src/core/vulkan/shaders/mul_mat_w4a8_convrot_cm2.comp

@@ -0,0 +1,164 @@
+#version 460
+
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_KHR_cooperative_matrix : require
+#extension GL_NV_cooperative_matrix2 : require
+#extension GL_KHR_memory_scope_semantics : require
+#extension GL_EXT_control_flow_attributes : require
+
+layout(local_size_x = 256) in;
+layout(constant_id = 0) const uint BM = 64;
+layout(constant_id = 1) const uint BN = 64;
+layout(constant_id = 2) const uint BK = 64;
+const uint STRIDE                     = BK / 2 + 4;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    uvec2 data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    uvec4 data_a[];
+};
+layout(binding = 2, std430) readonly buffer CB {
+    float data_cb[];
+};
+layout(binding = 3, std430) readonly buffer SC {
+    float data_sc[];
+};
+layout(binding = 4, std430) readonly buffer SR {
+    uint data_sr[];
+};
+layout(binding = 5, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 6, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 7, std430) writeonly buffer D {
+    int32_t data_d[];
+};
+shared uint lut[256];
+shared int cb[16];
+shared float cb_scale;
+shared float relative[256];
+shared float scales[BK / 16 * BM];
+
+float fp8_e4m3(uint v) {
+    const uint exponent = (v >> 3) & 15u, mantissa = v & 7u;
+    const float value = exponent == 15u && mantissa == 7u ? uintBitsToFloat(0x7fc00000u) : exponent == 0u ? ldexp(float(mantissa), -9)
+                                                                                                          : ldexp(1.0 + float(mantissa) / 8.0, int(exponent) - 7);
+    return (v & 128u) != 0u ? -value : value;
+}
+
+int32_t accumulate_group(const in uint row, const in uint col, const in int32_t value, const in int32_t previous, const in uint group) {
+    return floatBitsToInt(fma(float(value), scales[group * BM + row], intBitsToFloat(previous)));
+}
+
+shared uint weights[BM * STRIDE];
+const uint A_STRIDE = BK / 4 + 4;
+shared uint activations[BN * A_STRIDE];
+
+int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
+    const uint r = row_base + col, c = col_base + row;
+    if (r >= p.rows || c >= p.n) {
+        return 0;
+    }
+    precise float value = intBitsToFloat(acc) * data_as[r];
+    if (p.has_bias != 0) {
+        value += data_b[c];
+    }
+    return floatBitsToInt(value);
+}
+
+void main() {
+    const uint tid  = gl_LocalInvocationIndex;
+    const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y;
+    const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x;
+    // Reuse weight panels across activation tiles when the output expands.
+    const uint panel      = p.n > p.k ? 16u : cols;
+    const uint first_col  = (linear / (panel * rows)) * panel;
+    const uint group_cols = min(panel, cols - first_col);
+    const uint within     = linear - first_col * rows;
+    const uint col_base = (first_col + within % group_cols) * BM, row_base = (within / group_cols) * BN;
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout      = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    d_layout                                                                 = setTensorLayoutDimensionNV(d_layout, p.rows, p.n);
+    tensorViewNV<2, false, 1, 0> transpose                                   = createTensorViewNV(2, false, 1, 0);
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    if (tid == 0) {
+        float amax = 0.0;
+        for (uint i = 0; i < 16; ++i) {
+            amax = max(amax, abs(data_cb[i]));
+        }
+        cb_scale = amax / 127.0;
+        for (uint i = 0; i < 16; ++i) {
+            cb[i] = cb_scale == 0.0 ? 0 : clamp(int(roundEven(data_cb[i] / cb_scale)), -127, 127);
+        }
+    }
+    relative[tid] = fp8_e4m3(tid);
+    barrier();
+    for (uint i = tid; i < 256; i += 256) {
+        lut[i] = uint(cb[i & 15u] & 255) | (uint(cb[i >> 4] & 255) << 8);
+    }
+    barrier();
+    for (uint base = 0; base < p.k; base += BK) {
+        for (uint i = tid; i < BN * BK / 16; i += 256) {
+            const uint r = i / (BK / 16), t = i % (BK / 16);
+            const uvec4 value       = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0);
+            const uint offset       = r * A_STRIDE + t * 4;
+            activations[offset]     = value.x;
+            activations[offset + 1] = value.y;
+            activations[offset + 2] = value.z;
+            activations[offset + 3] = value.w;
+        }
+        for (uint i = tid; i < BM * BK / 16; i += 256) {
+            const uint c = i / (BK / 16), part = i % (BK / 16);
+            const uvec2 packed = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 16) + base / 16 + part] : uvec2(0);
+            // Complementary zero halves let adjacent groups reuse one activation fragment.
+            const uint odd         = part % 2;
+            const uint offset      = c * STRIDE + part * 8 + odd * 4;
+            const uint zero_offset = c * STRIDE + part * 8 + (1 - odd) * 4;
+            for (uint j = 0; j < 2; ++j) {
+                const uint v                     = packed[j];
+                weights[offset + 2 * j]          = lut[v & 255u] | (lut[(v >> 8) & 255u] << 16);
+                weights[offset + 2 * j + 1]      = lut[(v >> 16) & 255u] | (lut[v >> 24] << 16);
+                weights[zero_offset + 2 * j]     = 0u;
+                weights[zero_offset + 2 * j + 1] = 0u;
+            }
+        }
+        for (uint i = tid; i < BM * BK / 16; i += 256) {
+            const uint c = i / (BK / 16), g = i % (BK / 16);
+            float scale = 0.0;
+            if (col_base + c < p.n) {
+                const uint packed = data_sr[(col_base + c) * (p.k / 64) + base / 64 + g / 4];
+                scale             = cb_scale * (data_sc[col_base + c] * relative[(packed >> (8 * (g % 4))) & 255u]);
+            }
+            scales[g * BM + c] = scale;
+        }
+        barrier();
+        // The relative scale changes every 16 input values, so each group
+        // needs its own integer MMA and floating-point accumulation. Pad the
+        // weights with zeros to satisfy the device's K=32 matrix granularity.
+        [[unroll]] for (uint g = 0; g < BK / 16; g += 2) {
+            coopmat<int8_t, gl_ScopeWorkgroup, 32, BN, gl_MatrixUseB> a;
+            coopMatLoad(a, activations, g * 4, A_STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
+            [[unroll]] for (uint j = 0; j < 2; ++j) {
+                coopmat<int8_t, gl_ScopeWorkgroup, BM, 32, gl_MatrixUseA> w;
+                coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> part = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+                coopMatLoad(w, weights, (g + j) * 8, STRIDE, gl_CooperativeMatrixLayoutRowMajor);
+                part = coopMatMulAdd(w, a, part);
+                coopMatPerElementNV(sum, part, accumulate_group, sum, g + j);
+            }
+        }
+        barrier();
+    }
+    coopMatPerElementNV(sum, sum, dequantize, row_base, col_base);
+    coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row_base, BN, col_base, BM), transpose);
+}

src/core/vulkan/shaders/quantize_i8_convrot_h64.comp

@@ -0,0 +1,115 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+#ifndef GROUP_SIZE
+#define GROUP_SIZE 256
+#endif
+#if GROUP_SIZE != 64 && GROUP_SIZE != 256
+#error unsupported GROUP_SIZE
+#endif
+#define GROUP_THREADS (GROUP_SIZE / 4)
+
+#if GROUP_SIZE == 256
+#define HADAMARD_NORM (1.0 / 16.0)
+#else
+#define HADAMARD_NORM (1.0 / 8.0)
+#endif
+
+layout(local_size_x = GROUP_THREADS, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint k;
+    uint rows;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
+
+shared float values[GROUP_SIZE];
+shared float maxima[GROUP_THREADS];
+
+void transform_group(const uint row, const uint group) {
+    const uint tid         = gl_LocalInvocationID.x;
+    const uint base_offset = row * p.k + group * GROUP_SIZE;
+
+    [[unroll]] for (uint i = 0; i < 4; ++i) {
+        const uint index = tid + i * GROUP_THREADS;
+        values[index]    = data_a[base_offset + index] * HADAMARD_NORM;
+    }
+    barrier();
+
+    [[unroll]] for (uint stride = 1; stride < GROUP_SIZE; stride *= 4) {
+        const uint base = (tid / stride) * 4 * stride + tid % stride;
+        const uint i0   = base;
+        const uint i1   = i0 + stride;
+        const uint i2   = i1 + stride;
+        const uint i3   = i2 + stride;
+        const float a   = values[i0];
+        const float b   = values[i1];
+        const float c   = values[i2];
+        const float d   = values[i3];
+        values[i0]      = a + b + c - d;
+        values[i1]      = a + b - c + d;
+        values[i2]      = a - b + c + d;
+        values[i3]      = -a + b + c + d;
+        barrier();
+    }
+}
+
+uint quantize_value(const float value, const float inv_scale) {
+    const int quantized = clamp(int(roundEven(value * inv_scale)), -127, 127);
+    return uint(quantized) & 0xffu;
+}
+
+void quantize_row(const uint row) {
+    const uint tid    = gl_LocalInvocationID.x;
+    const uint groups = p.k / GROUP_SIZE;
+    float local_max   = 0.0;
+
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        [[unroll]] for (uint i = 0; i < 4; ++i) {
+            local_max = max(local_max, abs(values[tid + i * GROUP_THREADS]));
+        }
+        barrier();
+    }
+
+    maxima[tid] = local_max;
+    barrier();
+    [[unroll]] for (uint stride = GROUP_THREADS / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            maxima[tid] = max(maxima[tid], maxima[tid + stride]);
+        }
+        barrier();
+    }
+
+    if (tid == 0) {
+        data_s[row] = maxima[0] / 127.0;
+    }
+    barrier();
+
+    const float row_scale = maxima[0] / 127.0;
+    const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        const uint index                                   = tid * 4;
+        const uint packed                                  = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
+        data_d[(row * p.k + group * GROUP_SIZE) / 4 + tid] = packed;
+        barrier();
+    }
+}
+
+void main() {
+    for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
+        quantize_row(row);
+    }
+}

src/core/vulkan/shaders/rms_modulate.comp

@@ -0,0 +1,136 @@
+#version 450
+#extension GL_EXT_control_flow_attributes : enable
+#define BLOCK_SIZE 512
+#define FLOAT_TYPE float
+#define D_TYPE float
+layout(local_size_x = BLOCK_SIZE) in;
+layout(binding = 0) readonly buffer X {
+    float data_a[];
+};
+layout(binding = 1) readonly buffer W {
+    float data_b[];
+};
+layout(binding = 2) writeonly buffer D {
+    float data_d[];
+};
+layout(binding = 3) readonly buffer S {
+    float data_c[];
+};
+layout(binding = 4) readonly buffer H {
+    float data_e[];
+};
+layout(constant_id = 0) const bool has_shift   = true;
+layout(constant_id = 1) const bool do_multiply = true;
+layout(push_constant) uniform Parameters {
+    uint ne[4], stride[4], offsets[5];
+    float eps;
+    uint ne10;
+}
+p;
+shared float sumsh[BLOCK_SIZE];
+shared float modulation_input[BLOCK_SIZE];
+uint get_aoffset() {
+    return p.offsets[0];
+}
+uint get_boffset() {
+    return p.offsets[1];
+}
+uint get_doffset() {
+    return p.offsets[2];
+}
+uint src1_idx(uint a, uint b, uint c, uint d) {
+    return 0;
+}
+uint fastmod(uint a, uint b) {
+    return a % b;
+}
+float modulate(float value, uint col) {
+    // Keep precise modulation from changing the native normalization's FMA contraction.
+    modulation_input[gl_LocalInvocationIndex] = value;
+    memoryBarrierShared();
+    float weighted        = modulation_input[gl_LocalInvocationIndex];
+    precise float product = weighted * data_c[p.offsets[3] + col];
+    precise float result  = weighted + product;
+    if (has_shift)
+        result += data_e[p.offsets[4] + col];
+    return result;
+}
+void rms_norm(uint num_iters) {
+    const uint ncols     = p.ne[0];
+    const uint nrows     = gl_NumWorkGroups.x;
+    const uint nchannels = p.ne[2];
+    const uint nsamples  = p.ne[3];
+
+    const uint row = gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    const uint stride_row     = p.stride[1];
+    const uint stride_channel = p.stride[2];
+    const uint stride_sample  = p.stride[3];
+
+    // grid.y/z are clamped to the device workgroup limit, iterate over the excess channels/samples
+    for (uint samp = gl_WorkGroupID.z; samp < nsamples; samp += gl_NumWorkGroups.z) {
+        for (uint channel = gl_WorkGroupID.y; channel < nchannels; channel += gl_NumWorkGroups.y) {
+            barrier();
+
+            uint a_offset  = samp * stride_sample + channel * stride_channel + row * stride_row + get_aoffset();
+            uint b_offset  = src1_idx(0, row, channel, samp) + get_boffset();
+            uint d_offset  = ((samp * nchannels + channel) * nrows + row) * ncols + get_doffset();
+            FLOAT_TYPE sum = FLOAT_TYPE(0.0f);  // partial sum for thread in warp
+
+            [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+                FLOAT_TYPE xi = FLOAT_TYPE(0);
+                if (col < ncols) {
+                    xi = FLOAT_TYPE(data_a[a_offset + col]);
+                }
+                sum += xi * xi;
+            }
+
+            sumsh[tid] = sum;
+            // sum up partial sums and write back result
+            barrier();
+            [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+                if (tid < s) {
+                    sum += sumsh[tid + s];
+                    sumsh[tid] = sum;
+                }
+                barrier();
+            }
+            sum = sumsh[0];
+
+            const FLOAT_TYPE mean  = sum / FLOAT_TYPE(ncols);
+            const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.eps));
+
+            if (do_multiply) {
+                if (ncols > p.ne10) {
+                    [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+                        if (col >= ncols) {
+                            continue;
+                        }
+                        FLOAT_TYPE value       = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]);
+                        data_d[d_offset + col] = modulate(value, col);
+                    }
+                } else {
+                    [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+                        if (col >= ncols) {
+                            continue;
+                        }
+                        FLOAT_TYPE value       = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]);
+                        data_d[d_offset + col] = modulate(value, col);
+                    }
+                }
+            } else {
+                [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) {
+                    if (col >= ncols) {
+                        continue;
+                    }
+                    data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]));
+                }
+            }
+        }
+    }
+}
+
+void main() {
+    rms_norm(8);
+}

src/core/vulkan/shaders/rms_rope_apply.comp

@@ -12,6 +12,7 @@ layout(binding = 1) readonly buffer PE {
     float data_pe[];
 };
 #ifdef CAST_F16
+layout(constant_id = 0) const bool apply_output_scale = false;
 #extension GL_EXT_shader_16bit_storage : require
 #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
 layout(binding = 2) writeonly buffer D {
@@ -96,8 +97,10 @@ void main() {
                     }
                 }
 #ifdef CAST_F16
-                value *= pc.output_scale;
-                value += 0.0;
+                if (apply_output_scale) {
+                    value *= pc.output_scale;
+                    value += 0.0;
+                }
                 data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = float16_t(value);
 #else
                 data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = value;

src/generation_info.cpp

@@ -1,6 +1,7 @@
 #include "generation_info.h"
 
 #include <cmath>
+#include <cstdio>
 #include <cstdlib>
 #include <cstring>
 
@@ -14,6 +15,21 @@ namespace sd {
         return value ? value : "";
     }
 
+    static double shortest_f32_as_double(float value) {
+        if (!std::isfinite(value)) {
+            return static_cast<double>(value);
+        }
+        char buf[64];
+        buf[0] = '\0';
+        for (int prec = 1; prec <= 9; ++prec) {
+            std::snprintf(buf, sizeof(buf), "%.*g", prec, static_cast<double>(value));
+            if (std::strtof(buf, nullptr) == value) {
+                break;
+            }
+        }
+        return std::strtod(buf, nullptr);
+    }
+
     static json sampling_info(const sd_sample_params_t& p, bool high_noise = false) {
         const auto& g = p.guidance;
         json info     = {{"sample_steps", p.sample_steps}, {"guidance", {{"txt_cfg", g.txt_cfg}, {"distilled_guidance", g.distilled_guidance}}}};
@@ -117,8 +133,13 @@ namespace sd {
         }
         if (p.lora_count > 0) {
             info["loras"] = json::array();
-            for (uint32_t i = 0; i < p.lora_count; ++i)
-                info["loras"].push_back({{"path", text(p.loras[i].path)}, {"multiplier", p.loras[i].multiplier}, {"is_high_noise", p.loras[i].is_high_noise}});
+            for (uint32_t i = 0; i < p.lora_count; ++i) {
+                // nlohmann::json promotes float to double, so dump the shortest
+                // decimal that round-trips the float32 (e.g. 0.85f stays 0.85
+                // instead of 0.8500000238418579) to keep replay exact.
+                const double multiplier = shortest_f32_as_double(p.loras[i].multiplier);
+                info["loras"].push_back({{"path", text(p.loras[i].path)}, {"multiplier", multiplier}, {"is_high_noise", p.loras[i].is_high_noise}});
+            }
         }
         return info;
     }
@@ -127,8 +148,11 @@ namespace sd {
         auto info = common_info(p, "image");
         if (p.init_image.data)
             info["strength"] = p.strength;
-        if (p.control_image.data)
+        if (p.control_image.data) {
             info["control_strength"] = p.control_strength;
+            if (p.control_type != SD_CONTROL_TYPE_NONE)
+                info["control_type"] = sd_control_type_name(p.control_type);
+        }
         if (p.ip_adapter_image.data)
             info["ip_adapter_strength"] = p.ip_adapter_strength;
         if (p.pm_params.id_images && p.pm_params.id_images_count > 0) {
@@ -188,7 +212,7 @@ namespace sd {
         auto inpaint        = image_info(inpaint_params);
         inpaint["strength"] = inpaint_params.strength;
         // ADetailer clears these inputs before generating each crop.
-        for (const auto* key : {"control_strength", "pm_params", "pulid_params", "hires", "image_preprocess"})
+        for (const auto* key : {"control_strength", "control_type", "pm_params", "pulid_params", "hires", "image_preprocess"})
             inpaint.erase(key);
         sd_img_gen_params_t defaults;
         sd_img_gen_params_init(&defaults);

src/model.h

@@ -70,6 +70,7 @@ enum SDVersion {
     VERSION_MOSS_IMAGE,
     VERSION_FIBO2,
     VERSION_CHROMA,
+    VERSION_P_H3_IMAGE,
     VERSION_COUNT,
 };
 
@@ -166,7 +167,7 @@ static inline bool sd_version_is_ltxav(SDVersion version) {
 }
 
 static inline bool sd_version_is_minimax_h3(SDVersion version) {
-    return version == VERSION_MINIMAX_H3;
+    return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE;
 }
 
 static inline bool sd_version_is_wan(SDVersion version) {

src/model/adapter/ip_adapter.hpp

@@ -141,7 +141,7 @@ namespace IPAdapter {
                 }
                 auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
                 if (proj_in_iter != tensor_storage_map.end()) {
-                    embed_dim = proj_in_iter->second.ne[0];
+                    embed_dim = proj_in_iter->second.logical_ne0();
                 }
                 auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
                 if (proj_out_iter != tensor_storage_map.end()) {
@@ -168,7 +168,7 @@ namespace IPAdapter {
                 }
                 auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
                 if (proj_iter != tensor_storage_map.end()) {
-                    clip_dim = proj_iter->second.ne[0];
+                    clip_dim = proj_iter->second.logical_ne0();
                     out_dim  = proj_iter->second.ne[1];
                 }
                 num_tokens = out_dim / ctx_dim;

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,18 +1135,28 @@ 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;
+        if (forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_LINEAR &&
+            w->type == GGML_TYPE_I8 && x->ne[0] == 2 * w->ne[0]) {
+            logical_weight.ne[0] = x->ne[0];
+            for (int i = 1; i < GGML_MAX_DIMS; ++i)
+                logical_weight.nb[i] = logical_weight.nb[i - 1] * logical_weight.ne[i - 1];
+            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,
@@ -1140,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]
@@ -319,7 +317,7 @@ protected:
             tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
             has_ip            = true;
             int64_t inner_dim = d_head * n_head;
-            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
+            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").logical_ne0();
             blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
             blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
         }
@@ -576,6 +574,33 @@ public:
     }
 };
 
+class ResidualAttentionBlock : public GGMLBlock {
+public:
+    ResidualAttentionBlock(int64_t channels, int64_t n_head) {
+        blocks["ln_1"]       = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
+        blocks["attn"]       = std::shared_ptr<GGMLBlock>(new MultiheadAttention(channels, n_head, true, true, true));
+        blocks["ln_2"]       = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
+        blocks["mlp.c_fc"]   = std::shared_ptr<GGMLBlock>(new Linear(channels, channels * 4));
+        blocks["mlp.c_proj"] = std::shared_ptr<GGMLBlock>(new Linear(channels * 4, channels));
+    }
+
+    // x: [N, n_token, channels]
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        auto ln_1       = std::dynamic_pointer_cast<LayerNorm>(blocks["ln_1"]);
+        auto attn       = std::dynamic_pointer_cast<MultiheadAttention>(blocks["attn"]);
+        auto ln_2       = std::dynamic_pointer_cast<LayerNorm>(blocks["ln_2"]);
+        auto mlp_c_fc   = std::dynamic_pointer_cast<Linear>(blocks["mlp.c_fc"]);
+        auto mlp_c_proj = std::dynamic_pointer_cast<Linear>(blocks["mlp.c_proj"]);
+
+        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, h);
+        h      = mlp_c_proj->forward(ctx, h);
+        return ggml_add(ctx->ggml_ctx, x, h);
+    }
+};
+
 class AlphaBlender : public GGMLBlock {
 protected:
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
@@ -587,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,7 +147,9 @@ 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;
     int int8_convrot_group_size = 0;
     float scale;
@@ -158,16 +160,21 @@ protected:
         has_weight_scale        = false;
         int8_convrot            = false;
         int8_convrot_group_size = 0;
+        auto weight_storage     = tensor_storage_map.find(prefix + "weight");
+        w4_convrot_kind         = weight_storage != tensor_storage_map.end() ? weight_storage->second.w4_convrot_kind : 0;
         enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
-        if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
+        if (!w4_convrot_kind && (in_features % ggml_blck_size(wtype) != 0 || force_f32)) {
             wtype = GGML_TYPE_F32;
         }
-        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (w4_convrot_kind) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8 && in_features % 2 == 0);
+            GGML_ASSERT(weight_storage->second.ne[0] == in_features / 2 && weight_storage->second.ne[1] == out_features);
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, w4_convrot_kind ? in_features / 2 : in_features, out_features);
         if (bias) {
             enum ggml_type wtype = GGML_TYPE_F32;
             params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
         }
-        auto weight_storage           = tensor_storage_map.find(prefix + "weight");
         const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
         auto weight_scale_storage     = tensor_storage_map.find(prefix + "weight_scale");
         if (weight_scale_storage != tensor_storage_map.end()) {
@@ -176,9 +183,16 @@ protected:
             params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
             has_weight_scale       = true;
         }
-        if (is_int8_tensorwise) {
+        if (w4_convrot_kind == 2) {
+            params["weight_codebook"]  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16);
+            params["weight_s_channel"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
+#ifndef SD_USE_UPSTREAM_GGML
+            params["weight_s_rel"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F8_E4M3, in_features / 16, out_features);
+#endif
+        }
+        if (is_int8_tensorwise || w4_convrot_kind) {
             GGML_ASSERT(wtype == GGML_TYPE_I8);
-            GGML_ASSERT(has_weight_scale);
+            GGML_ASSERT(w4_convrot_kind == 2 || has_weight_scale);
             int8_convrot            = weight_storage->second.int8_convrot;
             int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
         }
@@ -202,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_;
     }
@@ -212,92 +230,60 @@ 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;
-                }
-            }
+            if (w4_convrot_kind) {
+                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,
-                                                x,
-                                                w,
-                                                weight_scale,
-                                                b,
-                                                int8_convrot ? int8_convrot_group_size : 0,
-                                                scale);
+                out = ggml_ext_linear_i8_tensorwise(ctx,
+                                                    x,
+                                                    w,
+                                                    weight_scale,
+                                                    b,
+                                                    int8_convrot ? int8_convrot_group_size : 0,
+                                                    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,
-                                                                                               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) {
@@ -310,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;
     }
@@ -345,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;
@@ -420,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);
     }
 };
 
@@ -497,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);
     }
 };
@@ -567,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,
@@ -656,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,
@@ -668,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;
@@ -712,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;
     }
@@ -792,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);
             }
@@ -800,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),
@@ -848,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);
     }
 };
 
@@ -896,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);
     }
 };
 
@@ -940,10 +849,11 @@ 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);
-        x = ggml_mul_inplace(ctx->ggml_ctx, x, w);
+        // The affine multiply may run on a different backend after normalization falls back.
+        x = ggml_mul(ctx->ggml_ctx, x, w);
         return x;
     }
 };
@@ -1000,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

@@ -0,0 +1,221 @@
+#ifndef __SD_MODEL_COMMON_VSR_HPP__
+#define __SD_MODEL_COMMON_VSR_HPP__
+
+#include <array>
+#include <limits>
+#include <map>
+#include <stdexcept>
+#include <unordered_map>
+#include "model/common/ggml_block.hpp"
+#include "model/common/rope.hpp"
+
+namespace VSR {
+
+    struct GraphInputs {
+        std::map<std::vector<int>, sd::Tensor<float>> rotations;
+        mutable std::unordered_map<const void*, ggml_tensor*> bindings;
+
+        void begin_graph() {
+            bindings.clear();
+        }
+
+        void clear() {
+            bindings.clear();
+            rotations.clear();
+        }
+
+        template <typename T>
+        ggml_tensor* input(GGMLRunnerContext* ctx, const sd::Tensor<T>& value) const {
+            auto existing = bindings.find(value.data());
+            if (existing != bindings.end())
+                return existing->second;
+            auto result = sd::make_ggml_tensor(ctx->ggml_ctx, value, false);
+            ctx->set_backend_tensor_data(result, value.data());
+            bindings.emplace(value.data(), result);
+            return result;
+        }
+
+        ggml_tensor* rope(GGMLRunnerContext* ctx, int width, int height, const std::vector<int>& times, const std::vector<int>& axes) {
+            std::vector<int> key = {width, height, axes[0], axes[1], axes[2]};
+            key.insert(key.end(), times.begin(), times.end());
+            auto existing = rotations.find(key);
+            if (existing != rotations.end())
+                return input(ctx, existing->second);
+            const int dim = axes[0] + axes[1] + axes[2];
+            auto& pe      = rotations.emplace(std::move(key), sd::Tensor<float>({2, 2, dim / 2, int64_t(width) * height * int64_t(times.size())})).first->second;
+            float* out    = pe.data();
+            for (int t : times) {
+                for (int y = 0; y < height; ++y) {
+                    for (int x = 0; x < width; ++x) {
+                        const int pos[] = {t, y, x};
+                        for (int a = 0; a < 3; ++a) {
+                            for (int i = 0; i < axes[a]; i += 2) {
+                                const double angle = pos[a] * std::pow(10000.0, -double(i) / axes[a]);
+                                *out++             = float(std::cos(angle));
+                                *out++             = float(-std::sin(angle));
+                                *out++             = float(std::sin(angle));
+                                *out++             = float(std::cos(angle));
+                            }
+                        }
+                    }
+                }
+            }
+            return input(ctx, pe);
+        }
+    };
+
+    static ggml_tensor* tokens(ggml_context* ctx, ggml_tensor* x) {
+        x = ggml_reshape_2d(ctx, x, x->ne[0] * x->ne[1], x->ne[2] * x->ne[3]);
+        return ggml_cont(ctx, ggml_transpose(ctx, x));
+    }
+
+    static ggml_tensor* image(ggml_context* ctx, ggml_tensor* x, int width, int height) {
+        const auto channels = x->ne[0];
+        x                   = ggml_cont(ctx, ggml_transpose(ctx, x));
+        return ggml_reshape_4d(ctx, x, width, height, channels, 1);
+    }
+
+    static ggml_tensor* pixel_shuffle(ggml_context* ctx, ggml_tensor* x, int width, int height, int scale, bool channel_first = true) {
+        const int64_t channels = x->ne[0] / (scale * scale);
+        if (!channel_first) {
+            x = ggml_reshape_4d(ctx, x, channels, scale, scale, width * height);
+            x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));
+        }
+        x = ggml_reshape_4d(ctx, x, scale, scale, channels, width * height);
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 3, 1));
+        x = ggml_reshape_4d(ctx, x, scale * width, height, scale, channels);
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
+        return ggml_reshape_4d(ctx, x, scale * width, scale * height, channels, 1);
+    }
+
+    struct AttentionTile {
+        int width;
+        int height;
+        sd::Tensor<int32_t> queries;
+        sd::Tensor<int32_t> keys;
+        sd::Tensor<float> mask;
+    };
+
+    struct LocalAttention {
+        int width;
+        int height;
+        std::vector<std::pair<int, int>> windows;
+        bool preserve_boundary;
+        mutable std::vector<AttentionTile> tiles;
+        GraphInputs inputs;
+
+        void begin_graph() {
+            inputs.begin_graph();
+        }
+
+        LocalAttention(int width, int height, const std::vector<std::pair<int, int>>& windows, bool preserve_boundary)
+            : width(width), height(height), windows(windows), preserve_boundary(preserve_boundary) {}
+
+        void prepare_tiles() const {
+            if (!tiles.empty())
+                return;
+            constexpr int tile_size = 16;
+            for (int y = 0; y < height; y += tile_size) {
+                for (int x = 0; x < width; x += tile_size) {
+                    AttentionTile tile;
+                    tile.width   = std::min(tile_size, width - x);
+                    tile.height  = std::min(tile_size, height - y);
+                    tile.queries = sd::Tensor<int32_t>({tile.width * tile.height});
+                    std::vector<int32_t> keys;
+                    auto bounds = [&](int qy, int qx, int plane) {
+                        auto [wh, ww] = windows[plane];
+                        if (preserve_boundary) {
+                            wh             = std::min(wh, height);
+                            ww             = std::min(ww, width);
+                            const int top  = std::clamp(qy - wh / 2, 0, height - wh);
+                            const int left = std::clamp(qx - ww / 2, 0, width - ww);
+                            return std::array<int, 4>{top, left, top + wh, left + ww};
+                        }
+                        return std::array<int, 4>{std::max(0, qy - wh / 2), std::max(0, qx - ww / 2),
+                                                  std::min(height, qy + wh / 2), std::min(width, qx + ww / 2)};
+                    };
+                    for (int p = 0; p < int(windows.size()); ++p) {
+                        auto begin = bounds(y, x, p);
+                        auto end   = bounds(y + tile.height - 1, x + tile.width - 1, p);
+                        for (int ky = begin[0]; ky < end[2]; ++ky)
+                            for (int kx = begin[1]; kx < end[3]; ++kx)
+                                keys.push_back(p * height * width + ky * width + kx);
+                    }
+                    tile.keys = sd::Tensor<int32_t>({int64_t(keys.size())});
+                    std::copy(keys.begin(), keys.end(), tile.keys.data());
+                    tile.mask = sd::Tensor<float>({int64_t(keys.size()), tile.width * tile.height});
+                    for (int q = 0; q < tile.width * tile.height; ++q) {
+                        int qy                 = y + q / tile.width;
+                        int qx                 = x + q % tile.width;
+                        tile.queries.data()[q] = qy * width + qx;
+                        for (size_t k = 0; k < keys.size(); ++k) {
+                            int p                                 = keys[k] / (width * height);
+                            int ky                                = (keys[k] % (width * height)) / width;
+                            int kx                                = keys[k] % width;
+                            auto b                                = bounds(qy, qx, p);
+                            tile.mask.data()[q * keys.size() + k] = ky >= b[0] && ky < b[2] && kx >= b[1] && kx < b[3]
+                                                                        ? 0.f
+                                                                        : -std::numeric_limits<float>::infinity();
+                        }
+                    }
+                    tiles.push_back(std::move(tile));
+                }
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, int heads) const {
+            auto c = ctx->ggml_ctx;
+#ifndef SD_USE_UPSTREAM_GGML
+            if (ctx->flash_attn_enabled && windows.size() <= 4) {
+                SDWindowAttentionParams params = {};
+                params.width                   = width;
+                params.height                  = height;
+                params.heads                   = heads;
+                params.planes                  = int(windows.size());
+                params.preserve_boundary       = preserve_boundary;
+                for (size_t i = 0; i < windows.size(); ++i) {
+                    params.window_height[i] = windows[i].first;
+                    params.window_width[i]  = windows[i].second;
+                }
+                if (auto out = ggml_ext_window_attn(ctx, q, k, v, params))
+                    return out;
+            }
+#endif
+            prepare_tiles();
+            ggml_tensor* output = nullptr;
+            ggml_tensor* row    = nullptr;
+            int row_width       = 0;
+            for (const auto& tile : tiles) {
+                auto qi   = inputs.input(ctx, tile.queries);
+                auto ki   = inputs.input(ctx, tile.keys);
+                auto tq   = ggml_get_rows(c, q, qi);
+                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(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;
+                if (row_width == width) {
+                    output    = output ? ggml_concat(c, output, row, 2) : row;
+                    row       = nullptr;
+                    row_width = 0;
+                }
+            }
+            return ggml_reshape_2d(c, output, output->ne[0], width * height);
+        }
+    };
+
+    static ggml_tensor* apply_rope(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, int heads) {
+        auto c              = ctx->ggml_ctx;
+        const int64_t dim   = x->ne[0];
+        const int64_t count = x->ne[1];
+        x                   = ggml_reshape_4d(c, x, dim / heads, heads, count, 1);
+        x                   = Rope::apply_rope(ctx, x, pe);
+        x                   = ggml_cont(c, ggml_permute(c, x, 0, 2, 1, 3));
+        return ggml_reshape_2d(c, x, dim, count);
+    }
+
+}  // namespace VSR
+
+#endif  // __SD_MODEL_COMMON_VSR_HPP__

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

@@ -72,10 +72,10 @@ namespace Boogu {
                 }
                 if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area = config.patch_size * config.patch_size;
-                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.in_channels = tensor_storage.logical_ne0() / patch_area;
                     config.hidden_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) {
-                    config.instruction_feat_dim = tensor_storage.ne[0];
+                    config.instruction_feat_dim = tensor_storage.logical_ne0();
                     config.hidden_size          = tensor_storage.ne[1];
                 } else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
                     detected_head_dim = tensor_storage.ne[0];
@@ -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

@@ -2,10 +2,47 @@
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"
+#include "stable-diffusion.h"
 
 // Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
 #define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
 
+/*
+    =============================== ControlNet Union ================================
+    Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnets/controlnet_union.py
+
+*/
+struct ControlNetUnionConfig {
+    int num_control_types = 0;  // 0: not a union model
+    int num_fusion_layers = 0;
+
+    static constexpr int CONTROL_TYPE_EMBED_DIM = 256;
+    static constexpr int FUSION_HEADS           = 8;
+
+    static ControlNetUnionConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                     const std::string& prefix,
+                                                     SDVersion version = VERSION_SD1) {
+        ControlNetUnionConfig config;
+        if (!sd_version_is_sdxl(version)) {
+            return config;
+        }
+        const std::string p = prefix.empty() ? "" : prefix + ".";
+        auto embedding      = tensor_storage_map.find(p + "control_add_embedding.linear_1.weight");
+        if (embedding == tensor_storage_map.end()) {
+            return config;
+        }
+        // control_add_embedding takes CONTROL_TYPE_EMBED_DIM inputs per control type
+        config.num_control_types = static_cast<int>(embedding->second.logical_ne0() / CONTROL_TYPE_EMBED_DIM);
+        while (tensor_storage_map.find(p + "transformer_layes." + std::to_string(config.num_fusion_layers) + ".ln_1.weight") != tensor_storage_map.end()) {
+            config.num_fusion_layers++;
+        }
+        LOG_VERBOSE("controlnet union: num_control_types = %d, num_fusion_layers = %d",
+                    config.num_control_types,
+                    config.num_fusion_layers);
+        return config;
+    }
+};
+
 /*
     =================================== ControlNet ===================================
     Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/cldm/cldm.py
@@ -13,6 +50,8 @@
 */
 class ControlNetBlock : public GGMLBlock {
 protected:
+    ControlNetUnionConfig config;
+
     SDVersion version = VERSION_SD1;
     // network hparams
     int in_channels                        = 4;
@@ -28,12 +67,19 @@ protected:
     int context_dim                        = 768;  // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
     bool use_linear_projection             = false;
 
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        if (is_union()) {
+            // [num_control_types, model_channels] in the checkpoint
+            params["task_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, model_channels, config.num_control_types);
+        }
+    }
+
 public:
     int model_channels  = 320;
     int adm_in_channels = 2816;  // only for VERSION_SDXL
 
-    ControlNetBlock(SDVersion version = VERSION_SD1)
-        : version(version) {
+    ControlNetBlock(SDVersion version = VERSION_SD1, ControlNetUnionConfig config = {})
+        : config(config), version(version) {
         if (sd_version_is_sd2(version)) {
             context_dim       = 1024;
             num_head_channels = 64;
@@ -162,6 +208,77 @@ public:
 
         // middle_block_out
         blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
+
+        if (is_union()) {
+            blocks["control_add_embedding.linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, time_embed_dim));
+            // nn.SiLU()
+            blocks["control_add_embedding.linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
+            for (int i = 0; i < config.num_fusion_layers; i++) {
+                blocks["transformer_layes." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResidualAttentionBlock(model_channels, ControlNetUnionConfig::FUSION_HEADS));
+            }
+            blocks["spatial_ch_projs"] = std::shared_ptr<GGMLBlock>(new Linear(model_channels, model_channels));
+        }
+    }
+
+    bool is_union() const {
+        return config.num_control_types > 0;
+    }
+
+    // control_types: [N*num_control_types], 1 for the types in use and 0 for the others
+    ggml_tensor* control_type_embedding(GGMLRunnerContext* ctx, ggml_tensor* control_types, int64_t batch) {
+        auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["control_add_embedding.linear_1"]);
+        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, 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);
+        return linear_2->forward(ctx, e);  // [N, time_embed_dim]
+    }
+
+    // [N, C, H, W] => [N, C]
+    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]);
+    }
+
+    // sample: [N, C, H, W], hint: [N, C, H, W] or [1, C, H, W]
+    ggml_tensor* fuse_condition(GGMLRunnerContext* ctx,
+                                ggml_tensor* sample,
+                                ggml_tensor* hint,
+                                int control_type_index) {
+        auto c           = ctx->ggml_ctx;
+        int64_t channels = sample->ne[2];
+        int64_t batch    = sample->ne[3];
+
+        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(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(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]
+        auto x = ggml_concat(c,
+                             ggml_reshape_3d(c, hint_feat, channels, batch, 1),
+                             ggml_reshape_3d(c, sample_feat, channels, batch, 1),
+                             2);
+        for (int i = 0; i < config.num_fusion_layers; i++) {
+            auto layer = std::dynamic_pointer_cast<ResidualAttentionBlock>(blocks["transformer_layes." + std::to_string(i)]);
+            x          = layer->forward(ctx, x);
+        }
+
+        auto projection = std::dynamic_pointer_cast<Linear>(blocks["spatial_ch_projs"]);
+        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);
+
+        return ggml_add(c, ggml_add(c, sample, hint), alpha);
     }
 
     ggml_tensor* resblock_forward(std::string name,
@@ -204,7 +321,9 @@ public:
                                       ggml_tensor* guided_hint,
                                       ggml_tensor* timesteps,
                                       ggml_tensor* context,
-                                      ggml_tensor* y = nullptr) {
+                                      ggml_tensor* y             = nullptr,
+                                      ggml_tensor* control_types = nullptr,
+                                      int control_type_index     = 0) {
         // x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
         // timesteps: [N,]
         // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
@@ -228,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);
@@ -246,6 +365,11 @@ public:
             emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
 
+        if (is_union()) {
+            GGML_ASSERT(control_types != nullptr);
+            emb = ggml_add(ctx->ggml_ctx, emb, control_type_embedding(ctx, control_types, x->ne[3]));
+        }
+
         std::vector<ggml_tensor*> outs;
 
         if (guided_hint == nullptr) {
@@ -257,7 +381,11 @@ public:
 
         // input block 0
         auto h = input_blocks_0_0->forward(ctx, x);
-        h      = ggml_add(ctx->ggml_ctx, h, guided_hint);
+        if (is_union()) {
+            h = fuse_condition(ctx, h, guided_hint, control_type_index);
+        } else {
+            h = ggml_add(ctx->ggml_ctx, h, guided_hint);
+        }
         outs.push_back(zero_convs_0->forward(ctx, h));
 
         // input block 1-11
@@ -308,9 +436,15 @@ public:
 
 struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
+    ControlNetUnionConfig config;
     ControlNetBlock control_net;
     std::string weight_prefix;
 
+    // ControlNet Union: index of the control type of the control image (-1: none yet)
+    // and the 0/1 vector of the types in use that the graph reads
+    int control_type_index = -1;
+    sd::Tensor<float> control_type_tensor;
+
     std::vector<ggml_tensor*> control_outputs_ggml;
     ggml_tensor* guided_hint_output_ggml = nullptr;
     std::vector<sd::Tensor<float>> controls;
@@ -320,12 +454,46 @@ struct ControlNet : public GGMLRunner {
         return "controlnet.guided_hint";
     }
 
+    bool is_union() const {
+        return config.num_control_types > 0;
+    }
+
+    // set the kind of the control image of the next generations, false (after saying why) when it
+    // can't be used: a union model needs one of the types it was trained on
+    bool set_control_type(sd_control_type_t type) {
+        if (!is_union()) {
+            if (type != SD_CONTROL_TYPE_NONE) {
+                LOG_WARN("control type '%s' ignored: the ControlNet is not a union model", sd_control_type_name(type));
+            }
+            control_type_index = -1;
+            return true;
+        }
+        if (type == SD_CONTROL_TYPE_NONE) {
+            LOG_ERROR("this ControlNet is a union model and needs a control type (openpose, depth, scribble, canny, normal, segment%s)",
+                      config.num_control_types > 6 ? ", tile, repaint" : "");
+            return false;
+        }
+        // the types of a union model, in training order: the enum without NONE
+        const int index = static_cast<int>(type) - 1;
+        if (index < 0 || index >= config.num_control_types) {
+            LOG_ERROR("control type '%s' is not supported by this ControlNet Union (it has %d control types)",
+                      sd_control_type_name(type), config.num_control_types);
+            return false;
+        }
+        control_type_index = index;
+        return true;
+    }
+
     ControlNet(ggml_backend_t backend,
                const String2TensorStorage& tensor_storage_map      = {},
                SDVersion version                                   = VERSION_SD1,
                const std::string& prefix                           = "",
                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) {
+        : GGMLRunner(backend, weight_manager),
+          version(version),
+          config(ControlNetUnionConfig::detect_from_weights(tensor_storage_map, prefix, version)),
+          control_net(version, config),
+          weight_prefix(prefix) {
         control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -376,6 +544,16 @@ struct ControlNet : public GGMLRunner {
             hint = make_input(hint_tensor);
         }
 
+        ggml_tensor* control_types = nullptr;
+        if (is_union()) {
+            const int64_t batch = x_tensor.dim() >= 4 ? x_tensor.shape()[3] : 1;
+            control_type_tensor = sd::zeros<float>({config.num_control_types * batch});
+            for (int64_t i = 0; i < batch; i++) {
+                control_type_tensor[i * config.num_control_types + control_type_index] = 1.f;
+            }
+            control_types = make_input(control_type_tensor);
+        }
+
         auto runner_ctx = get_context(gf);
 
         auto outs = control_net.forward(&runner_ctx,
@@ -384,7 +562,9 @@ struct ControlNet : public GGMLRunner {
                                         guided_hint_input,
                                         timesteps,
                                         context,
-                                        y);
+                                        y,
+                                        control_types,
+                                        control_type_index);
 
         if (guided_hint_input == nullptr && !outs.empty()) {
             guided_hint_output_ggml = outs[0];
@@ -413,6 +593,10 @@ struct ControlNet : public GGMLRunner {
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
+        if (is_union() && control_type_index < 0) {
+            LOG_ERROR("ControlNet Union: no control type set");
+            return std::nullopt;
+        }
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x, hint, timesteps, context, y);
         };

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

@@ -42,7 +42,7 @@ namespace ErnieImage {
                     config.in_channels = tensor_storage.ne[2];
                     config.hidden_size = tensor_storage.ne[3];
                 } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
-                    config.text_in_dim = tensor_storage.ne[0];
+                    config.text_in_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
                     detected_head_dim = tensor_storage.ne[0];
                 } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
@@ -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

@@ -31,12 +31,12 @@ namespace Fibo2 {
                 }
                 auto key = name.substr(prefix.size() + 1);
                 if (key == "x_embedder.weight") {
-                    config.channels = tensor.ne[0];
+                    config.channels = tensor.logical_ne0();
                     config.dim      = tensor.ne[1];
                 } else if (key == "perceiver.latent_queries") {
                     config.max_gist = tensor.ne[1];
                 } else if (key == "perceiver.layers.0.cross_attention.to_k.weight" || key == "perceiver.layers.0.cross_attention.to_kv.weight") {
-                    config.text_dim = tensor.ne[0];
+                    config.text_dim = tensor.logical_ne0();
                     config.fused_kv = key == "perceiver.layers.0.cross_attention.to_kv.weight";
                 } else if (key == "layers.0.attention.norm_q.weight") {
                     config.head_dim = tensor.ne[0];
@@ -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

@@ -158,7 +158,7 @@ namespace Flux {
                     }
                 }
                 if (ends_with(name, "txt_in.weight")) {
-                    config.context_in_dim = tensor_storage.ne[0];
+                    config.context_in_dim = tensor_storage.logical_ne0();
                     config.hidden_size    = tensor_storage.ne[1];
                 }
                 if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) {
@@ -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

@@ -27,12 +27,12 @@ namespace Flux3 {
                 std::string key = name.substr(prefix.size() + 1);
                 if (key == "txt_in.weight") {
                     detected           = true;
-                    config.context_dim = tensor.ne[0];
+                    config.context_dim = tensor.logical_ne0();
                     config.hidden_size = tensor.ne[1];
                 } else if (key == "emb_in.image.weight") {
-                    config.channels = tensor.ne[0];
+                    config.channels = tensor.logical_ne0();
                 } else if (key == "single_blocks.0.mlp_out.weight") {
-                    config.mlp_dim = tensor.ne[0];
+                    config.mlp_dim = tensor.logical_ne0();
                 } else if (key == "ref_time_in.in_layer.weight") {
                     config.flow_map = true;
                 }
@@ -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;
             };