fixed default gen amt not actually increased (+1 squashed commits)

Squashed commits:

[8e2417d83] fixed default gen amt not actually increased

CMakeLists.txt

@@ -519,6 +519,8 @@ add_library(sdtype_adapter
             otherarch/sdcpp/src/upscaler.cpp
             otherarch/sdcpp/src/runtime/guidance.cpp
             otherarch/sdcpp/src/runtime/guidance.h
+            otherarch/sdcpp/src/runtime/imatrix.cpp
+            otherarch/sdcpp/src/runtime/imatrix.h
             otherarch/sdcpp/src/stable-diffusion.cpp
             otherarch/sdcpp/thirdparty/zip.c
             otherarch/sdcpp/src/model_io/gguf_io.cpp
@@ -535,6 +537,8 @@ add_library(sdtype_adapter
             otherarch/sdcpp/src/tokenizers/gpt_oss_tokenizer.cpp
             otherarch/sdcpp/src/tokenizers/tokenizer.cpp
             otherarch/sdcpp/src/tokenizers/tokenize_util.cpp
+            otherarch/sdcpp/src/core/backend_fit.cpp
+            otherarch/sdcpp/src/core/layer_split_partition.cpp
             otherarch/sdcpp/src/core/ggml_extend_backend.cpp)
 target_include_directories(sdtype_adapter PUBLIC . ./ggml/include ./ggml/src ./ggml/src/ggml-cpu ./include ./otherarch ./otherarch/tools ./vendor/stb ./vendor/nlohmann ./vendor ./otherarch/sdcpp ./otherarch/sdcpp/include ./otherarch/sdcpp/src ./otherarch/sdcpp/examples ./tools ./common)
 target_compile_features(sdtype_adapter PUBLIC cxx_std_17) # don't bump
@@ -570,7 +574,10 @@ target_link_libraries(embeddings_adapter PRIVATE common2 ggml ${LLAMA_EXTRA_LIBS
 set_target_properties(embeddings_adapter PROPERTIES POSITION_INDEPENDENT_CODE ON)
 
 add_library(gpttype_adapter
-            gpttype_adapter.cpp)
+            gpttype_adapter.cpp
+            src/llama.cpp
+            common/chat.cpp
+            src/llama-model.cpp)
 target_include_directories(gpttype_adapter PUBLIC . ./src ./ggml/include ./ggml/src ./ggml/src/ggml-cpu ./include ./otherarch ./otherarch/tools ./vendor/stb ./vendor/nlohmann ./vendor ./otherarch/sdcpp ./otherarch/sdcpp/thirdparty ./tools ./common)
 target_compile_features(gpttype_adapter PUBLIC cxx_std_17) # don't bump
 target_link_libraries(gpttype_adapter PRIVATE common2 ggml ggml_v1 ggml_v2 ggml_v3 ${LLAMA_EXTRA_LIBS})

Makefile

@@ -101,9 +101,9 @@ NONECFLAGS =
 LLAMA_USE_BUNDLED_GLSLC := 1
 
 FAILSAFE_FLAGS = -DUSE_FAILSAFE
-VULKAN_FLAGS = -DGGML_USE_VULKAN -DSD_USE_VULKAN
+VULKAN_FLAGS = -DGGML_USE_VULKAN
 ifdef LLAMA_CUBLAS
-CUBLAS_FLAGS = -DGGML_USE_CUDA -DSD_USE_CUDA
+CUBLAS_FLAGS = -DGGML_USE_CUDA
 else
 CUBLAS_FLAGS =
 endif
@@ -215,7 +215,7 @@ OBJS_CUDA_TEMP_INST += \
     ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.o
 
 ifdef LLAMA_CUBLAS
-CUBLAS_FLAGS = -DGGML_USE_CUDA -DSD_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/x86_64-linux/include
+CUBLAS_FLAGS = -DGGML_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/x86_64-linux/include
 CUBLASLD_FLAGS = -lcuda -lcublas -lcudart -lcublasLt -lpthread -ldl -lrt -L/usr/local/cuda/lib64 -L/opt/cuda/lib64 -L$(CUDA_PATH)/targets/x86_64-linux/lib -L$(CUDA_PATH)/lib64/stubs -L/usr/local/cuda/targets/aarch64-linux/lib -L/usr/local/cuda/targets/sbsa-linux/lib -L/usr/lib/wsl/lib
 CUBLAS_OBJS = ggml-cuda.o ggml_v3-cuda.o ggml_v2-cuda.o ggml_v2-cuda-legacy.o
 CUBLAS_OBJS += $(patsubst %.cu,%.o,$(filter-out ggml/src/ggml-cuda/ggml-cuda.cu, $(wildcard ggml/src/ggml-cuda/*.cu)))
@@ -315,7 +315,7 @@ HIPFLAGS   += -DGGML_HIP_NO_ROCWMMA_FATTN
 endif
 endif
 
-HIPFLAGS   += -DGGML_USE_HIP -DGGML_HIP_NO_VMM -DGGML_USE_CUDA -DSD_USE_CUDA $(shell $(ROCM_PATH)/bin/hipconfig -C)
+HIPFLAGS   += -DGGML_USE_HIP -DGGML_HIP_NO_VMM -DGGML_USE_CUDA $(shell $(ROCM_PATH)/bin/hipconfig -C)
 HIPLDFLAGS    += -L$(ROCM_PATH)/lib -Wl,-rpath=$(ROCM_PATH)/lib
 HIPLDFLAGS    += -L$(ROCM_PATH)/lib64 -Wl,-rpath=$(ROCM_PATH)/lib64
 HIPLDFLAGS    += -lhipblas -lamdhip64 -lrocblas
@@ -339,8 +339,8 @@ endif # LLAMA_HIPBLAS
 
 
 ifdef LLAMA_METAL
-CFLAGS   += -DGGML_USE_METAL -DGGML_METAL_NDEBUG -DSD_USE_METAL
-CXXFLAGS += -DGGML_USE_METAL -DSD_USE_METAL
+CFLAGS   += -DGGML_USE_METAL -DGGML_METAL_NDEBUG
+CXXFLAGS += -DGGML_USE_METAL
 LDFLAGS  += -framework Foundation -framework Metal -framework MetalKit -framework MetalPerformanceShaders
 OBJS     += ggml-metal.o ggml-metal-device.o ggml-metal-device-m.o ggml-metal-context-m.o ggml-metal-common.o ggml-metal-ops.o
 
@@ -682,7 +682,9 @@ ggml-vulkan-shaders-noext.o: ggml/src/ggml-vulkan-shaders-noext.cpp ggml/include
 	$(CXX) $(CXXFLAGS) $(VKGEN_NOEXT_FORCE) $(VULKAN_FLAGS) -c $< -o $@
 
 # intermediate objects
-llama.o: src/llama.cpp ggml/include/ggml.h ggml/include/ggml-alloc.h ggml/include/ggml-backend.h ggml/include/ggml-cuda.h ggml/include/ggml-metal.h include/llama.h otherarch/llama-util.h
+llama.o: src/llama.cpp ggml/include/ggml.h ggml/include/ggml-alloc.h ggml/include/ggml-backend.h ggml/include/ggml-cuda.h ggml/include/ggml-metal.h include/llama.h otherarch/llama-util.h src/llama-chat.cpp src/llama-mmap.cpp src/llama-context.cpp src/llama-adapter.cpp src/llama-arch.cpp src/llama-batch.cpp src/llama-vocab.cpp src/llama-grammar.cpp src/llama-sampler.cpp src/llama-kv-cache.cpp src/llama-kv-cache-dsa.cpp src/llama-kv-cache-dsv4.cpp src/llama-kv-cache-iswa.cpp src/llama-memory-hybrid.cpp src/llama-memory-hybrid-iswa.cpp src/llama-memory-recurrent.cpp src/llama-model-loader.cpp src/llama-model-saver.cpp src/llama-quant.cpp src/llama-hparams.cpp src/llama-graph.cpp src/llama-io.cpp src/llama-memory.cpp common/fit.cpp ggml/include/ggml.h ggml/include/ggml-cpu.h ggml/include/ggml-cuda.h include/llama.h otherarch/llama-util.h
+	$(CXX) $(CXXFLAGS) -c $< -o $@
+llama-model.o: src/llama-model.cpp src/llama-model.h src/models/models.h ggml/include/ggml.h include/llama.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 common.o: common/common.cpp common/common.h common/log.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
@@ -698,8 +700,10 @@ llama-impl.o: src/llama-impl.cpp src/llama-impl.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
+chat.o: common/chat.cpp common/chat.h
+	$(CXX) $(CXXFLAGS) -c $< -o $@
 
-SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
+SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/streaming_writer.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
 
 SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp
 
@@ -732,7 +736,7 @@ OBJS_SDTYPE := otherarch/sdcpp/sdtype_adapter.o $(OBJS_SDCOMMON)
 
 LLAMASERVER_SRCS := tools/server/main.cpp tools/server/server.cpp tools/server/server-schema.cpp tools/server/server-chat.cpp tools/server/server-common.cpp tools/server/server-context.cpp tools/server/server-http.cpp tools/server/server-models.cpp tools/server/server-queue.cpp tools/server/server-task.cpp tools/server/server-tools.cpp tools/server/ui.cpp
 COMMON_DOWNLOAD_SRCS := common/download.cpp common/hf-cache.cpp vendor/cpp-httplib/httplib.cpp
-LLAMASERVER_COMMON_SRCS := common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS)
+LLAMASERVER_COMMON_SRCS := common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS)
 LLAMASERVER_CXXFLAGS := -I./tools/mtmd
 
 
@@ -755,7 +759,7 @@ music_default.o: otherarch/acestep/music_adapter.cpp
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 
 # idiotic "for easier compilation"
-GPTTYPE_ADAPTER = gpttype_adapter.cpp common/chat.cpp otherarch/llama_v2.cpp otherarch/llama_v3.cpp src/llama.cpp src/llama-chat.cpp src/llama-mmap.cpp src/llama-context.cpp src/llama-adapter.cpp src/llama-arch.cpp src/llama-batch.cpp src/llama-vocab.cpp src/llama-grammar.cpp src/llama-sampler.cpp src/llama-kv-cache.cpp src/llama-kv-cache-iswa.cpp src/llama-memory-hybrid.cpp src/llama-memory-hybrid-iswa.cpp src/llama-memory-recurrent.cpp src/llama-model-loader.cpp src/llama-model.cpp src/llama-quant.cpp src/llama-hparams.cpp otherarch/gptj_v1.cpp otherarch/gptj_v2.cpp otherarch/gptj_v3.cpp otherarch/gpt2_v1.cpp otherarch/gpt2_v2.cpp otherarch/gpt2_v3.cpp otherarch/rwkv_v2.cpp otherarch/rwkv_v3.cpp otherarch/neox_v2.cpp otherarch/neox_v3.cpp otherarch/mpt_v3.cpp ggml/include/ggml.h ggml/include/ggml-cpu.h ggml/include/ggml-cuda.h include/llama.h otherarch/llama-util.h
+GPTTYPE_ADAPTER = gpttype_adapter.cpp model_adapter.h otherarch/otherarch.h include/llama.h otherarch/llama_v2.cpp otherarch/llama_v3.cpp otherarch/gptj_v1.cpp otherarch/gptj_v2.cpp otherarch/gptj_v3.cpp otherarch/gpt2_v1.cpp otherarch/gpt2_v2.cpp otherarch/gpt2_v3.cpp otherarch/rwkv_v2.cpp otherarch/rwkv_v3.cpp otherarch/neox_v2.cpp otherarch/neox_v3.cpp otherarch/mpt_v3.cpp
 gpttype_adapter_failsafe.o: $(GPTTYPE_ADAPTER)
 	$(CXX) $(CXXFLAGS) $(FAILSAFE_FLAGS) -c $< -o $@
 gpttype_adapter.o: $(GPTTYPE_ADAPTER)
@@ -775,36 +779,36 @@ clean:
 	rm -vf otherarch/sdcpp/*.o otherarch/sdcpp/*/*.o otherarch/sdcpp/*/*/*.o otherarch/sdcpp/*/*/*/*.o
 
 # useful tools
-main: tools/completion/completion.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+main: tools/completion/main.cpp tools/completion/completion.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-mainvk: tools/completion/completion.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
-	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-fitparams: tools/fit-params/main.cpp tools/fit-params/fit-params.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
-	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-sdmain: $(OBJS_SDCOMMON) $(OBJS_SDMAIN) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+mainvk: tools/completion/main.cpp tools/completion/completion.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
+	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
+fitparams: tools/fit-params/main.cpp tools/fit-params/fit-params.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
+	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
+sdmain: $(OBJS_SDCOMMON) $(OBJS_SDMAIN) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-whispermain: otherarch/whispercpp/main.cpp otherarch/whispercpp/whisper.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+whispermain: otherarch/whispercpp/main.cpp otherarch/whispercpp/whisper.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-ttsmain: tools/tts/tts.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+ttsmain: tools/tts/tts.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-gguf-split: tools/gguf-split/gguf-split.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o build-info.h clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+gguf-split: tools/gguf-split/gguf-split.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o build-info.h clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-mtmd-cli: tools/mtmd/mtmd-cli.cpp tools/mtmd/clip.cpp common/debug.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h mtmd.o mtmd-helper.o mtmd-image.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+mtmd-cli: tools/mtmd/mtmd-cli.cpp tools/mtmd/clip.cpp common/debug.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h mtmd.o mtmd-helper.o mtmd-image.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-embedding: examples/embedding/embedding.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+embedding: examples/embedding/embedding.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-embeddingvk: examples/embedding/embedding.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
-	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-ttscppmain: otherarch/ttscpp/cli/cli.cpp otherarch/ttscpp/cli/playback.cpp otherarch/ttscpp/cli/playback.h otherarch/ttscpp/cli/write_file.cpp otherarch/ttscpp/cli/write_file.h otherarch/ttscpp/cli/vad.cpp otherarch/ttscpp/cli/vad.h otherarch/ttscpp/src/ttscpp.cpp otherarch/ttscpp/src/ttstokenizer.cpp otherarch/ttscpp/src/ttssampler.cpp otherarch/ttscpp/src/parler_model.cpp otherarch/ttscpp/src/dac_model.cpp otherarch/ttscpp/src/ttsutil.cpp otherarch/ttscpp/src/ttsargs.cpp otherarch/ttscpp/src/ttst5_encoder_model.cpp otherarch/ttscpp/src/phonemizer.cpp otherarch/ttscpp/src/tts_model.cpp otherarch/ttscpp/src/kokoro_model.cpp otherarch/ttscpp/src/dia_model.cpp otherarch/ttscpp/src/orpheus_model.cpp otherarch/ttscpp/src/snac_model.cpp otherarch/ttscpp/src/general_neural_audio_codec.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+embeddingvk: examples/embedding/embedding.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) src/llama-cparams.cpp build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
+	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
+ttscppmain: otherarch/ttscpp/cli/cli.cpp otherarch/ttscpp/cli/playback.cpp otherarch/ttscpp/cli/playback.h otherarch/ttscpp/cli/write_file.cpp otherarch/ttscpp/cli/write_file.h otherarch/ttscpp/cli/vad.cpp otherarch/ttscpp/cli/vad.h otherarch/ttscpp/src/ttscpp.cpp otherarch/ttscpp/src/ttstokenizer.cpp otherarch/ttscpp/src/ttssampler.cpp otherarch/ttscpp/src/parler_model.cpp otherarch/ttscpp/src/dac_model.cpp otherarch/ttscpp/src/ttsutil.cpp otherarch/ttscpp/src/ttsargs.cpp otherarch/ttscpp/src/ttst5_encoder_model.cpp otherarch/ttscpp/src/phonemizer.cpp otherarch/ttscpp/src/tts_model.cpp otherarch/ttscpp/src/kokoro_model.cpp otherarch/ttscpp/src/dia_model.cpp otherarch/ttscpp/src/orpheus_model.cpp otherarch/ttscpp/src/snac_model.cpp otherarch/ttscpp/src/general_neural_audio_codec.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-qwen3tts: otherarch/qwen3tts/q3ttsmain.cpp otherarch/qwen3tts/qwen3_tts.cpp otherarch/qwen3tts/text_tokenizer.cpp otherarch/qwen3tts/gguf_loader.cpp otherarch/qwen3tts/tts_transformer.cpp otherarch/qwen3tts/audio_tokenizer_decoder.cpp otherarch/qwen3tts/audio_tokenizer_encoder.cpp otherarch/qwen3tts/coreml_code_predictor_stub.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+qwen3tts: otherarch/qwen3tts/q3ttsmain.cpp otherarch/qwen3tts/qwen3_tts.cpp otherarch/qwen3tts/text_tokenizer.cpp otherarch/qwen3tts/gguf_loader.cpp otherarch/qwen3tts/tts_transformer.cpp otherarch/qwen3tts/audio_tokenizer_decoder.cpp otherarch/qwen3tts/audio_tokenizer_encoder.cpp otherarch/qwen3tts/coreml_code_predictor_stub.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-rpcserver: tools/rpc/rpc-server.cpp common/arg.cpp common/chat.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
-	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-llamaserver: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+rpcserver: tools/rpc/rpc-server.cpp common/arg.cpp common/preset.cpp $(COMMON_DOWNLOAD_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
+	$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
+llamaserver: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
-llamaservervk: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
-	$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
+llamaservervk: $(LLAMASERVER_SRCS) $(LLAMASERVER_COMMON_SRCS) build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o console.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
+	$(CXX) $(CXXFLAGS) $(LLAMASERVER_CXXFLAGS) -DGGML_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
 
 ggml/src/ggml-vulkan-shaders.cpp: ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
 ifdef VULKAN_BUILD
@@ -904,14 +908,14 @@ else
 endif
 
 #generated libraries
-koboldcpp_default: ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+koboldcpp_default: ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(DEFAULT_BUILD)
 
-koboldcpp_macos_failsafe: ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+koboldcpp_macos_failsafe: ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(DEFAULT_BUILD)
 
 ifdef FAILSAFE_BUILD
-koboldcpp_failsafe: ggml_v4_failsafe.o ggml-cpu_v4_failsafe.o ggml-ops-failsafe.o ggml-vec-failsafe.o ggml-binops.o ggml-unops.o ggml_v3_failsafe.o ggml_v2_failsafe.o ggml_v1_failsafe.o expose.o gpttype_adapter_failsafe.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FAILSAFE) $(OBJS)
+koboldcpp_failsafe: ggml_v4_failsafe.o ggml-cpu_v4_failsafe.o ggml-ops-failsafe.o ggml-vec-failsafe.o ggml-binops.o ggml-unops.o ggml_v3_failsafe.o ggml_v2_failsafe.o ggml_v1_failsafe.o expose.o gpttype_adapter_failsafe.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FAILSAFE) $(OBJS)
 	$(FAILSAFE_BUILD)
 else
 koboldcpp_failsafe:
@@ -919,7 +923,7 @@ koboldcpp_failsafe:
 endif
 
 ifdef NOAVX2_BUILD
-koboldcpp_noavx2: ggml_v4_noavx2.o ggml-cpu_v4_noavx2.o ggml-ops-noavx2.o ggml-vec-noavx2.o ggml-binops.o ggml-unops.o ggml_v3_noavx2.o ggml_v2_noavx2.o ggml_v1_failsafe.o expose.o gpttype_adapter_failsafe.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_SIMPLE) $(OBJS)
+koboldcpp_noavx2: ggml_v4_noavx2.o ggml-cpu_v4_noavx2.o ggml-ops-noavx2.o ggml-vec-noavx2.o ggml-binops.o ggml-unops.o ggml_v3_noavx2.o ggml_v2_noavx2.o ggml_v1_failsafe.o expose.o gpttype_adapter_failsafe.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_default.o tts_default.o music_default.o embeddings_default.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_SIMPLE) $(OBJS)
 	$(NOAVX2_BUILD)
 else
 koboldcpp_noavx2:
@@ -927,7 +931,7 @@ koboldcpp_noavx2:
 endif
 
 ifdef CUBLAS_BUILD
-koboldcpp_cublas: ggml_v4_cublas.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3_cublas.o ggml_v2_cublas.o ggml_v1.o expose.o gpttype_adapter_cublas.o $(OBJS_SDTYPE) whispercpp_cublas.o tts_default.o music_default.o embeddings_default.o clip_cublas.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_cublas.o ggml-repack.o $(CUBLAS_OBJS) $(OBJS_FULL) $(OBJS)
+koboldcpp_cublas: ggml_v4_cublas.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3_cublas.o ggml_v2_cublas.o ggml_v1.o expose.o gpttype_adapter_cublas.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_cublas.o tts_default.o music_default.o embeddings_default.o clip_cublas.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_cublas.o ggml-repack.o $(CUBLAS_OBJS) $(OBJS_FULL) $(OBJS)
 	$(CUBLAS_BUILD)
 else
 koboldcpp_cublas:
@@ -935,7 +939,7 @@ koboldcpp_cublas:
 endif
 
 ifdef HIPBLAS_BUILD
-koboldcpp_hipblas: ggml_v4_cublas.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3_cublas.o ggml_v2_cublas.o ggml_v1.o expose.o gpttype_adapter_cublas.o $(OBJS_SDTYPE) whispercpp_cublas.o tts_default.o music_default.o embeddings_default.o clip_cublas.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_cublas.o ggml-repack.o $(HIP_OBJS) $(OBJS_FULL) $(OBJS)
+koboldcpp_hipblas: ggml_v4_cublas.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3_cublas.o ggml_v2_cublas.o ggml_v1.o expose.o gpttype_adapter_cublas.o llama.o chat.o llama-model.o $(OBJS_SDTYPE) whispercpp_cublas.o tts_default.o music_default.o embeddings_default.o clip_cublas.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_cublas.o ggml-repack.o $(HIP_OBJS) $(OBJS_FULL) $(OBJS)
 	$(HIPBLAS_BUILD)
 else
 koboldcpp_hipblas:
@@ -943,12 +947,12 @@ koboldcpp_hipblas:
 endif
 
 ifdef VULKAN_BUILD
-koboldcpp_vulkan: ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+koboldcpp_vulkan: ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o ggml_v3.o ggml_v2.o ggml_v1.o expose.o gpttype_adapter_vulkan.o llama.o chat.o llama-model.o ggml-vulkan.o ggml-vulkan-shaders.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(VULKAN_BUILD)
 ifdef NOAVX2_BUILD
-koboldcpp_vulkan_noavx2: ggml_v4_vulkan_noavx2.o ggml-cpu_v4_noavx2.o ggml-ops-noavx2.o ggml-vec-noavx2.o ggml-binops.o ggml-unops.o ggml_v3_noavx2.o ggml_v2_noavx2.o ggml_v1_failsafe.o expose.o gpttype_adapter_vulkan_noavx2.o ggml-vulkan-noext.o ggml-vulkan-shaders-noext.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_SIMPLE) $(OBJS)
+koboldcpp_vulkan_noavx2: ggml_v4_vulkan_noavx2.o ggml-cpu_v4_noavx2.o ggml-ops-noavx2.o ggml-vec-noavx2.o ggml-binops.o ggml-unops.o ggml_v3_noavx2.o ggml_v2_noavx2.o ggml_v1_failsafe.o expose.o gpttype_adapter_vulkan_noavx2.o llama.o chat.o llama-model.o ggml-vulkan-noext.o ggml-vulkan-shaders-noext.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_SIMPLE) $(OBJS)
 	$(VULKAN_BUILD)
-koboldcpp_vulkan_failsafe: ggml_v4_vulkan_failsafe.o ggml-cpu_v4_failsafe.o ggml-ops-failsafe.o ggml-vec-failsafe.o ggml-binops.o ggml-unops.o ggml_v3_failsafe.o ggml_v2_failsafe.o ggml_v1_failsafe.o expose.o gpttype_adapter_vulkan_noavx2.o ggml-vulkan-noext.o ggml-vulkan-shaders-noext.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_SIMPLER) $(OBJS)
+koboldcpp_vulkan_failsafe: ggml_v4_vulkan_failsafe.o ggml-cpu_v4_failsafe.o ggml-ops-failsafe.o ggml-vec-failsafe.o ggml-binops.o ggml-unops.o ggml_v3_failsafe.o ggml_v2_failsafe.o ggml_v1_failsafe.o expose.o gpttype_adapter_vulkan_noavx2.o llama.o chat.o llama-model.o ggml-vulkan-noext.o ggml-vulkan-shaders-noext.o $(OBJS_SDTYPE) whispercpp_vulkan.o tts_default.o music_default.o embeddings_default.o clip_vulkan.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-repack.o $(OBJS_SIMPLER) $(OBJS)
 	$(VULKAN_BUILD)
 else
 koboldcpp_vulkan_noavx2:
@@ -966,17 +970,17 @@ koboldcpp_vulkan_failsafe:
 endif
 
 # tools
-quantize_gguf: tools/quantize/main.cpp tools/quantize/quantize.cpp common/imatrix-loader.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_gguf: tools/quantize/main.cpp tools/quantize/quantize.cpp common/imatrix-loader.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
-quantize_gptj: otherarch/tools/gptj_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_gptj: otherarch/tools/gptj_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
-quantize_gpt2: otherarch/tools/gpt2_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_gpt2: otherarch/tools/gpt2_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
-quantize_neox: otherarch/tools/neox_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_neox: otherarch/tools/neox_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
-quantize_mpt: otherarch/tools/mpt_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_mpt: otherarch/tools/mpt_quantize.cpp otherarch/tools/common-ggml.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o clip_default.o mtmd.o mtmd-helper.o mtmd-image.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
-quantize_ace: otherarch/acestep/quantize-acestep.cpp tools/mtmd/clip.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
+quantize_ace: otherarch/acestep/quantize-acestep.cpp tools/mtmd/clip.cpp ggml_v3.o ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o chat.o llama-model.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
 	$(CXX) $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
 
 

colab.ipynb

@@ -62,7 +62,7 @@
     "Model = \"https://huggingface.co/HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive/resolve/main/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_M.gguf\" #@param [\"https://huggingface.co/HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive/resolve/main/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_M.gguf\",\"https://huggingface.co/KoboldAI/LLaMA2-13B-Tiefighter-GGUF/resolve/main/LLaMA2-13B-Tiefighter.Q4_K_S.gguf\",\"https://huggingface.co/KoboldAI/LLaMA2-13B-Estopia-GGUF/resolve/main/LLaMA2-13B-Estopia.Q4_K_S.gguf\",\"https://huggingface.co/KoboldAI/Llama-3.1-8B-BookAdventures-GGUF/resolve/main/Llama-3.1-8B-BookAdventures.Q6_K.gguf\",\"https://huggingface.co/bartowski/TheDrummer_Cydonia-24B-v4.2.0-GGUF/resolve/main/TheDrummer_Cydonia-24B-v4.2.0-Q4_K_S.gguf\",\"https://huggingface.co/mradermacher/Broken-Tutu-24B-GGUF/resolve/main/Broken-Tutu-24B.Q4_K_S.gguf\",\"https://huggingface.co/bartowski/PocketDoc_Dans-PersonalityEngine-V1.3.0-24b-GGUF/resolve/main/PocketDoc_Dans-PersonalityEngine-V1.3.0-24b-Q4_K_S.gguf\",\"https://huggingface.co/LatitudeGames/Harbinger-24B-GGUF/resolve/main/Harbinger-24B-Q4_K_S.gguf\",\"https://huggingface.co/LatitudeGames/Muse-12B-GGUF/resolve/main/Muse-12B-Q4_K_S.gguf\",\"https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/resolve/main/Qwen3-VL-8B-Instruct-Q6_K.gguf\",\"https://huggingface.co/unsloth/Mistral-Small-3.2-24B-Instruct-2506-GGUF/resolve/main/Mistral-Small-3.2-24B-Instruct-2506-Q4_K_S.gguf\",\"https://huggingface.co/ggml-org/gpt-oss-20b-GGUF/resolve/main/gpt-oss-20b-mxfp4.gguf\",\"https://huggingface.co/KoboldAI/Llama-3.1-8B-BookAdventures-GGUF/resolve/main/Llama-3.1-8B-BookAdventures.Q6_K.gguf\",\"https://huggingface.co/bartowski/google_gemma-3-12b-it-GGUF/resolve/main/google_gemma-3-12b-it-Q4_K_S.gguf\",\"https://huggingface.co/unsloth/gemma-3n-E4B-it-GGUF/resolve/main/gemma-3n-E4B-it-Q6_K.gguf\",\"https://huggingface.co/unsloth/GLM-4-9B-0414-GGUF/resolve/main/GLM-4-9B-0414-Q6_K.gguf\",\"https://huggingface.co/mradermacher/Fimbulvetr-11B-v2-GGUF/resolve/main/Fimbulvetr-11B-v2.Q4_K_S.gguf\",\"https://huggingface.co/TheBloke/MythoMax-L2-13B-GGUF/resolve/main/mythomax-l2-13b.Q4_K_M.gguf\",\"https://huggingface.co/TheBloke/ReMM-SLERP-L2-13B-GGUF/resolve/main/remm-slerp-l2-13b.Q4_K_M.gguf\",\"https://huggingface.co/TheBloke/Xwin-LM-13B-v0.2-GGUF/resolve/main/xwin-lm-13b-v0.2.Q4_K_M.gguf\",\"https://huggingface.co/mradermacher/mini-magnum-12b-v1.1-GGUF/resolve/main/mini-magnum-12b-v1.1.Q4_K_S.gguf\",\"https://huggingface.co/TheBloke/Stheno-L2-13B-GGUF/resolve/main/stheno-l2-13b.Q4_K_M.gguf\",\"https://huggingface.co/TheBloke/MythoMax-L2-Kimiko-v2-13B-GGUF/resolve/main/mythomax-l2-kimiko-v2-13b.Q4_K_M.gguf\",\"https://huggingface.co/bartowski/Rocinante-12B-v1.1-GGUF/resolve/main/Rocinante-12B-v1.1-Q4_K_S.gguf\",\"https://huggingface.co/TheBloke/MistRP-Airoboros-7B-GGUF/resolve/main/mistrp-airoboros-7b.Q4_K_S.gguf\",\"https://huggingface.co/TheBloke/airoboros-mistral2.2-7B-GGUF/resolve/main/airoboros-mistral2.2-7b.Q4_K_S.gguf\",\"https://huggingface.co/concedo/KobbleTinyV2-1.1B-GGUF/resolve/main/KobbleTiny-Q4_K.gguf\",\"https://huggingface.co/grimjim/kukulemon-7B-GGUF/resolve/main/kukulemon-7B.Q8_0.gguf\",\"https://huggingface.co/mradermacher/LemonKunoichiWizardV3-GGUF/resolve/main/LemonKunoichiWizardV3.Q4_K_M.gguf\",\"https://huggingface.co/Lewdiculous/Kunoichi-DPO-v2-7B-GGUF-Imatrix/resolve/main/Kunoichi-DPO-v2-7B-Q4_K_M-imatrix.gguf\",\"https://huggingface.co/mradermacher/L3-8B-Stheno-v3.2-i1-GGUF/resolve/main/L3-8B-Stheno-v3.2.i1-Q4_K_M.gguf\",\"https://huggingface.co/Lewdiculous/Llama-3-Lumimaid-8B-v0.1-OAS-GGUF-IQ-Imatrix/resolve/main/v2-Llama-3-Lumimaid-8B-v0.1-OAS-Q4_K_M-imat.gguf\",\"https://huggingface.co/bartowski/NeuralDaredevil-8B-abliterated-GGUF/resolve/main/NeuralDaredevil-8B-abliterated-Q4_K_M.gguf\",\"https://huggingface.co/bartowski/L3-8B-Lunaris-v1-GGUF/resolve/main/L3-8B-Lunaris-v1-Q4_K_M.gguf\",\"https://huggingface.co/mradermacher/L3-Umbral-Mind-RP-v2.0-8B-GGUF/resolve/main/L3-Umbral-Mind-RP-v2.0-8B.Q4_K_M.gguf\",\"https://huggingface.co/bartowski/TheDrummer_Cydonia-24B-v2-GGUF/resolve/main/TheDrummer_Cydonia-24B-v2-Q4_K_S.gguf\",\"https://huggingface.co/bartowski/PocketDoc_Dans-PersonalityEngine-V1.2.0-24b-GGUF/resolve/main/PocketDoc_Dans-PersonalityEngine-V1.2.0-24b-IQ4_XS.gguf\",\"https://huggingface.co/mradermacher/Tlacuilo-12B-GGUF/resolve/main/Tlacuilo-12B.Q4_K_S.gguf\"] {\"allow-input\":true}\n",
     "MdCommand = \"\" #@markdown <br>\n",
     "Layers = \"Auto\" #@param [\"Auto\",\"999\"]{allow-input: true}\n",
-    "ContextSize = \"4096\" #@param [\"4096\",\"8192\",\"12288\",\"16384\"] {allow-input: true}\n",
+    "ContextSize = \"4096\" #@param [\"4096\",\"8192\",\"12288\",\"16384\",\"24576\",\"32768\",\"40960\"] {allow-input: true}\n",
     "\n",
     "#@markdown <hr>\n",
     "LoadVisionMMProjector = False #@param {type:\"boolean\"}\n",
@@ -130,7 +130,7 @@
     "  if Template == \"Gemma4 E4B Uncensored (General)\":\n",
     "    Customized = True\n",
     "    Model = \"https://huggingface.co/HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive/resolve/main/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_M.gguf\"\n",
-    "    CustomCtxSize = \"16384\"\n",
+    "    CustomCtxSize = \"40960\"\n",
     "    CustomMmproj = \"https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF/resolve/main/mmproj-BF16.gguf\"\n",
     "  if Template == \"Tiefighter 13B (General)\":\n",
     "    Customized = True\n",

common/arg.cpp

@@ -468,7 +468,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         // the first part is what gets loaded, so point params.model.path at it
         if (!url_tasks.empty()) {
             std::string first_path = url_tasks.front().local_path;
-            url_tasks.front().on_done = [&]() { params.model.path = first_path; };
+            url_tasks.front().on_done = [&, first_path]() { params.model.path = first_path; };
         }
         for (auto & task : url_tasks) {
             tasks.push_back(std::move(task));
@@ -497,13 +497,15 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     }
 
     // handle hf_plan tasks
-    auto add_tasks = [&opts, &tasks](const hf_cache::hf_files & model_files, common_params_model & model) {
+    auto add_tasks = [&opts, &tasks](const hf_cache::hf_files  & model_files,
+                                    const hf_cache::hf_file    & primary,
+                                    common_params_model        & model) {
         for (size_t i = 0; i < model_files.size(); ++i) {
             auto & model_file = model_files[i];
-            bool is_first = (i == 0);
-            tasks.emplace_back(model_file, opts, [&, is_first]() {
-                if (is_first) {
-                    // only use first part as model path
+            bool is_primary = (model_file.path == primary.path);
+            tasks.emplace_back(model_file, opts, [&, is_primary]() {
+                if (is_primary) {
+                    // the primary file is the first split (00001-of), use it as model path
                     model.path = hf_cache::finalize_file(model_file);
                 } else {
                     hf_cache::finalize_file(model_file);
@@ -512,7 +514,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         }
     };
     if (!plan.model_files.empty()) {
-        add_tasks(plan.model_files, params.model);
+        add_tasks(plan.model_files, plan.primary, params.model);
     }
     if (!plan.mmproj.local_path.empty()) {
         tasks.emplace_back(plan.mmproj, opts, [&]() {
@@ -540,12 +542,12 @@ void common_models_handler_apply(common_models_handler & handler, common_params
 
     // handle plan_spec (e.g. --spec-draft-hf)
     if (!plan_spec.model_files.empty()) {
-        add_tasks(plan_spec.model_files, params.speculative.draft.mparams);
+        add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
     }
 
     // handle vocoder plan (e.g. --hf-repo-v)
     if (!plan_voc.model_files.empty()) {
-        add_tasks(plan_voc.model_files, params.vocoder.model);
+        add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
     }
 
     // run all tasks in parallel
@@ -3297,6 +3299,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.sampling.reasoning_budget_message = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE"));
+    add_opt(common_arg(
+        {"--reasoning-preserve"},
+        {"--no-reasoning-preserve"},
+        "preserve reasoning trace in the full history, not just the last assistant message (default: template default)\n"
+        "compatible with certain templates having 'supports_preserve_reasoning' capability\n"
+        "example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking",
+        [](common_params & params, bool value) {
+            if (value) {
+                params.default_template_kwargs["preserve_reasoning"] = "true";
+            } else {
+                params.default_template_kwargs["preserve_reasoning"] = "false";
+            }
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING_PRESERVE"));
     add_opt(common_arg(
         {"--chat-template"}, "JINJA_TEMPLATE",
         string_format(
@@ -3472,7 +3488,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params) {
             params.offline = true;
         }
-    ).set_env("LLAMA_ARG_OFFLINE"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_OFFLINE"));
     add_opt(common_arg(
         {"-lv", "--verbosity", "--log-verbosity"}, "N",
         string_format("Set the verbosity threshold. Messages with a higher verbosity will be ignored. Values:\n"

common/chat.cpp

@@ -7,7 +7,6 @@
 #include "ggml.h"
 #include "json-schema-to-grammar.h"
 #include "log.h"
-#include "regex-partial.cpp"
 #include "reasoning-budget.h"
 #include "chat-auto-parser-generator.cpp"
 #include "chat-auto-parser-helpers.cpp"
@@ -926,6 +925,10 @@ static std::string common_chat_template_direct_apply_impl(
     if (inputs.add_generation_prompt) {
         inp["add_generation_prompt"] = true;
     }
+    if (inp.contains("preserve_reasoning") && inp["preserve_reasoning"].is_boolean()) {
+        bool enabled = inp["preserve_reasoning"].get<bool>();
+        jinja::caps_apply_preserve_reasoning(ctx, enabled);
+    }
 
     jinja::global_from_json(ctx, inp, inputs.mark_input);
 
@@ -2388,6 +2391,166 @@ static void func_args_not_string(json & messages) {
     }
 }
 
+// Trim leading/trailing whitespace from message contents before rendering. This
+// has to run on the messages (not on the rendered JSON) because templates with
+// string-only content caps concatenate typed content parts into a single string
+// during rendering, after which the per-part whitespace can no longer be reached.
+// Both the plain string content and the text of typed content parts are trimmed.
+static void trim_all_content(std::vector<common_chat_msg> & messages) {
+    for (auto & message : messages) {
+        message.content           = trim_whitespace(message.content);
+        message.reasoning_content = trim_whitespace(message.reasoning_content);
+        for (auto & part : message.content_parts) {
+            if (part.type == "text") {
+                part.text = trim_whitespace(part.text);
+            }
+        }
+    }
+}
+
+}
+
+// MiniCPM5 format:
+// - Reasoning: <think>{reasoning}</think> (optional)
+// - Tool calls: <function name="foo"><param name="bar">value</param></function>
+static common_chat_params common_chat_params_init_minicpm5(const common_chat_template &          tmpl,
+                                                           const autoparser::generation_params & inputs) {
+    common_chat_params data;
+
+    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs);
+    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
+    data.supports_thinking = true;
+    data.preserved_tokens  = {
+        "<function",
+        "<param",
+        "</function>",
+        "</param>",
+        "<think>",
+        "</think>",
+    };
+
+    data.thinking_start_tag = "<think>";
+    data.thinking_end_tag   = "</think>";
+
+    data.message_delimiters = {
+        { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant"             },
+        { COMMON_CHAT_ROLE_TOOL,      "<|im_start|>user\n<tool_response>" },
+        { COMMON_CHAT_ROLE_USER,      "<|im_start|>user"                  },
+        { COMMON_CHAT_ROLE_SYSTEM,    "<|im_start|>system"                },
+    };
+
+    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
+    auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty();
+    auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
+    auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
+
+    if (inputs.has_continuation()) {
+        const auto & msg = inputs.continue_msg;
+
+        data.generation_prompt = "<|im_start|>assistant\n<think>\n" + msg.reasoning_content;
+        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
+            data.generation_prompt += "\n</think>\n\n" + msg.render_content();
+        }
+
+        data.prompt += data.generation_prompt;
+    }
+
+    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
+        auto generation_prompt = p.literal("<|im_start|>assistant\n");
+
+        auto reasoning = p.eps();
+        if (extract_reasoning) {
+            reasoning = ("<think>" << p.reasoning(p.until("</think>")) << "</think>") + p.space();
+        }
+
+        // Response format parser
+        if (has_response_format) {
+            return generation_prompt + reasoning + p.content(p.schema(p.json(), "response-format", inputs.json_schema));
+        }
+
+        if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
+            // CDATA lets a value carry characters that would otherwise close the tag (e.g.
+            // </param>); capture the inner text only, excluding the CDATA markers.
+            auto string_value = p.choice({
+                p.literal("<![CDATA[") + p.ac(p.tool_arg_string_value(p.until("]]>")) + p.literal("]]>"), "]]>") + p.tool_arg_close(p.literal("</param>")),
+                p.negate(p.literal("<![CDATA[")) + p.ac(p.tool_arg_string_value(p.until("</param>")) + p.tool_arg_close(p.literal("</param>")), "</param>")
+            });
+
+            auto tool_choice = p.choice();
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto &      function = tool.at("function");
+                const std::string name     = function.at("name");
+                auto              params   = function.contains("parameters") ? function.at("parameters") : json::object();
+
+                auto args = p.eps();
+                if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) {
+                    auto schema_info = common_schema_info();
+                    schema_info.resolve_refs(params);
+
+                    auto arg_choice = p.choice();
+                    for (const auto & [prop_name, prop_schema] : params.at("properties").items()) {
+                        auto value_parser = p.eps();
+                        if (schema_info.resolves_to_string(prop_schema)) {
+                            value_parser = string_value;
+                        } else {
+                            value_parser = p.tool_arg_json_value(
+                                    p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false)
+                                ) + p.tool_arg_close(p.literal("</param>"));
+                        }
+
+                        auto arg_rule = p.tool_arg(
+                            p.tool_arg_open(p.literal("<param name=\"") + p.tool_arg_name(p.literal(prop_name)) + p.literal("\">")) +
+                            value_parser
+                        );
+
+                        arg_choice |= arg_rule;
+                    }
+                    args = p.zero_or_more(arg_choice + p.space());
+                }
+
+                auto tool_parser = p.tool(
+                    p.tool_open(p.literal("<function name=\"") + p.tool_name(p.literal(name)) + p.literal("\">"))
+                    << p.tool_args(args)
+                    << p.tool_close(p.literal("</function>")));
+
+                tool_choice |= p.rule("tool-" + name, tool_parser);
+            });
+
+            auto max_calls  = inputs.parallel_tool_calls ? -1 : 1;
+            auto tool_calls = p.trigger_rule("tool-call", p.repeat(tool_choice + p.space(), 1, max_calls));
+
+            auto content = p.content(p.until("<function"));
+
+            return generation_prompt + reasoning + content + tool_calls + p.end();
+        }
+
+        return generation_prompt + reasoning + p.content(p.rest()) + p.end();
+    });
+
+    data.parser = parser.save();
+
+    if (include_grammar) {
+        data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
+        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function = tool.at("function");
+                auto         schema   = function.contains("parameters") ? function.at("parameters") : json::object();
+                builder.resolve_refs(schema);
+            });
+            if (has_response_format) {
+                auto schema = inputs.json_schema;
+                builder.resolve_refs(schema);
+            }
+            parser.build_grammar(builder, data.grammar_lazy);
+        });
+
+        data.grammar_triggers = {
+            { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<function" },
+        };
+    }
+
+    return data;
 }
 
 static json common_chat_extra_context() {
@@ -2482,6 +2645,14 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_gemma4(tmpl, params);
     }
 
+    // MiniCPM5 - XML tool calls with <function name="..."><param name="...">...</param></function>
+    if (src.find("Tool usage guidelines:") != std::string::npos &&
+        src.find("<function name=\"") != std::string::npos &&
+        src.find("<param name=\"") != std::string::npos) {
+        LOG_DBG("Using specialized template: MiniCPM5\n");
+        return common_chat_params_init_minicpm5(tmpl, params);
+    }
+
     return std::nullopt;
 }
 
@@ -2493,7 +2664,16 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
         params.tools.is_array() && tmpls->template_tool_use ? *tmpls->template_tool_use : *tmpls->template_default;
     const auto & src             = tmpl.source();
     const auto & caps            = tmpl.original_caps();
-    params.messages              = render_message_to_json(inputs.messages, tmpl.original_caps());
+    std::vector<common_chat_msg>        trimmed_messages;
+    const std::vector<common_chat_msg> * messages_to_render = &inputs.messages;
+    if (src.find("You have access to the following functions in JSONSchema format") != std::string::npos) {
+        // StepFun: trim message contents (including typed content parts) before rendering,
+        // otherwise leftover whitespace drives the model into reasoning loops (issue #24181)
+        trimmed_messages   = inputs.messages;
+        workaround::trim_all_content(trimmed_messages);
+        messages_to_render = &trimmed_messages;
+    }
+    params.messages              = render_message_to_json(*messages_to_render, tmpl.original_caps());
     params.tool_choice           = inputs.tool_choice;
     params.reasoning_format      = inputs.reasoning_format;
     params.enable_thinking       = inputs.enable_thinking;

common/common.cpp

@@ -231,7 +231,7 @@ bool set_process_priority(enum ggml_sched_priority prio) {
     }
 
     if (!SetPriorityClass(GetCurrentProcess(), p)) {
-        LOG_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError());
+        COM_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError());
         return false;
     }
 
@@ -257,7 +257,7 @@ bool set_process_priority(enum ggml_sched_priority prio) {
     }
 
     if (setpriority(PRIO_PROCESS, 0, p) != 0) {
-        LOG_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno);
+        COM_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno);
         return false;
     }
     return true;
@@ -290,14 +290,14 @@ void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_para
 
     if (n_set && n_set < cpuparams.n_threads) {
         // Not enough set bits, may experience performance issues.
-        LOG_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads);
+        COM_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads);
     }
 }
 
 bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THREADS]) {
     size_t dash_loc = range.find('-');
     if (dash_loc == std::string::npos) {
-        LOG_ERR("Format of CPU range is invalid! Expected [<start>]-[<end>].\n");
+        COM_ERR("%s", "Format of CPU range is invalid! Expected [<start>]-[<end>].\n");
         return false;
     }
 
@@ -309,7 +309,7 @@ bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THRE
     } else {
         start_i = std::stoull(range.substr(0, dash_loc));
         if (start_i >= GGML_MAX_N_THREADS) {
-            LOG_ERR("Start index out of bounds!\n");
+            COM_ERR("%s", "Start index out of bounds!\n");
             return false;
         }
     }
@@ -319,7 +319,7 @@ bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THRE
     } else {
         end_i = std::stoull(range.substr(dash_loc + 1));
         if (end_i >= GGML_MAX_N_THREADS) {
-            LOG_ERR("End index out of bounds!\n");
+            COM_ERR("%s", "End index out of bounds!\n");
             return false;
         }
     }
@@ -339,7 +339,7 @@ bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREAD
     }
 
     size_t num_digits = mask.length() - start_i;
-    if (num_digits > 128) num_digits = 128;
+    num_digits = std::min<size_t>(num_digits, 128);
 
     size_t end_i = num_digits + start_i;
 
@@ -354,7 +354,7 @@ bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREAD
         } else if (c >= 'A' && c <= 'F') {
             id -= 'A' - 10;
         } else {
-            LOG_ERR("Invalid hex character '%c' at position %d\n", c, int32_t(i));
+            COM_ERR("Invalid hex character '%c' at position %d\n", c, int32_t(i));
             return false;
         }
 
@@ -385,21 +385,21 @@ void common_params_print_info(const common_params & params, bool print_devices)
 #else
     const char * build_type = " (debug)";
 #endif
-    LOG_TRC("%s: build %d (%s) with %s for %s%s\n", __func__, llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type);
+    COM_TRC("%s: build %d (%s) with %s for %s%s\n", __func__, llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type);
 
-    LOG_INF("log_info: verbosity = %d (adjust with the `-lv N` CLI arg)\n", common_log_get_verbosity_thold());
+    COM_INF("%s: verbosity = %d (adjust with the `-lv N` CLI arg)\n", __func__, common_log_get_verbosity_thold());
 
     // device enumeration creates a primary context on CUDA backends, skip it when the caller does not own any device
     if (print_devices) {
-        LOG_INF("device_info:\n");
+        COM_TRC("%s", "device_info:\n");
         for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
             auto * dev = ggml_backend_dev_get(i);
             size_t free, total;
             ggml_backend_dev_memory(dev, &free, &total);
-            LOG_INF("  - %-8s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024);
+            COM_TRC("  - %-8s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024);
         }
     }
-    LOG_INF("%s\n", common_params_get_system_info(params).c_str());
+    COM_TRC("%s\n", common_params_get_system_info(params).c_str());
 }
 
 std::string common_params_get_system_info(const common_params & params) {
@@ -666,7 +666,7 @@ void string_process_escapes(std::string & input) {
 bool string_parse_kv_override(const char * data, std::vector<llama_model_kv_override> & overrides) {
     const char * sep = strchr(data, '=');
     if (sep == nullptr || sep - data >= 128) {
-        LOG_ERR("%s: malformed KV override '%s'\n", __func__, data);
+        COM_ERR("%s: malformed KV override '%s'\n", __func__, data);
         return false;
     }
     llama_model_kv_override kvo;
@@ -689,20 +689,20 @@ bool string_parse_kv_override(const char * data, std::vector<llama_model_kv_over
         } else if (std::strcmp(sep, "false") == 0) {
             kvo.val_bool = false;
         } else {
-            LOG_ERR("%s: invalid boolean value for KV override '%s'\n", __func__, data);
+            COM_ERR("%s: invalid boolean value for KV override '%s'\n", __func__, data);
             return false;
         }
     } else if (strncmp(sep, "str:", 4) == 0) {
         sep += 4;
         kvo.tag = LLAMA_KV_OVERRIDE_TYPE_STR;
         if (strlen(sep) > 127) {
-            LOG_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data);
+            COM_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data);
             return false;
         }
         strncpy(kvo.val_str, sep, 127);
         kvo.val_str[127] = '\0';
     } else {
-        LOG_ERR("%s: invalid type for KV override '%s'\n", __func__, data);
+        COM_ERR("%s: invalid type for KV override '%s'\n", __func__, data);
         return false;
     }
     overrides.emplace_back(std::move(kvo));
@@ -1205,8 +1205,8 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     auto cparams = common_context_params_to_llama(params);
 
     if (params.fit_params) {
-        LOG_INF("%s: fitting params to device memory ...\n", __func__);
-        LOG_INF("%s: (for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n", __func__);
+        COM_TRC("%s", "fitting params to device memory ...\n");
+        COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
         common_fit_params(params.model.path.c_str(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
@@ -1233,7 +1233,7 @@ common_init_result::common_init_result(common_params & params, bool model_only)
         llama_adapter_lora_ptr lora;
         lora.reset(llama_adapter_lora_init(model, la.path.c_str()));
         if (lora == nullptr) {
-            LOG_ERR("%s: failed to load lora adapter '%s'\n", __func__, la.path.c_str());
+            COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str());
             pimpl->model.reset(model);
             return;
         }
@@ -1252,14 +1252,14 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     common_init_sampler_from_model(model, params.sampling);
 
     if (params.sampling.ignore_eos && llama_vocab_eos(vocab) == LLAMA_TOKEN_NULL) {
-        LOG_WRN("%s: warning: vocab does not have an EOS token, ignoring --ignore-eos\n", __func__);
+        COM_WRN("%s", "vocab does not have an EOS token, ignoring --ignore-eos\n");
         params.sampling.ignore_eos = false;
     }
 
     // initialize once
     for (llama_token i = 0; i < llama_vocab_n_tokens(vocab); i++) {
         if (llama_vocab_is_eog(vocab, i)) {
-            LOG_TRC("%s: added %s logit bias = %f\n", __func__, common_token_to_piece(vocab, i).c_str(), -INFINITY);
+            COM_TRC("added %s logit bias = %f\n", common_token_to_piece(vocab, i).c_str(), -INFINITY);
             params.sampling.logit_bias_eog.push_back({i, -INFINITY});
         }
     }
@@ -1297,7 +1297,7 @@ common_init_result::common_init_result(common_params & params, bool model_only)
 
     llama_context * lctx = llama_init_from_model(model, cparams);
     if (lctx == NULL) {
-        LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str());
+        COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str());
         return;
     }
 
@@ -1334,7 +1334,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 
     llama_model * model = res->model();
     if (model == NULL) {
-        LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str());
+        COM_ERR("failed to load model '%s'\n", params.model.path.c_str());
         return res;
     }
 
@@ -1344,14 +1344,14 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 
     llama_context * lctx = res->context();
     if (lctx == NULL) {
-        LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str());
+        COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str());
         return res;
     }
 
     const llama_vocab * vocab = llama_model_get_vocab(model);
 
     if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) {
-        LOG_WRN("%s: KV cache shifting is not supported for this context, disabling KV cache shifting\n", __func__);
+        COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n");
         params.ctx_shift = false;
     }
 
@@ -1380,7 +1380,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
         bool ok = true;
 
         if (llama_vocab_bos(vocab) == LLAMA_TOKEN_NULL) {
-            LOG_WRN("%s: warning: vocab does not have a  BOS token, reranking will not work\n", __func__);
+            COM_WRN("%s", "vocab does not have a  BOS token, reranking will not work\n");
             ok = false;
         }
 
@@ -1389,10 +1389,10 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
         bool has_rerank_prompt = llama_model_chat_template(model, "rerank") != NULL;
 
         if (!has_eos && !has_sep && !has_rerank_prompt) {
-            LOG_WRN("%s: warning: vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n", __func__);
+            COM_WRN("%s", "vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n");
             ok = false;
         } else if (!has_eos) {
-            LOG_WRN("%s: warning: vocab does not have an EOS token, using SEP token as fallback\n", __func__);
+            COM_WRN("%s", "vocab does not have an EOS token, using SEP token as fallback\n");
         }
 
         if (!ok) {
@@ -1405,7 +1405,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
     }
 
     if (params.warmup) {
-        LOG_INF("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__);
+        COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n");
 
         std::vector<llama_token> tmp;
         llama_token bos = llama_vocab_bos(vocab);
@@ -1479,20 +1479,20 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) {
 
     int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size()));
     if (ret != 0) {
-        LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret);
+        COM_ERR("llama_decode() failed: %d\n", ret);
         res = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
         goto done;
     }
 
     if (llama_n_rs_seq(ctx) > 0) {
-        LOG_INF("%s: the context supports bounded partial sequence removal\n", __func__);
+        COM_TRC("%s", "the context supports bounded partial sequence removal\n");
         res = COMMON_CONTEXT_SEQ_RM_TYPE_RS;
         goto done;
     }
 
     // try to remove the last tokens
     if (!llama_memory_seq_rm(mem, 0, 1, -1)) {
-        LOG_TRC("%s: the context does not support partial sequence removal\n", __func__);
+        COM_TRC("%s", "the context does not support partial sequence removal\n");
         res = COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
         goto done;
     }
@@ -1809,13 +1809,13 @@ static common_control_vector_data common_control_vector_load_one(const common_co
     };
     struct gguf_context * ctx_gguf = gguf_init_from_file(load_info.fname.c_str(), meta_gguf_params);
     if (!ctx_gguf) {
-        LOG_ERR("%s: failed to load control vector file from %s\n", __func__, load_info.fname.c_str());
+        COM_ERR("failed to load control vector file from %s\n", load_info.fname.c_str());
         return result;
     }
 
     int32_t n_tensors = gguf_get_n_tensors(ctx_gguf);
     if (n_tensors == 0) {
-        LOG_WRN("%s: no direction tensors found in %s\n", __func__, load_info.fname.c_str());
+        COM_WRN("no direction tensors found in %s\n", load_info.fname.c_str());
     }
 
     for (int i = 0; i < n_tensors; i++) {
@@ -1833,23 +1833,23 @@ static common_control_vector_data common_control_vector_load_one(const common_co
             }
         }
         if (layer_idx < 0) {
-            LOG_ERR("%s: invalid/unparsable direction tensor layer index in %s\n", __func__, load_info.fname.c_str());
+            COM_ERR("invalid/unparsable direction tensor layer index in %s\n", load_info.fname.c_str());
             result.n_embd = -1;
             break;
         } else if (layer_idx == 0) {
-            LOG_ERR("%s: invalid (zero) direction tensor layer index in %s\n", __func__, load_info.fname.c_str());
+            COM_ERR("invalid (zero) direction tensor layer index in %s\n", load_info.fname.c_str());
             result.n_embd = -1;
             break;
         }
 
         struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str());
         if (tensor->type != GGML_TYPE_F32) {
-            LOG_ERR("%s: invalid (non-F32) direction tensor type in %s\n", __func__, load_info.fname.c_str());
+            COM_ERR("invalid (non-F32) direction tensor type in %s\n", load_info.fname.c_str());
             result.n_embd = -1;
             break;
         }
         if (ggml_n_dims(tensor) != 1) {
-            LOG_ERR("%s: invalid (non-1D) direction tensor shape in %s\n", __func__, load_info.fname.c_str());
+            COM_ERR("invalid (non-1D) direction tensor shape in %s\n", load_info.fname.c_str());
             result.n_embd = -1;
             break;
         }
@@ -1857,7 +1857,7 @@ static common_control_vector_data common_control_vector_load_one(const common_co
         if (result.n_embd == -1) {
             result.n_embd = ggml_nelements(tensor);
         } else if (ggml_nelements(tensor) != result.n_embd) {
-            LOG_ERR("%s: direction tensor in %s does not match previous dimensions\n", __func__, load_info.fname.c_str());
+            COM_ERR("direction tensor in %s does not match previous dimensions\n", load_info.fname.c_str());
             result.n_embd = -1;
             break;
         }
@@ -1874,7 +1874,7 @@ static common_control_vector_data common_control_vector_load_one(const common_co
     }
 
     if (result.n_embd == -1) {
-        LOG_WRN("%s: skipping %s due to invalid direction tensors\n", __func__, load_info.fname.c_str());
+        COM_WRN("skipping %s due to invalid direction tensors\n", load_info.fname.c_str());
         result.data.clear();
     }
 
@@ -1895,7 +1895,7 @@ common_control_vector_data common_control_vector_load(const std::vector<common_c
             break;
         }
         if (result.n_embd != -1 && result.n_embd != cur.n_embd) {
-            LOG_ERR("%s: control vectors in %s does not match previous dimensions\n", __func__, info.fname.c_str());
+            COM_ERR("control vectors in %s does not match previous dimensions\n", info.fname.c_str());
             result.n_embd = -1;
             break;
         }
@@ -1911,7 +1911,7 @@ common_control_vector_data common_control_vector_load(const std::vector<common_c
     }
 
     if (result.n_embd == -1) {
-        LOG_ERR("%s: no valid control vector files passed\n", __func__);
+        COM_ERR("%s", "no valid control vector files passed\n");
         result.data.clear();
     }
 
@@ -2022,13 +2022,13 @@ bool common_prompt_batch_decode(
         // memory, so we can't just remove the last token from the memory and replay the last token which
         // is the reason for this logic.
         if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(all_tokens.data() + offset), n_tokens_before_last))) {
-            LOG_ERR("%s : failed to eval\n", __func__);
+            COM_ERR("%s", "failed to eval\n");
             return false;
         }
         n_past += n_tokens_before_last;
 
         llama_state_save_file(ctx, state_path.data(), all_tokens.data(), all_tokens.size());
-        LOG_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size());
+        COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size());
 
         llama_token last_token = all_tokens.back();
         llama_batch batch = llama_batch_get_one(&last_token, 1);
@@ -2036,13 +2036,13 @@ bool common_prompt_batch_decode(
         batch.pos = &pos;
 
         if (llama_decode(ctx, batch)) {
-            LOG_ERR("%s : failed to eval last token\n", __func__);
+            COM_ERR("%s", "failed to eval last token\n");
             return false;
         }
         n_past++;
     } else {
         if (llama_decode(ctx, llama_batch_get_one(const_cast<llama_token*>(all_tokens.data() + offset), n_new))) {
-            LOG_ERR("%s : failed to eval\n", __func__);
+            COM_ERR("%s", "failed to eval\n");
             return false;
         }
         n_past += n_new;

common/common.h

@@ -26,6 +26,13 @@
 #define DIRECTORY_SEPARATOR '/'
 #endif // _WIN32
 
+#define COM_DBG(fmt, ...) LOG_DBG("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define COM_TRC(fmt, ...) LOG_TRC("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define COM_INF(fmt, ...) LOG_INF("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define COM_WRN(fmt, ...) LOG_WRN("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define COM_ERR(fmt, ...) LOG_ERR("cmn  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define COM_CNT(fmt, ...) LOG_CNT(""              fmt,               __VA_ARGS__)
+
 #define die(msg)          do { fputs("error: " msg "\n", stderr);                exit(1); } while (0)
 #define die_fmt(fmt, ...) do { fprintf(stderr, "error: " fmt "\n", __VA_ARGS__); exit(1); } while (0)
 
@@ -163,6 +170,7 @@ enum common_speculative_type {
     COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE,  // standalone draft model speculative decoding
     COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3,  // Eagle3 speculative decoding
     COMMON_SPECULATIVE_TYPE_DRAFT_MTP,     // Multi-token prediction
+    COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH,  // DFlash speculative decoding
     COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE,  // simple self-speculative decoding based on n-grams
     COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K,   // self-speculative decoding with n-gram keys only
     COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
@@ -378,7 +386,7 @@ struct common_params_speculative {
 
     uint32_t need_n_rs_seq() const {
         bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
-            return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3;
+            return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH;
         });
 
         return needs_rs_seq ? draft.n_max : 0u;

common/fit.cpp

@@ -233,7 +233,7 @@ static void common_params_fit_impl(
         sum_projected_used = dmds_full.back().mb.total();
         sum_free           = dmds_full.back().total;
         sum_projected_free = sum_free - sum_projected_used;
-        LOG_INF("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n",
+        LOG_TRC("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n",
             __func__, sum_projected_used/MiB, sum_free/MiB);
         if (sum_projected_free >= margins[0]) {
             LOG_TRC("%s: will leave %" PRId64 " >= %" PRId64 " MiB of system memory, no changes needed\n",

common/http.h

@@ -11,6 +11,11 @@ struct common_http_url {
     std::string path;
 };
 
+// bracket an IPv6 literal host for a URL authority (RFC 3986)
+static std::string common_http_format_host(const std::string & host) {
+    return host.find(':') != std::string::npos ? "[" + host + "]" : host;
+}
+
 static common_http_url common_http_parse_url(const std::string & url) {
     common_http_url parts;
     auto scheme_end = url.find("://");
@@ -49,11 +54,28 @@ static common_http_url common_http_parse_url(const std::string & url) {
         parts.path = "/";
     }
 
-    auto colon_pos = parts.host.find(':');
+    // split the authority into host and optional port, a bracketed IPv6 literal keeps its inner colons (RFC 3986)
+    std::string port_str;
+    if (!parts.host.empty() && parts.host.front() == '[') {
+        auto close = parts.host.find(']');
+        if (close == std::string::npos) {
+            throw std::runtime_error("invalid IPv6 URL authority: " + parts.host);
+        }
+        auto after = parts.host.substr(close + 1);
+        if (!after.empty() && after.front() == ':') {
+            port_str = after.substr(1);
+        }
+        parts.host = parts.host.substr(1, close - 1);
+    } else {
+        auto colon_pos = parts.host.find(':');
+        if (colon_pos != std::string::npos) {
+            port_str = parts.host.substr(colon_pos + 1);
+            parts.host = parts.host.substr(0, colon_pos);
+        }
+    }
 
-    if (colon_pos != std::string::npos) {
-        parts.port = std::stoi(parts.host.substr(colon_pos + 1));
-        parts.host = parts.host.substr(0, colon_pos);
+    if (!port_str.empty()) {
+        parts.port = std::stoi(port_str);
     } else if (parts.scheme == "http") {
         parts.port = 80;
     } else if (parts.scheme == "https") {
@@ -83,7 +105,7 @@ static std::pair<httplib::Client, common_http_url> common_http_client(const std:
     }
 #endif
 
-    httplib::Client cli(parts.scheme + "://" + parts.host + ":" + std::to_string(parts.port));
+    httplib::Client cli(parts.scheme + "://" + common_http_format_host(parts.host) + ":" + std::to_string(parts.port));
 
     if (!parts.user.empty()) {
         cli.set_basic_auth(parts.user, parts.password);
@@ -95,5 +117,5 @@ static std::pair<httplib::Client, common_http_url> common_http_client(const std:
 }
 
 static std::string common_http_show_masked_url(const common_http_url & parts) {
-    return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + parts.host + parts.path;
+    return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + common_http_format_host(parts.host) + parts.path;
 }

common/jinja/caps.cpp

@@ -9,6 +9,9 @@
 #include <functional>
 #include <sstream>
 
+#ifdef FILENAME
+#undef FILENAME
+#endif
 #define FILENAME "jinja-caps"
 
 using json = nlohmann::ordered_json;
@@ -16,22 +19,34 @@ using json = nlohmann::ordered_json;
 namespace jinja {
 
 using caps_json_fn = std::function<json()>;
-using caps_analyze_fn = std::function<void(bool, value &, value &)>;
+using caps_ctx_fn = std::function<void(context &)>;
+using caps_analyze_fn = std::function<void(bool, value &, value &, const std::string &)>;
+
+void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) {
+    ctx.set_val("preserve_thinking",         mk_val<value_bool>(enabled));
+    ctx.set_val("clear_thinking",            mk_val<value_bool>(!enabled));
+    ctx.set_val("truncate_history_thinking", mk_val<value_bool>(!enabled));
+}
 
 static void caps_try_execute(jinja::program & prog,
                              const caps_json_fn & messages_fn,
+                             const caps_ctx_fn & ctx_fn,
                              const caps_json_fn & tools_fn,
                              const caps_analyze_fn & analyze_fn) {
     context ctx;
     ctx.is_get_stats = true;
     jinja::global_from_json(ctx, json{
         {"messages", messages_fn()},
-        {"tools", tools_fn()},
+        {"tools", tools_fn ? tools_fn() : json::array()},
         {"bos_token", ""},
         {"eos_token", ""},
         {"add_generation_prompt", true}
     }, true);
 
+    if (ctx_fn) {
+        ctx_fn(ctx);
+    }
+
     auto messages = ctx.get_val("messages");
     auto tools = ctx.get_val("tools");
 
@@ -49,7 +64,7 @@ static void caps_try_execute(jinja::program & prog,
         // ignore exceptions during capability analysis
     }
 
-    analyze_fn(success, messages, tools);
+    analyze_fn(success, messages, tools, result);
 }
 
 // for debugging only
@@ -109,11 +124,9 @@ caps caps_get(jinja::program & prog) {
                 }
             });
         },
-        [&]() {
-            // tools
-            return json{nullptr};
-        },
-        [&](bool success, value & messages, value &) {
+        nullptr, // ctx_fn
+        nullptr, // tools_fn
+        [&](bool success, value & messages, value &, const std::string &) {
             auto & content = messages->at(0)->at("content");
             caps_print_stats(content, "messages[0].content");
             if (has_op(content, "selectattr") || has_op(content, "array_access")) {
@@ -145,11 +158,9 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&]() {
-            // tools
-            return json::array();
-        },
-        [&](bool, value & messages, value &) {
+        nullptr, // ctx_fn
+        nullptr, // tools_fn
+        [&](bool, value & messages, value &, const std::string &) {
             auto & content = messages->at(0)->at("content");
             caps_print_stats(content, "messages[0].content");
             if (!content->stats.used) {
@@ -201,6 +212,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
+        nullptr, // ctx_fn
         [&]() {
             // tools
             return json::array({
@@ -224,7 +236,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&](bool success, value & messages, value & tools) {
+        [&](bool success, value & messages, value & tools, const std::string &) {
             if (!success) {
                 return; // Nothing can be inferred
             }
@@ -293,6 +305,7 @@ caps caps_get(jinja::program & prog) {
                     },
                 });
             },
+            nullptr, // ctx_fn
             [&]() {
                 // tools
                 return json::array({
@@ -316,7 +329,7 @@ caps caps_get(jinja::program & prog) {
                     },
                 });
             },
-            [&](bool success, value & messages, value & tools) {
+            [&](bool success, value & messages, value & tools, const std::string &) {
                 if (!success) {
                     result.supports_tool_calls = false;
                     result.supports_tools = false;
@@ -394,6 +407,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
+        nullptr, // ctx_fn
         [&]() {
             // tools
             return json::array({
@@ -417,7 +431,7 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&](bool success, value & messages, value & /*tools*/) {
+        [&](bool success, value & messages, value &, const std::string &) {
             if (!success) {
                 result.supports_parallel_tool_calls = false;
                 return;
@@ -438,11 +452,22 @@ caps caps_get(jinja::program & prog) {
     JJ_DEBUG("%s\n", ">>> Running capability check: preserve reasoning");
 
     // case: preserve reasoning content in chat history
+    const std::string reasoning_placeholder = "<REASONING_CONTENT_PLACEHOLDER>";
     caps_try_execute(
         prog,
         [&]() {
             // messages
             return json::array({
+                {
+                    {"role", "user"},
+                    {"content", "User message"}
+                },
+                {
+                    {"role", "assistant"},
+                    {"content", "Assistant message"},
+                    // check of reasoning_content deeper in the history, not just the last assistant message
+                    {"reasoning_content", reasoning_placeholder}
+                },
                 {
                     {"role", "user"},
                     {"content", "User message"}
@@ -458,14 +483,13 @@ caps caps_get(jinja::program & prog) {
                 },
             });
         },
-        [&]() {
-            // tools
-            return json::array();
+        [&](context & ctx) {
+            caps_apply_preserve_reasoning(ctx, true);
         },
-        [&](bool, value & messages, value &) {
-            auto & content = messages->at(1)->at("reasoning_content");
-            caps_print_stats(content, "messages[1].reasoning_content");
-            if (content->stats.used) {
+        nullptr, // tools_fn
+        [&](bool, value &, value &, const std::string & output) {
+            // note: we cannot use stats here because the reasoning_content may be used for "if" condition test, but not actually outputted in the final result
+            if (output.find(reasoning_placeholder) != std::string::npos) {
                 result.supports_preserve_reasoning = true;
             }
         }

common/jinja/caps.h

@@ -12,7 +12,9 @@ struct caps {
     bool supports_tool_calls = true;
     bool supports_system_role = true;
     bool supports_parallel_tool_calls = true;
-    bool supports_preserve_reasoning = false; // support assistant message with reasoning_content
+
+    // supports preserve reasoning trace in the full history, not just the last assistant message
+    bool supports_preserve_reasoning = false;
 
     // one of the 2 content capabilities must be true
     bool supports_string_content = true;
@@ -29,4 +31,6 @@ struct caps {
 
 caps caps_get(jinja::program & prog);
 
+void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled);
+
 } // namespace jinja

common/jinja/lexer.cpp

@@ -7,6 +7,9 @@
 #include <string>
 #include <vector>
 
+#ifdef FILENAME
+#undef FILENAME
+#endif
 #define FILENAME "jinja-lexer"
 
 namespace jinja {

common/jinja/parser.cpp

@@ -8,6 +8,9 @@
 #include <string>
 #include <vector>
 
+#ifdef FILENAME
+#undef FILENAME
+#endif
 #define FILENAME "jinja-parser"
 
 namespace jinja {

common/jinja/runtime.cpp

@@ -8,6 +8,9 @@
 #include <memory>
 #include <cmath>
 
+#ifdef FILENAME
+#undef FILENAME
+#endif
 #define FILENAME "jinja-runtime"
 
 bool g_jinja_debug = false;
@@ -954,4 +957,50 @@ value keyword_argument_expression::execute_impl(context & ctx) {
     return mk_val<value_kwarg>(k, v);
 }
 
+std::string runtime::debug_dump_program(const program & prog, const std::string & src) {
+    std::ostringstream oss;
+    size_t lvl = 0;
+    context ctx;
+    ctx.src.reset(new std::string(src));
+
+    auto indent = [](size_t lvl) -> std::string {
+        return std::string(lvl * 2, ' ');
+    };
+
+    ctx.visitor = [&](bool is_leaf, statement * node, std::vector<visitor_pair> children) {
+        oss << indent(lvl) << node->type() << ":\n";
+        lvl++;
+        if (is_leaf) {
+            const auto & pos = node->pos;
+            oss << indent(lvl) << "(leaf) at " << get_line_col(src, pos) << " in source:\n";
+            std::string snippet = peak_source(src, pos);
+            string_replace_all(snippet, "\n", "\n" + indent(lvl));
+            oss << indent(lvl) << snippet << "\n";
+        } else {
+            for (auto & [label, children_vec] : children) {
+                oss << indent(lvl) << label << ":\n";
+                lvl++;
+                if (children_vec.empty()) {
+                    oss << indent(lvl) << "<empty>\n\n";
+                } else {
+                    for (auto * child : children_vec) {
+                        if (!child) {
+                            continue;
+                        }
+                        child->visit(ctx);
+                    }
+                }
+                lvl--;
+            }
+        }
+        lvl--;
+    };
+
+    for (const auto & stmt : prog.body) {
+        stmt->visit(ctx);
+    }
+
+    return oss.str();
+}
+
 } // namespace jinja

common/jinja/runtime.h

@@ -47,12 +47,19 @@ const T * cast_stmt(const statement_ptr & ptr) {
 // not thread-safe
 void enable_debug(bool enable);
 
+// for visiting AST nodes
+// function signature: void(bool is_leaf, statement * node, pair of <label, children>)
+using visitor_pair = std::pair<std::string, std::vector<statement *>>;
+using visitor_fn = std::function<void(bool, statement *, std::vector<visitor_pair>)>;
+
 struct context {
     std::shared_ptr<std::string> src; // for debugging; use shared_ptr to avoid copying on scope creation
     std::time_t current_time; // for functions that need current time
 
     bool is_get_stats = false; // whether to collect stats
 
+    visitor_fn visitor;
+
     // src is optional, used for error reporting
     context(std::string src = "") : src(std::make_shared<std::string>(std::move(src))) {
         env = mk_val<value_object>();
@@ -99,6 +106,15 @@ private:
     value_object env;
 };
 
+// utils for visiting AST nodes
+static std::vector<statement *> stmts_to_ptr(const statements & stmts) {
+    std::vector<statement *> children;
+    for (const auto & stmt : stmts) {
+        children.push_back(stmt.get());
+    }
+    return children;
+}
+
 /**
  * Base class for all nodes in the AST.
  */
@@ -106,6 +122,7 @@ struct statement {
     size_t pos; // position in source, for debugging
     virtual ~statement() = default;
     virtual std::string type() const { return "Statement"; }
+    virtual void visit(context & ctx) { ctx.visitor(true, this, {}); }
 
     // execute_impl must be overridden by derived classes
     virtual value execute_impl(context &) { throw_exec_error(); }
@@ -166,6 +183,13 @@ struct if_statement : public statement {
 
     std::string type() const override { return "If"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"test", {test.get()}},
+            {"body", stmts_to_ptr(body)},
+            {"alternate", stmts_to_ptr(alternate)}
+        });
+    }
 };
 
 struct identifier;
@@ -190,6 +214,14 @@ struct for_statement : public statement {
 
     std::string type() const override { return "For"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"loopvar", {loopvar.get()}},
+            {"iterable", {iterable.get()}},
+            {"body", stmts_to_ptr(body)},
+            {"default_block", stmts_to_ptr(default_block)}
+        });
+    }
 };
 
 struct break_statement : public statement {
@@ -241,6 +273,13 @@ struct set_statement : public statement {
 
     std::string type() const override { return "Set"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"assignee", {assignee.get()}},
+            {"value", {val.get()}},
+            {"body", stmts_to_ptr(body)}
+        });
+    }
 };
 
 struct macro_statement : public statement {
@@ -256,6 +295,13 @@ struct macro_statement : public statement {
 
     std::string type() const override { return "Macro"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"name", {name.get()}},
+            {"args", stmts_to_ptr(args)},
+            {"body", stmts_to_ptr(body)}
+        });
+    }
 };
 
 struct comment_statement : public statement {
@@ -289,6 +335,12 @@ struct member_expression : public expression {
     }
     std::string type() const override { return "MemberExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"object", {object.get()}},
+            {"property", {property.get()}}
+        });
+    }
 };
 
 struct call_expression : public expression {
@@ -302,6 +354,12 @@ struct call_expression : public expression {
     }
     std::string type() const override { return "CallExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"callee", {callee.get()}},
+            {"args", stmts_to_ptr(args)}
+        });
+    }
 };
 
 /**
@@ -405,6 +463,12 @@ struct binary_expression : public expression {
     }
     std::string type() const override { return "BinaryExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"left", {left.get()}},
+            {"right", {right.get()}}
+        });
+    }
 };
 
 /**
@@ -431,6 +495,12 @@ struct filter_expression : public expression {
 
     std::string type() const override { return "FilterExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"operand", {operand.get()}},
+            {"filter", {filter.get()}}
+        });
+    }
 };
 
 struct filter_statement : public statement {
@@ -443,6 +513,12 @@ struct filter_statement : public statement {
     }
     std::string type() const override { return "FilterStatement"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"filter", {filter.get()}},
+            {"body", stmts_to_ptr(body)}
+        });
+    }
 };
 
 /**
@@ -468,6 +544,12 @@ struct select_expression : public expression {
         }
         return lhs->execute_impl(ctx);
     }
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"lhs", {lhs.get()}},
+            {"test", {test.get()}}
+        });
+    }
 };
 
 /**
@@ -486,6 +568,12 @@ struct test_expression : public expression {
     }
     std::string type() const override { return "TestExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"operand", {operand.get()}},
+            {"test", {test.get()}}
+        });
+    }
 };
 
 /**
@@ -501,6 +589,11 @@ struct unary_expression : public expression {
     }
     std::string type() const override { return "UnaryExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"argument", {argument.get()}}
+        });
+    }
 };
 
 struct slice_expression : public expression {
@@ -518,6 +611,13 @@ struct slice_expression : public expression {
     [[noreturn]] value execute_impl(context &) override {
         throw std::runtime_error("must be handled by MemberExpression");
     }
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"start_expr", {start_expr.get()}},
+            {"stop_expr", {stop_expr.get()}},
+            {"step_expr", {step_expr.get()}}
+        });
+    }
 };
 
 struct keyword_argument_expression : public expression {
@@ -531,6 +631,12 @@ struct keyword_argument_expression : public expression {
     }
     std::string type() const override { return "KeywordArgumentExpression"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"key", {key.get()}},
+            {"val", {val.get()}}
+        });
+    }
 };
 
 struct spread_expression : public expression {
@@ -539,6 +645,11 @@ struct spread_expression : public expression {
         chk_type<expression>(this->argument);
     }
     std::string type() const override { return "SpreadExpression"; }
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"argument", {argument.get()}}
+        });
+    }
 };
 
 struct call_statement : public statement {
@@ -553,6 +664,13 @@ struct call_statement : public statement {
     }
     std::string type() const override { return "CallStatement"; }
     value execute_impl(context & ctx) override;
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"call", {call.get()}},
+            {"caller_args", stmts_to_ptr(caller_args)},
+            {"body", stmts_to_ptr(body)}
+        });
+    }
 };
 
 struct ternary_expression : public expression {
@@ -575,6 +693,13 @@ struct ternary_expression : public expression {
             return false_expr->execute(ctx);
         }
     }
+    void visit(context & ctx) override {
+        ctx.visitor(false, this, {
+            {"condition", {condition.get()}},
+            {"true_expr", {true_expr.get()}},
+            {"false_expr", {false_expr.get()}}
+        });
+    }
 };
 
 struct raised_exception : public std::exception {
@@ -648,6 +773,8 @@ struct runtime {
         }
         return parts;
     }
+
+    static std::string debug_dump_program(const program & prog, const std::string & src);
 };
 
 } // namespace jinja

common/jinja/value.cpp

@@ -12,6 +12,9 @@
 #include <optional>
 #include <algorithm>
 
+#ifdef FILENAME
+#undef FILENAME
+#endif
 #define FILENAME "jinja-value"
 
 namespace jinja {
@@ -1108,6 +1111,50 @@ const func_builtins & value_array_t::get_builtins() const {
             std::reverse(arr.begin(), arr.end());
             return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(arr)) : mk_val<value_array>(std::move(arr));
         }},
+        {"min", [](const func_args & args) -> value {
+            args.ensure_count(1, 4);
+            args.ensure_vals<value_array>();
+            value val_case    = args.get_kwarg_or_pos("case_sensitive", 1);
+            value attribute   = args.get_kwarg_or_pos("attribute",      2);
+            if (!attribute->is_undefined()) {
+                throw not_implemented_exception("min: attribute not implemented");
+            }
+            // FIXME: min is currently always case sensitive
+            (void) val_case;
+            const auto & arr = args.get_pos(0)->as_array();
+            if (arr.empty()) {
+                return mk_val<value_undefined>();
+            }
+            value result = arr[0];
+            for (size_t i = 1; i < arr.size(); ++i) {
+                if (value_compare(arr[i], result, value_compare_op::lt)) {
+                    result = arr[i];
+                }
+            }
+            return result;
+        }},
+        {"max", [](const func_args & args) -> value {
+            args.ensure_count(1, 4);
+            args.ensure_vals<value_array>();
+            value val_case    = args.get_kwarg_or_pos("case_sensitive", 1);
+            value attribute   = args.get_kwarg_or_pos("attribute",      2);
+            if (!attribute->is_undefined()) {
+                throw not_implemented_exception("max: attribute not implemented");
+            }
+            // FIXME: max is currently always case sensitive
+            (void) val_case;
+            const auto & arr = args.get_pos(0)->as_array();
+            if (arr.empty()) {
+                return mk_val<value_undefined>();
+            }
+            value result = arr[0];
+            for (size_t i = 1; i < arr.size(); ++i) {
+                if (value_compare(arr[i], result, value_compare_op::gt)) {
+                    result = arr[i];
+                }
+            }
+            return result;
+        }},
         {"unique", array_unique_not_implemented},
     };
     return builtins;

common/preset.cpp

@@ -7,6 +7,7 @@
 #include <fstream>
 #include <sstream>
 #include <filesystem>
+#include <regex>
 
 static std::string rm_leading_dashes(const std::string & str) {
     size_t pos = 0;
@@ -16,6 +17,23 @@ static std::string rm_leading_dashes(const std::string & str) {
     return str.substr(pos);
 }
 
+static std::string canonical_tag(const std::string & tag) {
+    static const std::regex re_tag("[-.]([A-Z0-9_]+)$", std::regex::icase);
+    std::smatch m;
+    if (std::regex_search(tag, m, re_tag)) {
+        std::string canon = m[1].str();
+        for (char & c : canon) {
+            c = (char) std::toupper((unsigned char) c);
+        }
+        return canon;
+    }
+    std::string upper = tag;
+    for (char & c : upper) {
+        c = (char) std::toupper((unsigned char) c);
+    }
+    return upper;
+}
+
 std::vector<std::string> common_preset::to_args(const std::string & bin_path) const {
     std::vector<std::string> args;
 
@@ -270,11 +288,18 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co
 
     for (auto section : ini_data) {
         common_preset preset;
-        if (section.first.empty()) {
-            preset.name = COMMON_PRESET_DEFAULT_NAME;
-        } else {
-            preset.name = section.first;
+        std::string section_name = section.first.empty() ? std::string(COMMON_PRESET_DEFAULT_NAME) : section.first;
+        if (section_name != "*" && section_name != COMMON_PRESET_DEFAULT_NAME) {
+            auto colon_idx = section_name.rfind(':');
+            if (colon_idx != std::string::npos) {
+                std::string tag       = section_name.substr(colon_idx + 1);
+                std::string canon_tag = canonical_tag(tag);
+                if (canon_tag != tag) {
+                    section_name = section_name.substr(0, colon_idx + 1) + canon_tag;
+                }
+            }
         }
+        preset.name = section_name;
         LOG_DBG("loading preset: %s\n", preset.name.c_str());
         for (const auto & [key, value] : section.second) {
             if (key == "version") {

common/reasoning-budget.cpp

@@ -65,12 +65,12 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
             if (ctx->start_matcher.advance(token)) {
                 ctx->state = REASONING_BUDGET_COUNTING;
                 ctx->remaining = ctx->budget;
-                LOG_INF("reasoning-budget: activated, budget=%d tokens\n", ctx->budget);
+                COM_TRC("activated, budget=%d tokens\n", ctx->budget);
 
                 if (ctx->remaining <= 0) {
                     ctx->state = REASONING_BUDGET_FORCING;
                     ctx->force_pos = 0;
-                    LOG_INF("reasoning-budget: budget=0, forcing immediately\n");
+                    COM_TRC("%s", "budget=0, forcing immediately\n");
                 }
             }
             break;
@@ -80,7 +80,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
         {
             if (ctx->end_matcher.advance(token)) {
                 ctx->state = REASONING_BUDGET_DONE;
-                LOG_INF("reasoning-budget: deactivated (natural end)\n");
+                COM_TRC("%s", "deactivated (natural end)\n");
                 break;
             }
 
@@ -95,7 +95,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
                     ctx->state = REASONING_BUDGET_FORCING;
                     ctx->force_pos = 0;
                     ctx->end_matcher.reset();
-                    LOG_INF("reasoning-budget: UTF-8 complete, now forcing end sequence\n");
+                    COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n");
                 }
             } else if (ctx->state == REASONING_BUDGET_COUNTING) {
                 ctx->remaining--;
@@ -104,11 +104,11 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
                         ctx->state = REASONING_BUDGET_FORCING;
                         ctx->force_pos = 0;
                         ctx->end_matcher.reset();
-                        LOG_INF("reasoning-budget: budget exhausted, forcing end sequence\n");
+                        COM_TRC("%s", "budget exhausted, forcing end sequence\n");
                     } else {
                         ctx->state = REASONING_BUDGET_WAITING_UTF8;
                         ctx->end_matcher.reset();
-                        LOG_INF("reasoning-budget: budget exhausted, waiting for UTF-8 completion\n");
+                        COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n");
                     }
                 }
             }
@@ -118,7 +118,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
             ctx->force_pos++;
             if (ctx->force_pos >= ctx->forced_tokens.size()) {
                 ctx->state = REASONING_BUDGET_DONE;
-                LOG_INF("reasoning-budget: forced sequence complete, done\n");
+                COM_TRC("%s", "forced sequence complete, done\n");
             }
             break;
         case REASONING_BUDGET_DONE:
@@ -128,12 +128,12 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
                 ctx->state = REASONING_BUDGET_COUNTING;
                 ctx->remaining = ctx->budget;
                 ctx->end_matcher.reset();
-                LOG_INF("reasoning-budget: re-activated on new start tag, budget=%d tokens\n", ctx->budget);
+                COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
 
                 if (ctx->remaining <= 0) {
                     ctx->state = REASONING_BUDGET_FORCING;
                     ctx->force_pos = 0;
-                    LOG_INF("reasoning-budget: budget=0, forcing immediately\n");
+                    COM_TRC("%s", "budget=0, forcing immediately\n");
                 }
             }
             break;
@@ -264,7 +264,7 @@ bool common_reasoning_budget_force(struct llama_sampler * smpl) {
     ctx->state = REASONING_BUDGET_FORCING;
     ctx->force_pos = 0;
     ctx->end_matcher.reset();
-    LOG_INF("reasoning-budget: forced into forcing state (manual transition)\n");
+    COM_TRC("%s", "forced into forcing state (manual transition)\n");
 
     return true;
 }

common/regex-partial.cpp

@@ -1,204 +0,0 @@
-#include "regex-partial.h"
-#include "common.h"
-#include <functional>
-#include <optional>
-
-common_regex::common_regex(const std::string & pattern) :
-    pattern(pattern),
-    rx(pattern),
-    rx_reversed_partial(regex_to_reversed_partial_regex(pattern)) {}
-
-common_regex_match common_regex::search(const std::string & input, size_t pos, bool as_match) const {
-    std::smatch match;
-    if (pos > input.size()) {
-        throw std::runtime_error("Position out of bounds");
-    }
-    auto start = input.begin() + pos;
-    auto found = as_match
-        ? std::regex_match(start, input.end(), match, rx)
-        : std::regex_search(start, input.end(), match, rx);
-    if (found) {
-        common_regex_match res;
-        res.type = COMMON_REGEX_MATCH_TYPE_FULL;
-        for (size_t i = 0; i < match.size(); ++i) {
-            auto begin = pos + match.position(i);
-            res.groups.emplace_back(begin, begin + match.length(i));
-        }
-        return res;
-    }
-    std::match_results<std::string::const_reverse_iterator> srmatch;
-    if (std::regex_search(input.rbegin(), input.rend() - pos, srmatch, rx_reversed_partial, std::regex_constants::match_continuous)) {
-        auto group = srmatch[1].str();
-        if (group.length() != 0) {
-            auto it = srmatch[1].second.base();
-            // auto position = static_cast<size_t>(std::distance(input.begin(), it));
-            if ((!as_match) || it == input.begin()) {
-                common_regex_match res;
-                res.type = COMMON_REGEX_MATCH_TYPE_PARTIAL;
-                const size_t begin = std::distance(input.begin(), it);
-                const size_t end = input.size();
-                if (begin == std::string::npos || end == std::string::npos || begin > end) {
-                    throw std::runtime_error("Invalid range");
-                }
-                res.groups.push_back({begin, end});
-                return res;
-            }
-        }
-    }
-    return {};
-}
-
-/*
-  Transforms a regex pattern to a partial match pattern that operates on a reversed input string to find partial final matches of the original pattern.
-
-  Ideally we'd like to use boost::match_partial (https://beta.boost.org/doc/libs/1_59_0/libs/regex/doc/html/boost_regex/partial_matches.html)
-  to see if a string ends with a partial regex match, but but it's not in std::regex yet.
-  Instead, we'll the regex into a partial match regex operating as a full match on the reverse iterators of the input.
-
-  - /abcd/ -> ^(dcba|cba|ba|a) -> ^((?:(?:(?:(?:d)?c)?b)?a)
-  - /a|b/ -> ^(a|b)
-  - /a*?/ -> error, could match ""
-  - /a*b/ -> ^((?:b)?a*+) (final repetitions become eager)
-  - /.*?ab/ -> ^((?:b)?a) (omit .*)
-  - /a.*?b/ -> ^((?:b)?.*?a) (keep reluctant matches)
-  - /a(bc)d/ -> ^((?:(?:d)?(?:(?:c)?b))?a)
-  - /a(bc|de)/ -> ^((?:(?:(?:e)?d)?|(?:(?:c)?b)?)?a)
-  - /ab{2,4}c/ -> ^cbbb?b?a -> ^((?:(?:(?:(?:(?:c)?b)?b)?b?)?b?)?a)
-
-  The regex will match a reversed string fully, and the end of the first (And only) capturing group will indicate the reversed start of the original partial pattern.
-  All other groups are turned into non-capturing groups, and reluctant quantifiers are ignored.
-*/
-std::string regex_to_reversed_partial_regex(const std::string & pattern) {
-    auto it = pattern.begin();
-    const auto end = pattern.end();
-
-    std::function<std::string()> process = [&]() {
-        std::vector<std::vector<std::string>> alternatives(1);
-        std::vector<std::string> * sequence = &alternatives.back();
-
-        while (it != end) {
-            if (*it == '[') {
-                auto start = it;
-                ++it;
-                while (it != end) {
-                    if ((*it == '\\') && (++it != end)) {
-                        ++it;
-                    } else if ((it != end) && (*it == ']')) {
-                        break;
-                    } else {
-                        ++it;
-                    }
-                }
-                if (it == end) {
-                    throw std::runtime_error("Unmatched '[' in pattern");
-                }
-                ++it;
-                sequence->push_back(std::string(start, it));
-            } else if (*it == '*' || *it == '?' || *it == '+') {
-                if (sequence->empty()) {
-                    throw std::runtime_error("Quantifier without preceding element");
-                }
-                sequence->back() += *it;
-                auto is_star = *it == '*';
-                ++it;
-                if (is_star) {
-                    if (it != end && *it == '?') {
-                        ++it;
-                    }
-                }
-            } else if (*it == '{') {
-                if (sequence->empty()) {
-                    throw std::runtime_error("Repetition without preceding element");
-                }
-                ++it;
-                auto start = it;
-                while (it != end && *it != '}') {
-                    ++it;
-                }
-                if (it == end) {
-                    throw std::runtime_error("Unmatched '{' in pattern");
-                }
-                auto parts = string_split(std::string(start, it), ",");
-                ++it;
-                if (parts.size() > 2) {
-                    throw std::runtime_error("Invalid repetition range in pattern");
-                }
-
-                auto parseOptInt = [&](const std::string & s, const std::optional<int> & def = std::nullopt) -> std::optional<int> {
-                    if (s.empty()) {
-                        return def;
-                    }
-                    return std::stoi(s);
-                };
-                auto min = parseOptInt(parts[0], 0);
-                auto max = parts.size() == 1 ? min : parseOptInt(parts[1]);
-                if (min && max && *max < *min) {
-                    throw std::runtime_error("Invalid repetition range in pattern");
-                }
-                // Brutal but... let's repeat at least min times, then ? for the delta between min & max (or * for unbounded)
-                auto part = sequence->back();
-                sequence->pop_back();
-                for (int i = 0; i < *min; i++) {
-                    sequence->push_back(part);
-                }
-                if (max) {
-                    for (int i = *min; i < *max; i++) {
-                        sequence->push_back(part + "?");
-                    }
-                } else {
-                    sequence->push_back(part + "*");
-                }
-            } else if (*it == '(') {
-                ++it;
-                if (it != end && *it == '?' && (it + 1 != end) && *(it + 1) == ':') {
-                    it += 2;
-                }
-                auto sub = process();
-                if (*it != ')') {
-                    throw std::runtime_error("Unmatched '(' in pattern");
-                }
-                ++it;
-                auto & part = sequence->emplace_back("(?:");
-                part += sub;
-                part += ")";
-            } else if (*it == ')') {
-                break;
-            } else if (*it == '|') {
-                ++it;
-                alternatives.emplace_back();
-                sequence = &alternatives.back();
-            } else if (*it == '\\' && (++it != end)) {
-                auto str = std::string("\\") + *it;
-                sequence->push_back(str);
-                ++it;
-            } else if (it != end) {
-                sequence->push_back(std::string(1, *it));
-                ++it;
-            }
-        }
-
-        // /abcd/ -> ^(dcba|cba|ba|a) -> ^((?:(?:(?:d)?c)?b)?a)
-        // if n(=4) parts, opening n-1(=3) non-capturing groups after the 1 capturing group
-        // We'll do the outermost capturing group and final .* in the enclosing function.
-        std::vector<std::string> res_alts;
-        for (const auto & parts : alternatives) {
-            auto & res = res_alts.emplace_back();
-            for (size_t i = 0; i < parts.size() - 1; i++) {
-                res += "(?:";
-            }
-            for (auto it = parts.rbegin(); it != parts.rend(); ++it) {
-                res += *it;
-                if (it != parts.rend() - 1) {
-                    res += ")?";
-                }
-            }
-        }
-        return string_join(res_alts, "|");
-    };
-    auto res = process();
-    if (it != end) {
-        throw std::runtime_error("Unmatched '(' in pattern");
-    }
-
-    return "^(" + res + ")";
-}

common/regex-partial.h

@@ -1,56 +0,0 @@
-#pragma once
-
-#include <regex>
-#include <string>
-
-enum common_regex_match_type {
-    COMMON_REGEX_MATCH_TYPE_NONE,
-    COMMON_REGEX_MATCH_TYPE_PARTIAL,
-    COMMON_REGEX_MATCH_TYPE_FULL,
-};
-
-struct common_string_range {
-    size_t begin;
-    size_t end;
-    common_string_range(size_t begin, size_t end) : begin(begin), end(end) {
-        if (begin > end) {
-            throw std::runtime_error("Invalid range");
-        }
-    }
-    // prevent default ctor
-    common_string_range() = delete;
-    bool empty() const {
-        return begin == end;
-    }
-    bool operator==(const common_string_range & other) const {
-        return begin == other.begin && end == other.end;
-    }
-};
-
-struct common_regex_match {
-    common_regex_match_type type = COMMON_REGEX_MATCH_TYPE_NONE;
-    std::vector<common_string_range> groups;
-
-    bool operator==(const common_regex_match & other) const {
-        return type == other.type && groups == other.groups;
-    }
-    bool operator!=(const common_regex_match & other) const {
-        return !(*this == other);
-    }
-};
-
-class common_regex {
-    std::string pattern;
-    std::regex rx;
-    std::regex rx_reversed_partial;
-
-  public:
-    explicit common_regex(const std::string & pattern);
-
-    common_regex_match search(const std::string & input, size_t pos, bool as_match = false) const;
-
-    const std::string & str() const { return pattern; }
-};
-
-// For testing only (pretty print of failures).
-std::string regex_to_reversed_partial_regex(const std::string & pattern);

common/speculative.cpp

@@ -18,6 +18,13 @@
 #include <map>
 #include <cinttypes>
 
+#define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define SPC_INF(fmt, ...) LOG_INF("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define SPC_WRN(fmt, ...) LOG_WRN("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define SPC_ERR(fmt, ...) LOG_ERR("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+#define SPC_CNT(fmt, ...) LOG_CNT(""              fmt,               __VA_ARGS__)
+
 #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  128
 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5
 
@@ -26,6 +33,7 @@ const std::map<std::string, common_speculative_type> common_speculative_type_fro
     {"draft-simple",  COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE},
     {"draft-eagle3",  COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
     {"draft-mtp",     COMMON_SPECULATIVE_TYPE_DRAFT_MTP},
+    {"draft-dflash",  COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},
     {"ngram-simple",  COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
     {"ngram-map-k",   COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
     {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
@@ -60,21 +68,20 @@ static bool common_speculative_are_compatible(
     const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft);
 
     const auto vocab_type_tgt = llama_vocab_type(vocab_tgt);
-    LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt);
+    SPC_DBG("vocab_type tgt: %d\n", vocab_type_tgt);
 
     const auto vocab_type_dft = llama_vocab_type(vocab_dft);
-    LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft);
+    SPC_DBG("vocab_type dft: %d\n", vocab_type_dft);
 
     if (vocab_type_tgt != vocab_type_dft) {
-        LOG_WRN("%s: draft model vocab type must match target model to use speculation but "
-                "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt);
+        SPC_WRN("draft model vocab type must match target model to use speculation but "
+                "vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt);
         return false;
     }
 
     if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||
         (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) {
-        LOG_WRN("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",
-                __func__,
+        SPC_WRN("draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",
                 llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft),
                 llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft));
         return false;
@@ -82,8 +89,7 @@ static bool common_speculative_are_compatible(
 
     if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||
         (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) {
-        LOG_WRN("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",
-                __func__,
+        SPC_WRN("draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n",
                 llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft),
                 llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft));
         return false;
@@ -97,8 +103,8 @@ static bool common_speculative_are_compatible(
             : n_vocab_dft - n_vocab_tgt;
 
         if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {
-            LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__);
-            LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",
+            SPC_DBG("draft model vocab must closely match target model to use speculation but "
+                    "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",
                     n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);
             return false;
         }
@@ -108,8 +114,8 @@ static bool common_speculative_are_compatible(
             const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);
 
             if (std::strcmp(token_text_tgt, token_text_dft) != 0) {
-                LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__);
-                LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i,
+                SPC_DBG("draft model vocab must match target model to use speculation but "
+                        "token %d content differs - target '%s', draft '%s'\n", i,
                         common_token_to_piece(vocab_tgt, i).c_str(),
                         common_token_to_piece(vocab_dft, i).c_str());
                 return false;
@@ -186,9 +192,9 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
         auto * ctx_dft = this->params.ctx_dft;
         auto * ctx_tgt = this->params.ctx_tgt;
 
-        LOG_INF("%s: adding speculative implementation 'draft-simple'\n", __func__);
-        LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
-        LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__,
+        SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n");
+        SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min);
+        SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",
                 this->params.n_gpu_layers,
                 ggml_type_name(this->params.cache_type_k),
                 ggml_type_name(this->params.cache_type_v),
@@ -228,16 +234,16 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
         }
 
         const bool vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft));
-        LOG_DBG("%s: vocab_cmpt = %d\n", __func__, vocab_cmpt);
+        SPC_DBG("vocab_cmpt = %d\n", vocab_cmpt);
 
         if (!vocab_cmpt) {
-            LOG_ERR("%s: the target and draft vocabs are not compatible\n", __func__);
+            SPC_ERR("%s", "the target and draft vocabs are not compatible\n");
 
             throw std::runtime_error("draft model vocab type must match target model to use speculation");
         }
 
         if (n_seq != llama_n_seq_max(ctx_dft)) {
-            LOG_ERR("%s: n_seq mismatch: %d != %d\n", __func__, n_seq, llama_n_seq_max(ctx_dft));
+            SPC_ERR("n_seq mismatch: %d != %d\n", n_seq, llama_n_seq_max(ctx_dft));
 
             throw std::runtime_error("the draft model number of sequences is incompatible with the speculative n_seq");
         }
@@ -257,7 +263,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
         const int ret = llama_decode(ctx_dft, batch);
 
         if (ret != 0) {
-            LOG_ERR("%s: failed to decode draft batch, ret = %d\n", __func__, ret);
+            SPC_ERR("failed to decode draft batch, ret = %d\n", ret);
 
             return false;
         }
@@ -290,7 +296,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
 
         int ret = llama_decode(ctx_dft, batch);
         if (ret != 0) {
-            LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);
+            SPC_ERR("llama_decode returned %d\n", ret);
             return;
         }
 
@@ -314,7 +320,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
                 const auto * cur_p = common_sampler_get_candidates(smpl, true);
 
                 for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
-                    LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+                    SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
                             seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
                             common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
                 }
@@ -354,7 +360,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
             // evaluate the drafted tokens on the draft model
             ret = llama_decode(ctx_dft, batch);
             if (ret != 0) {
-                LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                 break;
             }
 
@@ -449,8 +455,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
         : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq)
         , params(params.draft)
     {
-        LOG_INF("%s: adding speculative implementation 'draft-eagle3'\n", __func__);
-        LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", __func__, params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling);
+        SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n");
+        SPC_TRC("- n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling);
 
         auto * ctx_tgt = this->params.ctx_tgt;
         auto * ctx_dft = this->params.ctx_dft;
@@ -493,7 +499,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
                 llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
 
                 if (!llama_set_sampler(ctx_dft, seq_id, chain)) {
-                    LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id);
+                    SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id);
                     llama_sampler_free(chain);
                     chain = nullptr;
                 }
@@ -548,9 +554,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
         auto * ctx_dft = this->params.ctx_dft;
         const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);
         if (pos_max < N - 2) {
-            LOG_WRN("%s: ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. "
+            SPC_WRN("ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. "
                     "Drafts may degrade.\n",
-                    __func__, (int) pos_max, N - 2);
+                    (int) pos_max, N - 2);
         }
     }
 
@@ -621,8 +627,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
             };
             const int32_t rc = llama_encode(ctx_dft, enc_batch);
             if (rc != 0) {
-                LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
-                        __func__, rc, (int) n_chunk, (int) i);
+                SPC_ERR("llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                        rc, (int) n_chunk, (int) i);
                 return false;
             }
 
@@ -692,8 +698,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
         if (batch.n_tokens > 0) {
             const int32_t rc = llama_decode(ctx_dft, batch);
             if (rc != 0) {
-                LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",
-                        __func__, rc, (int) batch.n_tokens, (int) batch_in.pos[0]);
+                SPC_ERR("llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",
+                        rc, (int) batch.n_tokens, (int) batch_in.pos[0]);
                 return false;
             }
         }
@@ -744,7 +750,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
 
         int ret = llama_decode(ctx_dft, batch);
         if (ret != 0) {
-            LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);
+            SPC_ERR("llama_decode returned %d\n", ret);
             return;
         }
 
@@ -770,7 +776,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
                 const auto * cur_p = common_sampler_get_candidates(smpl, true);
 
                 for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
-                    LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+                    SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
                             seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
                             common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
                 }
@@ -809,7 +815,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
 
             ret = llama_decode(ctx_dft, batch);
             if (ret != 0) {
-                LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                 break;
             }
 
@@ -893,6 +899,305 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
     }
 };
 
+// DFlash: block-diffusion drafting with a draft-side KV cache injection
+struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+    common_params_speculative_draft params;
+
+    llama_batch batch;        // noise tokens
+    llama_batch batch_inject; // target features for KV cache injection
+
+    std::vector<common_sampler_ptr> smpls;
+
+    int32_t n_embd_dec = 0;  // draft hidden size
+    int32_t n_embd_enc = 0;  // target_layer_ids_n * target_hidden_size
+    int32_t n_embd_tgt = 0;  // target model hidden size
+
+    int32_t     block_size    = 0;
+    llama_token mask_token_id = 0;
+
+    const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
+    uint32_t        target_layer_ids_n = 0;
+
+    // scratch buffer for concatenated target features [n_tokens, n_embd_enc]
+    std::vector<float> features_buf;
+
+    common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq)
+        , params(params.draft)
+    {
+        auto * ctx_tgt = this->params.ctx_tgt;
+        auto * ctx_dft = this->params.ctx_dft;
+        GGML_ASSERT(ctx_tgt && ctx_dft && "DFlash requires ctx_tgt and ctx_dft to be set");
+
+        const llama_model * model_dft = llama_get_model(ctx_dft);
+        const llama_model * model_tgt = llama_get_model(ctx_tgt);
+
+        target_layer_ids   = llama_model_target_layer_ids  (model_dft);
+        target_layer_ids_n = llama_model_target_layer_ids_n(model_dft);
+        GGML_ASSERT(target_layer_ids_n > 0 && "DFlash model has no target_layer_ids");
+
+        n_embd_tgt    = llama_model_n_embd(model_tgt);
+        n_embd_dec    = llama_model_n_embd(model_dft);
+        n_embd_enc    = (int32_t) target_layer_ids_n * n_embd_tgt;
+
+        // read the trained block size from the dflash.block_size metadata key
+        block_size = 16;
+        {
+            char buf[32] = {};
+            if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) {
+                block_size = std::atoi(buf);
+            }
+        }
+        mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
+
+        LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__);
+        LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
+        LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n);
+
+        // DFlash input is [id_last, <mask> * (block_size-1)], so it can draft at most block_size-1 tokens per step
+        if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) {
+            LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n",
+                    __func__, this->params.n_max, this->params.n_min, block_size, block_size - 1);
+            this->params.n_max = std::min(this->params.n_max, block_size - 1);
+            this->params.n_min = std::min(this->params.n_min, block_size - 1);
+        }
+
+        batch        = llama_batch_init(llama_n_batch(ctx_dft), 0,          n_seq);
+        batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq);
+
+        smpls.resize(n_seq);
+        for (auto & s : smpls) {
+            common_params_sampling sparams;
+            sparams.no_perf  = false;
+            sparams.top_k    = 10;
+            sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };
+            s.reset(common_sampler_init(model_dft, sparams));
+        }
+
+        // turn on extraction of the target layers' input embeddings
+        for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
+            llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
+        }
+
+        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
+        llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
+    }
+
+    ~common_speculative_impl_draft_dflash() override {
+        llama_batch_free(batch);
+        llama_batch_free(batch_inject);
+    }
+
+    void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return;
+        }
+
+        const int32_t N = (int32_t) prompt.size();
+        if (N <= 0) {
+            return;
+        }
+
+        const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(params.ctx_dft), seq_id);
+        if (pos_max < N - 1) {
+            LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - process() did not run on every prefill ubatch. "
+                    "Drafts may degrade.\n",
+                    __func__, (int) pos_max, N - 1);
+        }
+    }
+
+    bool process(const llama_batch & batch_in) override {
+        if (batch_in.n_tokens <= 0) {
+            return true;
+        }
+
+        if (batch_in.token == nullptr || batch_in.embd != nullptr) {
+            return true;
+        }
+
+        const int32_t n_tokens = batch_in.n_tokens;
+
+        // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch)
+        std::vector<int32_t> i_batch_beg(n_seq, -1);
+        std::vector<int32_t> i_batch_end(n_seq, -1);
+        for (int32_t k = 0; k < n_tokens; ++k) {
+            GGML_ASSERT(batch_in.n_seq_id[k] == 1);
+            const llama_seq_id seq_id = batch_in.seq_id[k][0];
+            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+                continue;
+            }
+            i_batch_end[seq_id] = k;
+            if (i_batch_beg[seq_id] < 0) {
+                i_batch_beg[seq_id] = k;
+            }
+        }
+
+        auto * ctx_tgt = this->params.ctx_tgt;
+        auto * ctx_dft = this->params.ctx_dft;
+
+        const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft);
+
+        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+            if (i_batch_beg[seq_id] < 0) {
+                continue;
+            }
+            const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
+
+            for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
+                const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
+
+                // gather this chunk's target features, interleaved by extract layer
+                features_buf.resize((size_t) n_chunk * n_embd_enc);
+                for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
+                    const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
+                    if (!layer) {
+                        GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
+                    }
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+                        const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
+                        std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
+                    }
+                }
+
+                // fuse extracted features through DFlash encoder
+                llama_batch enc_batch = {
+                    /*.n_tokens =*/ n_chunk,
+                    /*.token    =*/ nullptr,
+                    /*.embd     =*/ features_buf.data(),
+                    /*.pos      =*/ nullptr,
+                    /*.n_seq_id =*/ nullptr,
+                    /*.seq_id   =*/ nullptr,
+                    /*.logits   =*/ nullptr,
+                };
+
+                int32_t rc = llama_encode(ctx_dft, enc_batch);
+                if (rc != 0) {
+                    LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                            __func__, rc, (int) n_chunk, (int) offset);
+                    return false;
+                }
+
+                const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
+                GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
+
+                // inject the DFlash decoder K/V cache at the tokens' target positions
+                batch_inject.n_tokens = n_chunk;
+                std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
+
+                for (int32_t i = 0; i < n_chunk; ++i) {
+                    batch_inject.pos[i]       = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                    batch_inject.n_seq_id[i]  = 1;
+                    batch_inject.seq_id[i][0] = seq_id;
+                    batch_inject.logits[i]    = false;
+                }
+                rc = llama_decode(ctx_dft, batch_inject);
+                if (rc != 0) {
+                    LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                            __func__, rc, (int) n_chunk, (int) offset);
+                    return false;
+                }
+            }
+        }
+
+        return true;
+    }
+
+    void draft(common_speculative_draft_params_vec & dparams) override {
+        auto & ctx_dft = params.ctx_dft;
+
+        common_batch_clear(batch);
+
+        // build one batch holding every drafting sequence's noise block into a single decode)
+        // record where each block starts and its size
+        std::vector<int32_t> i_block_beg(n_seq, -1);
+        std::vector<int32_t> n_block    (n_seq,  0);
+
+        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+            auto & dp = dparams[seq_id];
+            if (!dp.drafting) {
+                continue;
+            }
+
+            common_sampler_reset(smpls[seq_id].get());
+
+            const int32_t n = (int32_t) dp.n_past;
+
+            int32_t n_draft = params.n_max;
+            if (dp.n_max > 0) {
+                n_draft = std::min(n_draft, dp.n_max);
+            }
+
+            const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * <mask>
+            i_block_beg[seq_id] = batch.n_tokens;
+            n_block    [seq_id] = n_block_tokens;
+            for (int32_t i = 0; i < n_block_tokens; ++i) {
+                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true);
+            }
+        }
+
+        if (batch.n_tokens == 0) {
+            return;
+        }
+
+        // decode all sequence's noise block in a single batch
+        int ret = llama_decode(ctx_dft, batch);
+        if (ret != 0) {
+            LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);
+            return;
+        }
+
+        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+            if (i_block_beg[seq_id] < 0) {
+                continue;
+            }
+            auto & dp = dparams[seq_id];
+
+            const int32_t beg            = i_block_beg[seq_id];
+            const int32_t n_block_tokens = n_block[seq_id];
+
+            auto * smpl = smpls[seq_id].get();
+
+            auto & result = *dp.result;
+
+            // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
+            for (int32_t i = 1; i < n_block_tokens; ++i) {
+                common_sampler_sample(smpl, ctx_dft, beg + i, true);
+
+                const auto * cur_p = common_sampler_get_candidates(smpl, true);
+
+                for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
+                    LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+                            seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
+                            common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
+                }
+
+                const llama_token id = cur_p->data[0].id;
+
+                if (cur_p->data[0].p < params.p_min) {
+                    break;
+                }
+
+                common_sampler_accept(smpl, id, true);
+
+                result.push_back(id);
+            }
+
+            if (result.size() < (size_t) params.n_min) {
+                result.clear();
+            }
+        }
+    }
+
+    void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
+        // noop
+    }
+
+    bool need_embd() const override {
+        return false;
+    }
+};
+
 struct common_speculative_impl_draft_mtp : public common_speculative_impl {
     common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft)
 
@@ -942,9 +1247,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                 "MTP input row width must match the target h_nextn width");
         n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft)));
 
-        LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__);
-        LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling);
-        LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__,
+        SPC_TRC("%s", "adding speculative implementation 'draft-mtp'\n");
+        SPC_TRC("- n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling);
+        SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",
                 this->params.n_gpu_layers,
                 ggml_type_name(this->params.cache_type_k),
                 ggml_type_name(this->params.cache_type_v),
@@ -975,7 +1280,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                 llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
 
                 if (!llama_set_sampler(ctx_dft, seq_id, chain)) {
-                    LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id);
+                    SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id);
                     llama_sampler_free(chain);
                     chain = nullptr;
                 }
@@ -1038,11 +1343,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);
 
         if (pos_max < N - 1 && !is_mem_shared) {
-            LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - "
+            SPC_WRN("ctx_dft pos_max=%d < N-1=%d - "
                     "process() hook may not have run on every prefill ubatch "
                     "(need_embd / logits=1 on every prompt position?). "
                     "Drafts may degrade.\n",
-                    __func__, (int) pos_max, N - 1);
+                    (int) pos_max, N - 1);
         }
     }
 
@@ -1128,8 +1433,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
                 const int32_t rc = llama_decode(ctx_dft, batch);
                 if (rc != 0) {
-                    LOG_ERR("%s: llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n",
-                            __func__, head, (int) rc, (int) batch_in.pos[0]);
+                    SPC_ERR("llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n",
+                            head, (int) rc, (int) batch_in.pos[0]);
                     ok = false;
                     break;
                 }
@@ -1217,7 +1522,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
             int ret = llama_decode(ctx_dft, batch);
             if (ret != 0) {
-                LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                 break;
             }
 
@@ -1239,7 +1544,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                 const auto * cur_p = common_sampler_get_candidates(smpl, true);
 
                 for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
-                    LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+                    SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
                             seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
                             common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
                 }
@@ -1353,8 +1658,8 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
         , params(params.ngram_simple)
         , config(config)
     {
-        LOG_INF("%s: adding speculative implementation 'ngram-simple'\n", __func__);
-        LOG_INF("%s: - size_n=%d, size_m=%d, min_hits=%d\n", __func__,
+        SPC_TRC("%s", "adding speculative implementation 'ngram-simple'\n");
+        SPC_TRC("- size_n=%d, size_m=%d, min_hits=%d\n",
                 this->params.size_n, this->params.size_m, this->params.min_hits);
     }
 
@@ -1403,8 +1708,8 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
             this->config.push_back(config);
         }
 
-        LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(this->type).c_str());
-        LOG_INF("%s: - size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n", __func__,
+        SPC_TRC("adding speculative implementation '%s'\n", common_speculative_type_to_str(this->type).c_str());
+        SPC_TRC("- size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n",
                 config.size_key, config.size_value, config.key_only, config.min_hits);
     }
 
@@ -1478,15 +1783,15 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
         , verbose(std::getenv("LLAMA_TRACE") != nullptr) {
         static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t));
 
-        LOG_INF("%s: adding speculative implementation 'ngram-mod'\n", __func__);
-        LOG_INF("%s: - n_match=%d, n_max=%d, n_min=%d\n", __func__,
+        SPC_TRC("%s", "adding speculative implementation 'ngram-mod'\n");
+        SPC_TRC("- n_match=%d, n_max=%d, n_min=%d\n",
                 this->params.n_match, this->params.n_max, this->params.n_min);
-        LOG_INF("%s: - mod size=%zu (%.3f MB)\n", __func__,
+        SPC_TRC("- mod size=%zu (%.3f MB)\n",
                 mod.size(), (float)(mod.size_bytes())/1024/1024);
 
         if (this->params.n_match < 16) {
-            LOG_WRN("%s: ngram_mod n_match=%d is too small - poor quality is possible, "
-                    "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, this->params.n_match);
+            SPC_WRN("ngram_mod n_match=%d is too small - poor quality is possible, "
+                    "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", this->params.n_match);
         }
 
         sinfos.resize(n_seq);
@@ -1510,11 +1815,11 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
         sinfo.i_last = prompt.size() - n;
 
         const double f = (double)mod.get_used() / (double)mod.size();
-        LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f);
+        SPC_TRC("ngram_mod occupancy = %zu/%zu (%.2f)\n", mod.get_used(), mod.size(), f);
 
         constexpr double f_thold = 0.25;
         if (f > f_thold) {
-            LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold);
+            SPC_WRN("ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", f, f_thold);
 
             mod.reset();
         }
@@ -1608,7 +1913,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
                 sinfo.n_low++;
                 if (sinfo.n_low >= 5) {
                     if (verbose) {
-                        LOG_WRN("%s: low acceptance streak (%d) - resetting ngram_mod\n", __func__, sinfo.n_low);
+                        SPC_TRC("low acceptance streak (%d) - resetting ngram_mod\n", sinfo.n_low);
                     }
 
                     mod.reset();
@@ -1658,8 +1963,8 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
         , save_dynamic(save_dynamic)
         , save_static(save_static)
     {
-        LOG_INF("%s: adding speculative implementation 'ngram-cache'\n", __func__);
-        LOG_INF("%s: - n_draft=%d, cache_static=%s, cache_dynamic=%s\n", __func__,
+        SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n");
+        SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n",
                 n_draft,
                 path_static.empty() ? "none" : path_static.c_str(),
                 path_dynamic.empty() ? "none" : path_dynamic.c_str());
@@ -1674,7 +1979,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
                     sinfo.ngram_cache_static = ngram_cache_static;
                 }
             } catch (...) {
-                LOG_ERR("failed to open static lookup cache: %s", path_static.c_str());
+                SPC_ERR("failed to open static lookup cache: %s", path_static.c_str());
                 GGML_ABORT("Couldn't read static lookup cache");
             }
         }
@@ -1687,7 +1992,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
                     sinfo.ngram_cache_dynamic = ngram_cache_dynamic;
                 }
             } catch (...) {
-                LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
+                SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
                 GGML_ABORT("Couldn't read dynamic lookup cache");
             }
         }
@@ -1836,6 +2141,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) {
         case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE:  return "draft-simple";
         case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3:  return "draft-eagle3";
         case COMMON_SPECULATIVE_TYPE_DRAFT_MTP:     return "draft-mtp";
+        case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:  return "draft-dflash";
         case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:  return "ngram-simple";
         case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:   return "ngram-map-k";
         case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
@@ -1888,6 +2194,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
             case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE:
             case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3:
             case COMMON_SPECULATIVE_TYPE_DRAFT_MTP:
+            case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:
                 n_max = std::max(n_max, std::max(0, spec->draft.n_max));
                 break;
             case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
@@ -1925,6 +2232,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         bool has_draft_simple = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE));
         bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr;
         bool has_draft_mtp    = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP))    && params.draft.ctx_dft != nullptr;
+        bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr;
 
 
 
@@ -1935,7 +2243,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         bool has_ngram_mod     = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD));
 
         // when adding a new type - update here the logic above
-        static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 9);
+        static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10);
 
         // this list here defines the priority of the speculators
         // the one with highest priority are listed first
@@ -1965,6 +2273,9 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         if (has_draft_mtp) {
             configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, params));
         }
+        if (has_draft_dflash) {
+            configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params));
+        }
     }
 
     std::vector<std::unique_ptr<common_speculative_impl>> impls = {};
@@ -1985,6 +2296,10 @@ common_speculative * common_speculative_init(common_params_speculative & params,
                 impls.push_back(std::make_unique<common_speculative_impl_draft_mtp>(config.params, n_seq));
                 break;
             }
+            case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: {
+                impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(config.params, n_seq));
+                break;
+            }
             case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
                 common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple);
 
@@ -2034,7 +2349,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
     }
 
     if (impls.empty()) {
-        LOG_WRN("%s: no implementations specified for speculative decoding\n", __func__);
+        SPC_TRC("%s", "no implementations specified for speculative decoding\n");
         return nullptr;
     }
 
@@ -2161,13 +2476,13 @@ void common_speculative_draft(common_speculative * spec) {
 
                 if (dp.n_max > 0) {
                     if (!result.empty() && (int) result.size() > dp.n_max) {
-                        LOG_DBG("%s: truncating draft to %d tokens\n", __func__, dp.n_max);
+                        SPC_DBG("truncating draft to %d tokens\n", dp.n_max);
                         result.resize(dp.n_max);
                     }
                 }
 
                 if (!result.empty()) {
-                    LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__,
+                    SPC_DBG("called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n",
                             common_speculative_type_to_str(impl.get()->type).c_str(), dp.prompt->size(),
                             impl.get()->n_call_draft, result.size());
 
@@ -2291,7 +2606,7 @@ void common_speculative_print_stats(const common_speculative * spec) {
             str_stats = ", #mean acc len = " + oss.str() + ", #acc rate/pos = (" + tmp.str() + ")";
         }
 
-        LOG_INF("statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = %6zu, #acc tokens = %5zu%s%s\n",
+        SPC_TRC("statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = %6zu, #acc tokens = %5zu%s%s\n",
                 common_speculative_type_to_str(impl->type).c_str(),
                 impl->n_call_begin, impl->n_call_draft, impl->n_call_accept,
                 impl->n_gen_drafts,

conversion/__init__.py

@@ -50,6 +50,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "DeepseekV2ForCausalLM": "deepseek",
     "DeepseekV3ForCausalLM": "deepseek",
     "DeepseekV32ForCausalLM": "deepseek",
+    "DFlashDraftModel": "qwen",
+    "DeepseekV4ForCausalLM": "deepseek",
     "DistilBertForMaskedLM": "bert",
     "DistilBertForSequenceClassification": "bert",
     "DistilBertModel": "bert",

conversion/base.py

@@ -1273,7 +1273,7 @@ class TextModel(ModelBase):
         if (f_norm_eps := self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon"], optional=True)) is not None:
             self.gguf_writer.add_layer_norm_eps(f_norm_eps)
             logger.info(f"gguf: layer norm epsilon = {f_norm_eps}")
-        if (n_experts := self.find_hparam(["num_local_experts", "num_experts"], optional=True)) is not None:
+        if (n_experts := self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True)) is not None:
             self.gguf_writer.add_expert_count(n_experts)
             logger.info(f"gguf: expert count = {n_experts}")
         if (n_experts_used := self.find_hparam(["num_experts_per_tok", "num_experts_per_token", "top_k_experts"], optional=True)) is not None:
@@ -1291,6 +1291,8 @@ class TextModel(ModelBase):
                 self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
             elif score_func == "softmax":
                 self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX)
+            elif score_func == "sqrtsoftplus":
+                self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS)
             else:
                 raise ValueError(f"Unsupported expert score gating function value: {score_func}")
             logger.info(f"gguf: expert score gating function = {score_func}")
@@ -2600,6 +2602,17 @@ class LazyTorchTensor(gguf.LazyBase):
         return cls._wrap_fn(func)(*args, **kwargs)
 
 
+if hasattr(torch, "float8_e8m0fnu"):
+    _torch_float8_e8m0 = torch.float8_e8m0fnu
+    LazyTorchTensor._dtype_map[_torch_float8_e8m0] = np.uint8
+    LazyTorchTensor._dtype_byteswap_map[_torch_float8_e8m0] = np.uint8
+    LazyTorchTensor._dtype_str_map["F8_E8M0"] = _torch_float8_e8m0
+else:
+    # Older torch builds do not expose F8_E8M0. Keep the raw bytes so callers
+    # that know the format can decode them explicitly.
+    LazyTorchTensor._dtype_str_map["F8_E8M0"] = torch.uint8
+
+
 def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> str:
     # TODO @ngxson : this won't work correctly if the model has both audio & vision encoders
     # maybe we should fallback to text model's arch in that case, since not many models have both

conversion/deepseek.py

@@ -1,15 +1,18 @@
 from __future__ import annotations
 
+import json
 import re
+from pathlib import Path
 
 from typing import Any, Callable, Iterable, TYPE_CHECKING
 
+import numpy as np
 import torch
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf, logger
+from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger
 
 from .qwen import QwenModel
 
@@ -467,3 +470,307 @@ class DeepseekV32Model(DeepseekV2Model):
         self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
         self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
         self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
+
+
+@ModelBase.register("DeepseekV4ForCausalLM")
+class DeepseekV4Model(TextModel):
+    model_arch = gguf.MODEL_ARCH.DEEPSEEK4
+    _skipped_mtp_tensors = 0
+
+    def __init__(self, *args, **kwargs):
+        type(self)._skipped_mtp_tensors = 0
+        super().__init__(*args, **kwargs)
+
+        with open(self.dir_model / "config.json", "r", encoding="utf-8") as f:
+            raw_hparams = json.load(f)
+        for key, value in raw_hparams.items():
+            self.hparams.setdefault(key, value)
+
+        self.block_count = self.hparams["num_hidden_layers"]
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        self._dsv4_fp8_dequantized: set[str] = set()
+        self._dsv4_bf16_tensors: set[str] = set()
+        self._dsv4_f32_tensors: set[str] = set()
+        self._dsv4_mxfp4_generated = False
+        self._collect_source_dtypes()
+
+        if type(self)._skipped_mtp_tensors:
+            logger.info("Skipping %d DeepSeek-V4 MTP tensor(s) for conversion v0", type(self)._skipped_mtp_tensors)
+
+        # add a default chat template; if the model has a built-in template, it will be overridden later
+        template_path = Path(__file__).parent.parent / "models" / "templates" / "deepseek-ai-DeepSeek-V4.jinja"
+        if template_path.is_file():
+            with open(template_path, "r", encoding="utf-8") as f:
+                self.gguf_writer.add_chat_template(f.read())
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if name.startswith("mtp."):
+            cls._skipped_mtp_tensors += 1
+            return None
+        return super().filter_tensors(item)
+
+    @staticmethod
+    def _float8_dtypes() -> tuple[torch.dtype, ...]:
+        return tuple(
+            dtype for dtype in (
+                getattr(torch, "float8_e4m3fn", None),
+                getattr(torch, "float8_e5m2", None),
+            ) if dtype is not None
+        )
+
+    @staticmethod
+    def _e8m0_to_float(scale: Tensor) -> Tensor:
+        torch_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None)
+        if torch_float8_e8m0 is not None and scale.dtype == torch_float8_e8m0:
+            return scale.float()
+
+        bits = scale.view(torch.uint8).float()
+        return torch.exp2(bits - 127.0)
+
+    def _collect_source_dtypes(self) -> None:
+        for name, gen in self.model_tensors.items():
+            dtype = gen().dtype
+            if dtype == torch.bfloat16:
+                self._dsv4_bf16_tensors.add(name)
+            elif dtype == torch.float32:
+                self._dsv4_f32_tensors.add(name)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"])
+        self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"])
+        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
+
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"])
+        self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"])
+        self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])
+        self.gguf_writer.add_swiglu_clamp_exp([hparams["swiglu_limit"]] * self.block_count)
+        self.gguf_writer.add_swiglu_clamp_shexp([hparams["swiglu_limit"]] * self.block_count)
+
+        self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"])
+        self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"])
+        self.gguf_writer.add_indexer_top_k(hparams["index_topk"])
+
+        self.gguf_writer.add_attention_output_group_count(hparams["o_groups"])
+        self.gguf_writer.add_attention_output_lora_rank(hparams["o_lora_rank"])
+        self.gguf_writer.add_attention_compress_ratios(hparams["compress_ratios"])
+        self.gguf_writer.add_attention_compress_rope_freq_base(hparams["compress_rope_theta"])
+        self.gguf_writer.add_hyper_connection_count(hparams["hc_mult"])
+        self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"])
+        self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
+        self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"])
+
+    def dequant_model(self):
+        fp8_dtypes = self._float8_dtypes()
+        tensors_to_remove: list[str] = []
+
+        def dequant_fp8_weight(weight: Tensor, scale: Tensor) -> Tensor:
+            out_features, in_features = weight.shape
+            scale_f = self._e8m0_to_float(scale)
+            scale_f = scale_f.repeat_interleave(128, 0)[:out_features]
+            scale_f = scale_f.repeat_interleave(128, 1)[:, :in_features]
+            return weight.float() * scale_f
+
+        for name in list(self.model_tensors.keys()):
+            if not name.endswith(".scale"):
+                continue
+            weight_name = name.removesuffix(".scale") + ".weight"
+            if weight_name not in self.model_tensors:
+                continue
+
+            weight = self.model_tensors[weight_name]
+            scale = self.model_tensors[name]
+            if weight().dtype not in fp8_dtypes:
+                continue
+
+            self.model_tensors[weight_name] = lambda w=weight, s=scale: dequant_fp8_weight(w(), s())
+            self._dsv4_fp8_dequantized.add(weight_name)
+            tensors_to_remove.append(name)
+
+        for name in tensors_to_remove:
+            del self.model_tensors[name]
+
+    @staticmethod
+    def _pack_mxfp4_blocks(weight: Tensor, scale: Tensor) -> np.ndarray:
+        packed = weight.contiguous().view(torch.uint8)
+        scale_u8 = scale.contiguous().view(torch.uint8)
+
+        out_features, packed_cols = packed.shape
+        logical_cols = packed_cols * 2
+        if logical_cols % 32 != 0:
+            raise ValueError(f"MXFP4 source row has {logical_cols} values, expected a multiple of 32")
+
+        n_blocks = logical_cols // 32
+        if tuple(scale_u8.shape) != (out_features, n_blocks):
+            raise ValueError(f"MXFP4 scale shape {tuple(scale_u8.shape)} does not match {(out_features, n_blocks)}")
+
+        src = packed.reshape(out_features, n_blocks, 16)
+        low = src & 0x0F
+        high = (src >> 4) & 0x0F
+
+        # The safetensors bytes store adjacent values as low/high nibbles.
+        # ggml MXFP4 blocks store values 0..15 in low nibbles and 16..31 in high nibbles.
+        vals = torch.stack((low, high), dim=-1).reshape(out_features, n_blocks, 32)
+        qs = vals[:, :, :16] | (vals[:, :, 16:] << 4)
+        raw = torch.cat((scale_u8.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
+        return raw.reshape(out_features, n_blocks * 17).cpu().numpy()
+
+    def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]:
+        n_experts = self.hparams["n_routed_experts"]
+        data: np.ndarray | None = None
+        consumed: list[str] = []
+
+        for eid in range(n_experts):
+            weight_name = f"layers.{bid}.ffn.experts.{eid}.{proj}.weight"
+            scale_name = f"layers.{bid}.ffn.experts.{eid}.{proj}.scale"
+            if weight_name not in self.model_tensors or scale_name not in self.model_tensors:
+                raise KeyError(f"Missing routed expert tensors for {weight_name}")
+
+            weight = LazyTorchTensor.to_eager(self.model_tensors[weight_name]())
+            scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())
+            packed = self._pack_mxfp4_blocks(weight, scale)
+            if data is None:
+                data = np.empty((n_experts, *packed.shape), dtype=packed.dtype)
+            data[eid] = packed
+            consumed.extend((weight_name, scale_name))
+
+        assert data is not None
+        new_name = self.format_tensor_name(tensor_key, bid)
+        shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4)
+        logger.info(f"{new_name}: repacked routed experts to MXFP4, shape = {{{', '.join(str(n) for n in reversed(shape))}}}")
+        self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
+
+        return consumed
+
+    def _write_hash_routing_tensors(self) -> list[str]:
+        consumed: list[str] = []
+
+        for bid in range(self.hparams["num_hash_layers"]):
+            name = f"layers.{bid}.ffn.gate.tid2eid"
+            if name not in self.model_tensors:
+                raise KeyError(f"Missing hash routing tensor {name}")
+
+            data_torch = LazyTorchTensor.to_eager(self.model_tensors[name]())
+            data = data_torch.to(torch.int32).cpu().numpy()
+            new_name = self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_TID2EID, bid, ".weight")
+            logger.info(f"{new_name}: converted hash routing table to I32, shape = {{{', '.join(str(n) for n in reversed(data.shape))}}}")
+            self.gguf_writer.add_tensor(new_name, data)
+            consumed.append(name)
+
+        return consumed
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        if self._dsv4_mxfp4_generated:
+            return ()
+
+        consumed: list[str] = self._write_hash_routing_tensors()
+        for bid in range(self.block_count):
+            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP))
+            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP))
+            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP))
+
+        for name in consumed:
+            del self.model_tensors[name]
+
+        self._dsv4_mxfp4_generated = True
+        return ()
+
+    def _format_dsv4_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> str:
+        return self.format_tensor_name(key, bid, suffix)
+
+    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
+        root_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+            "embed.weight": (gguf.MODEL_TENSOR.TOKEN_EMBD, ".weight"),
+            "norm.weight": (gguf.MODEL_TENSOR.OUTPUT_NORM, ".weight"),
+            "head.weight": (gguf.MODEL_TENSOR.OUTPUT, ".weight"),
+            "hc_head_fn": (gguf.MODEL_TENSOR.HC_HEAD_FN, ".weight"),
+            "hc_head_base": (gguf.MODEL_TENSOR.HC_HEAD_BASE, ".weight"),
+            "hc_head_scale": (gguf.MODEL_TENSOR.HC_HEAD_SCALE, ".weight"),
+        }
+        if name in root_map:
+            return root_map[name]
+
+        match = re.match(r"layers\.(\d+)\.(.+)$", name)
+        if match is None:
+            raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
+
+        layer = int(match.group(1))
+        if bid != layer:
+            raise ValueError(f"Tensor {name!r} parsed bid {bid} but layer name has {layer}")
+
+        layer_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+            "hc_attn_fn": (gguf.MODEL_TENSOR.HC_ATTN_FN, ".weight"),
+            "hc_attn_base": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"),
+            "hc_attn_scale": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"),
+            "hc_ffn_fn": (gguf.MODEL_TENSOR.HC_FFN_FN, ".weight"),
+            "hc_ffn_base": (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"),
+            "hc_ffn_scale": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"),
+            "attn.attn_sink": (gguf.MODEL_TENSOR.ATTN_SINKS, ".weight"),
+            "attn.wq_a.weight": (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"),
+            "attn.wq_b.weight": (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"),
+            "attn.q_norm.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"),
+            "attn.wkv.weight": (gguf.MODEL_TENSOR.ATTN_KV, ".weight"),
+            "attn.kv_norm.weight": (gguf.MODEL_TENSOR.ATTN_KV_NORM, ".weight"),
+            "attn.wo_a.weight": (gguf.MODEL_TENSOR.ATTN_OUT_A, ".weight"),
+            "attn.wo_b.weight": (gguf.MODEL_TENSOR.ATTN_OUT_B, ".weight"),
+            "attn.compressor.ape": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_APE, ".weight"),
+            "attn.compressor.wkv.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_WKV, ".weight"),
+            "attn.compressor.wgate.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, ".weight"),
+            "attn.compressor.norm.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_NORM, ".weight"),
+            "attn.indexer.wq_b.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_Q_B, ".weight"),
+            "attn.indexer.weights_proj.weight": (gguf.MODEL_TENSOR.INDEXER_PROJ, ".weight"),
+            "attn.indexer.compressor.ape": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, ".weight"),
+            "attn.indexer.compressor.wkv.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_WKV, ".weight"),
+            "attn.indexer.compressor.wgate.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, ".weight"),
+            "attn.indexer.compressor.norm.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_NORM, ".weight"),
+            "attn_norm.weight": (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"),
+            "ffn_norm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"),
+            "ffn.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"),
+            "ffn.gate.bias": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"),
+            "ffn.gate.tid2eid": (gguf.MODEL_TENSOR.FFN_GATE_TID2EID, ".weight"),
+            "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
+            "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
+            "ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
+        }
+
+        tensor_name = match.group(2)
+        if tensor_name in layer_map:
+            return layer_map[tensor_name]
+
+        if re.match(r"ffn\.experts\.\d+\.w[123]\.(weight|scale)$", tensor_name):
+            return gguf.MODEL_TENSOR.FFN_GATE_EXP, ".weight"
+
+        raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if re.match(r"layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(weight|scale)$", name):
+            return []
+
+        tensor_key, suffix = self._map_dsv4_tensor_name(name, bid)
+        if tensor_key == gguf.MODEL_TENSOR.FFN_GATE_TID2EID:
+            return []
+
+        return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
+
+    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+        del new_name, bid  # unused
+
+        if name in self._dsv4_fp8_dequantized and n_dims >= 2:
+            return gguf.GGMLQuantizationType.Q8_0
+        if name in self._dsv4_f32_tensors:
+            return gguf.GGMLQuantizationType.F32
+        if name in self._dsv4_bf16_tensors and n_dims >= 2:
+            return gguf.GGMLQuantizationType.BF16
+
+        return False
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        self._is_mxfp4 = True
+        self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE

conversion/llama.py

@@ -73,7 +73,7 @@ class LlamaModel(TextModel):
             target_num_layers = target_config["num_hidden_layers"]
             target_layers = [2, target_num_layers // 2, target_num_layers - 3]
             logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")
-            self.gguf_writer.add_array(f"{self.gguf_writer.arch}.target_layers", target_layers)
+            self.gguf_writer.add_target_layers(target_layers)
 
             # target_hidden_size: prefer eagle3 config, fallback to target config
             if eagle3_raw_config.get("target_hidden_size") is not None:
@@ -83,12 +83,12 @@ class LlamaModel(TextModel):
                 target_hidden_size = target_config["hidden_size"]
                 src = "target model config"
             logger.info(f"EAGLE-3: target_hidden_size = {target_hidden_size} (from {src})")
-            self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.target_hidden_size", target_hidden_size)
+            self.gguf_writer.add_target_hidden_size(target_hidden_size)
 
             # norm_before_residual (RedHat-style eagle3 specific)
             norm_before_residual = eagle3_raw_config.get("norm_before_residual", False)
             logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")
-            self.gguf_writer.add_bool(f"{self.gguf_writer.arch}.norm_before_residual", norm_before_residual)
+            self.gguf_writer.add_norm_before_residual(norm_before_residual)
 
     def set_vocab(self):
         # eagle3: use tokenizer from target model if provided

conversion/qwen.py

@@ -625,3 +625,51 @@ class Qwen3_5TextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReor
 @ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")
 class Qwen3_5MoeTextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase):
     model_arch = gguf.MODEL_ARCH.QWEN35MOE
+
+
+@ModelBase.register("DFlashDraftModel")
+class DFlashModel(Qwen3Model):
+    model_arch = gguf.MODEL_ARCH.DFLASH
+
+    def set_vocab(self):
+        if self.target_model_dir is None:
+            raise ValueError(
+                "DFlash draft model requires --target-model-dir to be specified. "
+                "Please provide the path to the target model directory containing the tokenizer."
+            )
+        logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}")
+        original_dir = self.dir_model
+        self.dir_model = self.target_model_dir
+        super().set_vocab()
+        self.dir_model = original_dir
+
+        mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id")
+        if mask_token_id is not None:
+            self.gguf_writer.add_mask_token_id(mask_token_id)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        block_size = self.hparams.get("block_size", 16)
+        self.gguf_writer.add_block_size(block_size)
+        dflash_config = self.hparams.get("dflash_config", {})
+
+        target_layer_ids = dflash_config.get("target_layer_ids", [])
+        if target_layer_ids:
+            extract_layer_ids = [i + 1 for i in target_layer_ids]
+            self.gguf_writer.add_target_layers(extract_layer_ids)
+
+        use_sliding_window = self.hparams.get("use_sliding_window", False)
+        sliding_window = self.hparams.get("sliding_window")
+        layer_types = self.hparams.get("layer_types")
+        if use_sliding_window and sliding_window and layer_types:
+            is_swa = [lt == "sliding_attention" for lt in layer_types]
+            self.gguf_writer.add_sliding_window(sliding_window)
+            self.gguf_writer.add_sliding_window_pattern(is_swa)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("model."):
+            name = "model." + name
+        return super().filter_tensors((name, gen))

expose.h

@@ -186,14 +186,12 @@ struct sd_load_model_inputs
 {
     const char * model_filename = nullptr;
     const char * executable_path = nullptr;
-    const int kcpp_main_device = -1;
+    const char * backend = nullptr;
     const int threads = 0;
     const int quant = 0;
     const bool flash_attention = false;
-    const bool offload_cpu = false;
+    const char * params_backend = nullptr;
     const bool use_mmap = false;
-    const int kcpp_vae_device = -1;
-    const int kcpp_clip_device = -1;
     const bool diffusion_conv_direct = false;
     const bool vae_conv_direct = false;
     const bool taesd = false;
@@ -211,8 +209,10 @@ struct sd_load_model_inputs
     const char * upscaler_filename = nullptr;
     const int img_hard_limit = 0;
     const int img_soft_limit = 0;
-    const float max_vram = 0.f;
+    const char * max_vram = nullptr;
+    const char * split_mode = nullptr;
     const bool stream_layers = false;
+    const bool auto_fit = false;
     const char * devices_override = nullptr;
     const bool quiet = false;
     const int debugmode = 0;

ggml/src/ggml-backend-meta.cpp

@@ -1144,6 +1144,11 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m
         ggml_context          * simple_ctx = stc.ctxs[j].get();
         ggml_backend_buffer_t   simple_buf = buf_ctx->bufs[j].get();
 
+        if ((simple_buf != nullptr) && ggml_backend_buffer_is_multi_buffer(simple_buf)) {
+            // see https://github.com/ggml-org/llama.cpp/issues/22197
+            GGML_ABORT("multi buffers are not supported by the meta backend");
+        }
+
         if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) {
             // TODO: the following assert fails for llama-parallel even though the results are correct:
             // GGML_ASSERT(ggml_is_contiguously_allocated(tensor));
@@ -1245,9 +1250,8 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer
 
 static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
-    GGML_ASSERT(ggml_is_contiguous(tensor));
-
     const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
+    GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
     if (split_state.n_segments != 1 || split_state.nr[0] != 1) {
         GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS);
@@ -1360,9 +1364,8 @@ static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 
 static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
     const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer);
-    GGML_ASSERT(ggml_is_contiguous(tensor));
-
     const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
+    GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
     if (split_state.n_segments != 1 || split_state.nr[0] != 1) {
         GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS);

ggml/src/ggml-common.h

@@ -1111,11 +1111,12 @@ GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16)
     -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113,
 GGML_TABLE_END()
 
-// e2m1 values (doubled)
+// e2m1 values (doubled), shared by MXFP4 and NVFP4
 // ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
-GGML_TABLE_BEGIN(int8_t, kvalues_mxfp4, 16)
+GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16)
     0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12,
 GGML_TABLE_END()
+#define kvalues_mxfp4 kvalues_fp4
 
 #define NGRID_IQ1S 2048
 #define IQ1S_DELTA 0.125f

ggml/src/ggml-cpu/arch-fallback.h

@@ -72,7 +72,6 @@
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
 #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
 // quants.c
-#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4

ggml/src/ggml-cpu/arch/arm/quants.c

@@ -812,10 +812,10 @@ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
         const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
         const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
         const float32x4_t nvsc = {
-            ggml_ue4m3_to_fp32(x[ib].d[0]),
-            ggml_ue4m3_to_fp32(x[ib].d[1]),
-            ggml_ue4m3_to_fp32(x[ib].d[2]),
-            ggml_ue4m3_to_fp32(x[ib].d[3])
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[3])
         };
         const float32x4_t scales = vmulq_f32(nvsc, (float32x4_t){dy0, dy0, dy1, dy1});
 

ggml/src/ggml-cpu/arch/x86/quants.c

@@ -935,7 +935,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
 
 #if defined __AVX2__
 
-    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_mxfp4);
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
     const __m128i m4b  = _mm_set1_epi8(0x0f);
     const __m256i mone = _mm256_set1_epi16(1);
 
@@ -964,7 +964,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
     sumf = hsum_float_8(_mm256_add_ps(accum1, accum2));
 
 #elif defined __AVX__
-    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_mxfp4);
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
     const __m128i m4b  = _mm_set1_epi8(0x0f);
 
     __m256 accum = _mm256_setzero_ps();
@@ -994,14 +994,152 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
         int sumi1 = 0;
         int sumi2 = 0;
         for (int j = 0; j < QK_MXFP4/2; ++j) {
-            sumi1 += y[ib].qs[j +          0] * kvalues_mxfp4[x[ib].qs[j] & 0xf];
-            sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >>  4];
+            sumi1 += y[ib].qs[j +          0] * kvalues_fp4[x[ib].qs[j] & 0xf];
+            sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_fp4[x[ib].qs[j] >>  4];
         }
         sumf += d * (sumi1 + sumi2);
     }
     *s = sumf;
 }
 
+void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK_NVFP4 == 0);
+
+    const block_nvfp4 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK_NVFP4;
+    int ib = 0;
+    float sumf = 0;
+
+#if defined(__AVX2__)
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+    const __m256i mone = _mm256_set1_epi16(1);
+
+    __m256 accum = _mm256_setzero_ps();
+    for(; ib < nb; ib++){
+
+        const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0));
+        const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16));
+
+        const __m256i q8_01 = _mm256_loadu_si256((const __m256i *)y[2*ib + 0].qs);
+        const __m256i q8_23 = _mm256_loadu_si256((const __m256i *)y[2*ib + 1].qs);
+
+        const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b));
+        const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b));
+        const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b));
+        const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b));
+
+        //reordering
+        const __m256i q4_01 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_01_lo,q4_01_hi), _mm_unpacklo_epi64(q4_01_lo,q4_01_hi));
+        const __m256i q4_23 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_23_lo,q4_23_hi),_mm_unpacklo_epi64(q4_23_lo,q4_23_hi));
+
+        const __m256i p01 = mul_add_epi8(q4_01,q8_01);
+        const __m256i p_1 = _mm256_madd_epi16(p01, mone);
+
+        const __m256i p23 = mul_add_epi8(q4_23,q8_23);
+        const __m256i p_2 = _mm256_madd_epi16(p23, mone);
+
+        const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
+        const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
+
+        const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0;
+        const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0;
+        const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1;
+        const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1;
+
+        const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0));
+        const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2));
+
+        accum = _mm256_fmadd_ps(scales01, _mm256_cvtepi32_ps(p_1), accum);
+        accum = _mm256_fmadd_ps(scales23, _mm256_cvtepi32_ps(p_2), accum);
+    }
+    sumf = hsum_float_8(accum);
+
+#elif defined(__AVX__)
+
+    const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4);
+    const __m128i m4b  = _mm_set1_epi8(0x0f);
+
+    __m256 accum = _mm256_setzero_ps();
+    for(; ib < nb; ib++){
+
+        const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0));
+        const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16));
+
+        const __m128i q8_0 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 0));
+        const __m128i q8_1 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 16));
+        const __m128i q8_2 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 0));
+        const __m128i q8_3 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 16));
+
+        const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b));
+        const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b));
+        const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b));
+        const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b));
+
+        const __m128i q4_0 = _mm_unpacklo_epi64(q4_01_lo, q4_01_hi);
+        const __m128i q4_1 = _mm_unpackhi_epi64(q4_01_lo, q4_01_hi);
+        const __m128i q4_2 = _mm_unpacklo_epi64(q4_23_lo, q4_23_hi);
+        const __m128i q4_3 = _mm_unpackhi_epi64(q4_23_lo, q4_23_hi);
+
+        const __m128i p0_i32 = mul_sum_i8_pairs(q4_0, q8_0);
+        const __m128i p1_i32 = mul_sum_i8_pairs(q4_1, q8_1);
+        const __m128i p2_i32 = mul_sum_i8_pairs(q4_2, q8_2);
+        const __m128i p3_i32 = mul_sum_i8_pairs(q4_3, q8_3);
+
+        const __m128 p0 = _mm_cvtepi32_ps(p0_i32);
+        const __m128 p1 = _mm_cvtepi32_ps(p1_i32);
+        const __m128 p2 = _mm_cvtepi32_ps(p2_i32);
+        const __m128 p3 = _mm_cvtepi32_ps(p3_i32);
+
+        const __m256 p01 = _mm256_set_m128(p1, p0);
+        const __m256 p23 = _mm256_set_m128(p3, p2);
+
+        const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
+        const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
+
+        const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0;
+        const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0;
+        const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1;
+        const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1;
+
+        const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0));
+        const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2));
+
+        accum = _mm256_add_ps(accum, _mm256_mul_ps(p01, scales01));
+        accum = _mm256_add_ps(accum, _mm256_mul_ps(p23, scales23));
+    }
+    sumf = hsum_float_8(accum);
+
+#endif
+
+    for (;ib < nb; ++ib) {
+        for (int s_idx = 0; s_idx < 4; ++s_idx) {
+            const float d = GGML_CPU_UE4M3_TO_FP32(x[ib].d[s_idx]);
+            const int q8_block = s_idx / 2;
+            const int q8_off   = (s_idx % 2) * QK_NVFP4_SUB;
+            const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d);
+
+            int sumi_lo = 0, sumi_hi = 0;
+            for (int j = 0; j < QK_NVFP4_SUB/2; ++j) {
+                const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j];
+                sumi_lo += y[2*ib + q8_block].qs[q8_off + j +               0] * kvalues_fp4[qv & 0xf];
+                sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_fp4[qv >>  4];
+            }
+
+            sumf += dy * d * (sumi_lo + sumi_hi);
+        }
+    }
+    *s = sumf;
+}
+
 void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
     const int nb = n / qk;

ggml/src/ggml-cpu/ggml-cpu.c

@@ -83,6 +83,9 @@ float ggml_table_f32_f16[1 << 16];
 // precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h)
 float ggml_table_f32_e8m0_half[1 << 8];
 
+// precomputed f32 table for ue4m3 (1 KB) (simd-mappings.h)
+float ggml_table_f32_ue4m3[1 << 8];
+
 #if defined(__ARM_ARCH)
 struct ggml_arm_arch_features_type {
     int sve_cnt;
@@ -4647,6 +4650,11 @@ void ggml_cpu_init(void) {
                 ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i);
             }
 
+            // initialize UE4M3 table (256 entries)
+            for (int i = 0; i < (1 << 8); ++i) {
+                ggml_table_f32_ue4m3[i] = ggml_ue4m3_to_fp32(i);
+            }
+
             const uint64_t t_end = ggml_time_us(); UNUSED(t_end);
 
             GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0);

ggml/src/ggml-cpu/llamafile/sgemm.cpp

@@ -2321,24 +2321,28 @@ class tinyBLAS_Q0_PPC {
     }
 
     void matmul(int64_t m, int64_t n) {
-    #if defined(_AIX) || defined(__BIG_ENDIAN__)
-        mnpack(0, m, 0, n);
-    #else
-        const int64_t mc = 64;
-        const int64_t kc = 64;
+        int64_t mc = 64;
         int64_t nc = 64;
+        int64_t kc = 64;
+        int64_t n_chunk = 64;
+    #if defined(_AIX) || defined(__BIG_ENDIAN__)
+        mc = 32;
+        nc = 32;
+        kc = 32;
+        n_chunk = 32
+    #endif
         int64_t n_aligned = 0;
-        if (n % 64 == 0) {
+        if (n % n_chunk == 0) {
             n_aligned = n;
         } else if (n == 4) {
             n_aligned = 4;
-        } else if (n < 64) {
+        } else if (n < n_chunk) {
             n_aligned = (n / 8) * 8;
         } else {
-            n_aligned = (n / 64) * 64;
+            n_aligned = (n / n_chunk) * n_chunk;
         }
         if (n_aligned > 0) {
-            if (n_aligned % 64 == 0)      nc = 64;
+            if (n_aligned % n_chunk == 0) nc = n_chunk;
             else if (n_aligned == n)      nc = n;
             else if (n_aligned % 32 == 0) nc = 32;
             else if (n_aligned % 24 == 0) nc = 24;
@@ -2354,7 +2358,6 @@ class tinyBLAS_Q0_PPC {
         } else {
             mnpack(0, m, 0, n);
         }
-    #endif
     }
 
   private:
@@ -3195,16 +3198,19 @@ class tinyBLAS_PPC {
     }
 
     void matmul(int64_t m, int64_t n) {
+        int64_t mc = 256;
+        int64_t nc = 256;
+        int64_t kc = 256;
     #if defined(_AIX) || defined(__BIG_ENDIAN__)
-        mnpack(0, m, 0, n);
-    #else
-        int64_t mc = 256; int64_t nc = 256; int64_t kc = 256;
+        mc = 128;
+        nc = 128;
+        kc = 128;
+    #endif
         if (m % mc == 0 && n % nc == 0 && k % kc == 0) {
             matmul_tiled(m, n, mc, nc, kc);
         } else {
             mnpack(0, m, 0, n);
         }
-    #endif
     }
 
   private:

ggml/src/ggml-cpu/ops.cpp

@@ -1913,7 +1913,11 @@ static void ggml_compute_forward_concat_any(
     GGML_ASSERT(dim >= 0 && dim < 4);
 
     int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
+    if (dim == 0) {
+        o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
+    } else {
+        o[dim] = src0->ne[dim];
+    }
 
     const char * x;
 
@@ -1921,8 +1925,8 @@ static void ggml_compute_forward_concat_any(
     for (int i3 = 0; i3 < ne3; i3++) {
         for (int i2 = ith; i2 < ne2; i2 += nth) {
             for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
+                for (int i0 = 0; i0 < ne0/ggml_blck_size(dst->type); i0++) {
+                    if (i0 < ne00/ggml_blck_size(src0->type) && i1 < ne01 && i2 < ne02 && i3 < ne03) {
                         x = (const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03;
                     } else {
                         x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13;
@@ -2071,6 +2075,14 @@ void ggml_compute_forward_concat(
     ggml_tensor * dst) {
 
     const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    if (ggml_is_quantized(src0->type)) {
+        GGML_ASSERT(ggml_is_contiguous(src0));
+        GGML_ASSERT(ggml_is_contiguous(src1));
+        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
+        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
+    }
 
     switch (src0->type) {
         case GGML_TYPE_F16:

ggml/src/ggml-cpu/simd-mappings.h

@@ -120,6 +120,10 @@ extern float ggml_table_f32_f16[1 << 16];
 // defined in ggml-cpu.c, initialized in ggml_cpu_init()
 extern float ggml_table_f32_e8m0_half[1 << 8];
 
+// precomputed f32 table for ue4m3 (1 KB)
+// defined in ggml-cpu.c, initialized in ggml_cpu_init()
+extern float ggml_table_f32_ue4m3[1 << 8];
+
 // Use lookup table for E8M0 on x86 (faster than bit manipulation)
 #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
 #define GGML_CPU_E8M0_TO_FP32_HALF(x) ggml_table_f32_e8m0_half[(uint8_t)(x)]
@@ -127,6 +131,13 @@ extern float ggml_table_f32_e8m0_half[1 << 8];
 #define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x)
 #endif
 
+// Use lookup table for UE4M3 on x86 and ARM (faster than bit manipulation)
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__ARM_NEON)
+#define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)]
+#else
+#define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x)
+#endif
+
 // On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32,
 // so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON.
 // This is also true for POWER9.

ggml/src/ggml-cuda/concat.cu

@@ -152,8 +152,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml
                         src0_d + i3*(src0->nb[3] / sizeof(T)),
                         src1_d + i3*(src1->nb[3] / sizeof(T)),
                         dst_d  + i3*( dst->nb[3] / sizeof(T)),
-                        src0->ne[0], src0->ne[1], src0->ne[2],
-                        dst->ne[0],  dst->ne[1],  dst->ne[2], dim, stream);
+                        ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
+                        ggml_row_size(dst->type, dst->ne[0])/sizeof(T),  dst->ne[1],  dst->ne[2], dim, stream);
             }
         } else {
             const size_t size0 = ggml_nbytes(src0);
@@ -163,6 +163,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml
             CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
         }
     } else {
+        GGML_ASSERT(!ggml_is_quantized(src0->type));
+
         dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]);
         auto launch_kernel = [&](auto dim) {
             concat_non_cont<T, dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
@@ -204,24 +206,34 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 
     GGML_ASSERT(src0->type == src1->type);
     GGML_ASSERT(dst->type  == src0->type);
-    GGML_ASSERT(!ggml_is_quantized(src0->type));
-    GGML_ASSERT(ggml_blck_size(src0->type) == 1);
 
-    switch (ggml_type_size(src0->type)) {
-        case 1:
-            concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
-            break;
-        case 2:
-            concat_cuda<uint16_t>(src0, src1, dst, dim, stream);
-            break;
-        case 4:
-            concat_cuda<uint32_t>(src0, src1, dst, dim, stream);
-            break;
-        case 8:
-            concat_cuda<uint64_t>(src0, src1, dst, dim, stream);
-            break;
-        default:
-            GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type));
-            break;
+    if (ggml_is_quantized(src0->type)) {
+        GGML_ASSERT(ggml_is_contiguous(src0));
+        GGML_ASSERT(ggml_is_contiguous(src1));
+        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
+        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
+
+        // if tensors are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
+        concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
+    } else {
+        GGML_ASSERT(ggml_blck_size(src0->type) == 1);
+
+        switch (ggml_type_size(src0->type)) {
+            case 1:
+                concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
+                break;
+            case 2:
+                concat_cuda<uint16_t>(src0, src1, dst, dim, stream);
+                break;
+            case 4:
+                concat_cuda<uint32_t>(src0, src1, dst, dim, stream);
+                break;
+            case 8:
+                concat_cuda<uint64_t>(src0, src1, dst, dim, stream);
+                break;
+            default:
+                GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type));
+                break;
+        }
     }
 }

ggml/src/ggml-cuda/conv-transpose-1d.cu

@@ -11,31 +11,32 @@ static  __global__ void conv_transpose_1d_kernel(
         return;
     }
 
-    int out_index = global_index / dst_ne0;
+    int out_t = global_index % dst_ne0;
+    int out_ch = (global_index / dst_ne0) % dst_ne1;
+    int plane = global_index / (dst_ne0 * dst_ne1);
 
     float accumulator = 0;
 
     for (int c = 0; c < src0_ne2; c++) {
-        int idx = global_index % dst_ne0;
+        int kernel_offset = src0_ne0 * (out_ch + src0_ne1 * c);
+        int input_offset = src1_ne0 * (c + src1_ne1 * plane);
 
-        int kernel_offset = (src0_ne0 * src0_ne1 * c) + (out_index * src0_ne0);
-        int input_offset = src1_ne0 * c;
+        for (int k = 0; k < src0_ne0; k++) {
+            int input_numer = out_t + p0 - k*d0;
+            if (input_numer < 0 || input_numer % s0 != 0) {
+                continue;
+            }
 
-        int i_min = (idx >= src0_ne0) ? ((idx - src0_ne0 + s0) / s0) : 0;
-        int i_max_val = idx / s0;
-        int i_max = (i_max_val < src1_ne0) ? i_max_val : (src1_ne0 - 1);
+            int input_t = input_numer / s0;
+            if (input_t >= src1_ne0) {
+                continue;
+            }
 
-        for (int i = i_min; i <= i_max; i++) {
-            int weight_idx = idx - i*s0;
-
-            float kernel_weight = src0[kernel_offset + weight_idx];
-            float input_value =  src1[input_offset+i];
-
-            accumulator += kernel_weight * input_value;
+            accumulator += src0[kernel_offset + k] * src1[input_offset + input_t];
         }
     }
     dst[global_index] = accumulator;
-    GGML_UNUSED_VARS(p0, d0, src0_ne3, src1_ne3, dst_ne3, src1_ne1, dst_ne1, src1_ne2, dst_ne2);
+    GGML_UNUSED_VARS(src0_ne3, src1_ne2, src1_ne3, dst_ne2, dst_ne3);
 }
 
 static void conv_transpose_1d_f32_f32_cuda(

ggml/src/ggml-cuda/cpy.cu

@@ -386,6 +386,46 @@ static void ggml_cpy_f32_iq4_nl_cuda(
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
+// check if a same-type copy reduces to a 2D strided copy (height rows of width
+// contiguous bytes), so it can use cudaMemcpy2DAsync instead of the scalar kernel
+static bool ggml_cuda_cpy_as_memcpy_2d(const ggml_tensor * src0, const ggml_tensor * src1,
+        size_t & width, size_t & height, size_t & spitch, size_t & dpitch) {
+    // require matching shape: a reshaped copy maps elements by flat order, which the
+    // prefix walk below does not handle
+    if (src0->type != src1->type || !ggml_are_same_shape(src0, src1)) {
+        return false;
+    }
+
+    // grow the contiguous prefix block shared by both tensors
+    size_t block_nb = ggml_element_size(src0);
+    int d = 0;
+    for (; d < GGML_MAX_DIMS; ++d) {
+        if (src0->nb[d] != block_nb || src1->nb[d] != block_nb) {
+            break;
+        }
+        block_nb *= src0->ne[d];
+    }
+
+    // d == 0: nothing contiguous; d == GGML_MAX_DIMS: fully contiguous (handled by memcpy)
+    if (d == 0 || d == GGML_MAX_DIMS) {
+        return false;
+    }
+
+    // dim d carries the rows; everything above it must be a single element
+    for (int i = d + 1; i < GGML_MAX_DIMS; ++i) {
+        if (src0->ne[i] != 1) {
+            return false;
+        }
+    }
+
+    width  = block_nb;
+    height = src0->ne[d];
+    spitch = src0->nb[d];
+    dpitch = src1->nb[d];
+
+    return spitch >= width && dpitch >= width;
+}
+
 void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, ggml_tensor * src1) {
     const int64_t ne = ggml_nelements(src0);
     GGML_ASSERT(ne == ggml_nelements(src1));
@@ -421,6 +461,8 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
     const bool can_be_transposed = nb01 == (int64_t)ggml_element_size(src0) &&
         src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0);
 
+    size_t mc_width = 0, mc_height = 0, mc_spitch = 0, mc_dpitch = 0;
+
     if (src0->type == src1->type && contiguous_srcs) {
         GGML_ASSERT(ggml_nbytes(src0) == ggml_nbytes(src1));
 #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY)
@@ -431,6 +473,9 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
         {
             CUDA_CHECK(cudaMemcpyAsync(src1_ddc, src0_ddc, ggml_nbytes(src0), cudaMemcpyDeviceToDevice, main_stream));
         }
+    } else if (ggml_cuda_cpy_as_memcpy_2d(src0, src1, mc_width, mc_height, mc_spitch, mc_dpitch)) {
+        CUDA_CHECK(cudaMemcpy2DAsync(src1_ddc, mc_dpitch, src0_ddc, mc_spitch,
+                                     mc_width, mc_height, cudaMemcpyDeviceToDevice, main_stream));
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) {
         if (can_be_transposed) {
             ggml_cpy_scalar_cuda<float, float, true>

ggml/src/ggml-cuda/fattn-common.cuh

@@ -664,7 +664,7 @@ constexpr __device__ dequantize_V_t get_dequantize_V() {
 template <int ncols1>
 __launch_bounds__(FATTN_KQ_STRIDE/2, 1)
 static __global__ void flash_attn_mask_to_KV_max(
-        const half2 * __restrict__ mask, int * __restrict__ KV_max, const int ne30, const int s31, const int s33) {
+        const half2 * __restrict__ mask, int * __restrict__ KV_max, const int ne30, const int64_t s31, const int64_t s33) {
     const int ne31     = gridDim.x;
     const int tid      = threadIdx.x;
     const int sequence = blockIdx.y;
@@ -1089,8 +1089,8 @@ void launch_fattn(
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
     if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
-        const int s31 = mask->nb[1] / sizeof(half2);
-        const int s33 = mask->nb[3] / sizeof(half2);
+        const int64_t s31 = mask->nb[1] / sizeof(half2);
+        const int64_t s33 = mask->nb[3] / sizeof(half2);
 
         const dim3 blocks_num_KV_max(ntiles_x, Q->ne[3], 1);
         const dim3 block_dim_KV_max(FATTN_KQ_STRIDE/2, 1, 1);

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -2003,6 +2003,10 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112,  64)
 DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128,  64)
 DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256,  64)
 
+extern DECL_FATTN_MMA_F16_CASE(512, 512,  4,  2);
+extern DECL_FATTN_MMA_F16_CASE(512, 512,  8,  2);
+extern DECL_FATTN_MMA_F16_CASE(512, 512, 16,  2);
+extern DECL_FATTN_MMA_F16_CASE(512, 512, 32,  2);
 extern DECL_FATTN_MMA_F16_CASE(512, 512,  2,  4);
 extern DECL_FATTN_MMA_F16_CASE(512, 512,  4,  4);
 extern DECL_FATTN_MMA_F16_CASE(512, 512,  8,  4);

ggml/src/ggml-cuda/fattn-tile.cuh

@@ -76,6 +76,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 16, 256, 2,  64,  64)
 
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  2,  64, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 2,  64,  64)
@@ -144,6 +145,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_nv
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 16, 256, 2,  32,  64)
 
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  2,  64, 2,  32,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  32,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  32,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 2,  32,  64)
@@ -219,6 +221,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_am
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 32, 512, 1, 128,  64)
 
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  2,  64, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 2,  64,  64)
@@ -296,6 +299,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_am
 
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(320, 256, 32, 256, 2, 128,  64)
 
+    GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  2,  64, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  4, 128, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512,  8, 256, 2,  64,  64)
     GGML_CUDA_FATTN_TILE_CONFIG_CASE(512, 512, 16, 256, 4,  64,  64)
@@ -1308,12 +1312,12 @@ static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggm
             return;
         }
 
-        if constexpr (DV <= 256) {
-            if (use_gqa_opt && gqa_ratio % 2 == 0) {
-                launch_fattn_tile_switch_ncols1<DKQ, DV, 2, use_logit_softcap>(ctx, dst);
-                return;
-            }
+        if (use_gqa_opt && gqa_ratio % 2 == 0) {
+            launch_fattn_tile_switch_ncols1<DKQ, DV, 2, use_logit_softcap>(ctx, dst);
+            return;
+        }
 
+        if constexpr (DV <= 256) {
             launch_fattn_tile_switch_ncols1<DKQ, DV, 1, use_logit_softcap>(ctx, dst);
             return;
         }

ggml/src/ggml-cuda/fattn.cu

@@ -99,12 +99,12 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
         return;
     }
 
-    if constexpr (DKQ <= 256) {
-        if (use_gqa_opt && gqa_ratio > 1) {
-            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst);
-            return;
-        }
+    if (use_gqa_opt && gqa_ratio > 1) {
+        ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst);
+        return;
+    }
 
+    if constexpr (DKQ <= 256) {
         ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 1>(ctx, dst);
     } else {
         GGML_ABORT("fatal error");
@@ -338,6 +338,26 @@ enum best_fattn_kernel {
     BEST_FATTN_KERNEL_MMA_F16  = 400,
 };
 
+static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+        case GGML_TYPE_F16:
+            return true;
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+#ifndef GGML_CUDA_FA_ALL_QUANTS
+            return false;
+#endif // GGML_CUDA_FA_ALL_QUANTS
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q5_1: // kcpp: support q5_1 kv
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_BF16:
+            return true;
+        default:
+            return false;
+    }
+}
+
 static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
 #ifndef FLASH_ATTN_AVAILABLE
     GGML_UNUSED(device); GGML_UNUSED(dst);
@@ -428,22 +448,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     }
 #endif // GGML_CUDA_FA_ALL_QUANTS
 
-    switch (K->type) {
-        case GGML_TYPE_F32:
-        case GGML_TYPE_F16:
-            break;
-        case GGML_TYPE_Q4_1:
-        case GGML_TYPE_Q5_0:
-#ifndef GGML_CUDA_FA_ALL_QUANTS
-            return BEST_FATTN_KERNEL_NONE;
-#endif // GGML_CUDA_FA_ALL_QUANTS
-        case GGML_TYPE_Q4_0:
-        case GGML_TYPE_Q5_1: //kcpp: support q5_1 kv
-        case GGML_TYPE_Q8_0:
-        case GGML_TYPE_BF16:
-            break;
-        default:
-            return BEST_FATTN_KERNEL_NONE;
+    if (!ggml_cuda_fattn_kv_type_supported(K->type) || !ggml_cuda_fattn_kv_type_supported(V->type)) {
+        return BEST_FATTN_KERNEL_NONE;
     }
 
     if (mask && mask->ne[2] != 1) {

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -10,6 +10,7 @@ gated_delta_net_cuda(const float * q,
                                      const float * beta,
                                      const float * curr_state,
                                      float *       dst,
+                                     float *       state,
                                      int64_t       H,
                                      int64_t       n_tokens,
                                      int64_t       n_seqs,
@@ -25,6 +26,7 @@ gated_delta_net_cuda(const float * q,
                                      const uint3   neqk1_magic,
                                      const uint3   rq3_magic,
                                      float         scale,
+                                     int64_t       state_slot_stride,
                                      int           K) {
     const uint32_t h_idx    = blockIdx.x;
     const uint32_t sequence = blockIdx.y;
@@ -35,9 +37,7 @@ gated_delta_net_cuda(const float * q,
     const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
     const uint32_t iq3 = fastdiv(sequence, rq3_magic);
 
-    const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
     float *       attn_data        = dst;
-    float *       state            = dst + attn_score_elems;
 
     // input state holds s0 only: [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v.
     // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
@@ -145,10 +145,9 @@ gated_delta_net_cuda(const float * q,
         if constexpr (keep_rs_t) {
             // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
             // When n_tokens < K only slots 0..n_tokens-1 are written; older slots are caller-owned.
-            const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
             const int target_slot = (int) n_tokens - 1 - t;
             if (target_slot >= 0 && target_slot < K) {
-                float * curr_state = (dst + attn_score_elems) + target_slot * state_size_per_token + state_out_offset;
+                float * curr_state = state + target_slot * state_slot_stride;
 #pragma unroll
                 for (int r = 0; r < rows_per_lane; r++) {
                     const int i = r * warp_size + lane;
@@ -171,13 +170,13 @@ template <bool KDA, bool keep_rs_t>
 static void launch_gated_delta_net(
         const float * q_d, const float * k_d, const float * v_d,
         const float * g_d, const float * b_d, const float * s_d,
-        float * dst_d,
+        float * dst_d, float * state_d,
         int64_t S_v,   int64_t H, int64_t n_tokens, int64_t n_seqs,
         int64_t sq1,   int64_t sq2, int64_t sq3,
         int64_t sv1,   int64_t sv2, int64_t sv3,
         int64_t sb1,   int64_t sb2, int64_t sb3,
         int64_t neqk1, int64_t rq3,
-        float scale, int K, cudaStream_t stream) {
+        float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
     //TODO: Add chunked kernel for even faster pre-fill
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
     const int num_warps = 4;
@@ -187,34 +186,32 @@ static void launch_gated_delta_net(
     const uint3 neqk1_magic = init_fastdiv_values(neqk1);
     const uint3 rq3_magic   = init_fastdiv_values(rq3);
 
-    int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
-
     const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream);
     switch (S_v) {
         case 16:
             ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t>, launch_params,
-                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
             break;
         case 32:
             ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params,
-                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
             break;
         case 64: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params,
-                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
             break;
         }
         case 128: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t>, launch_params,
-                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
+                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
             break;
         }
         default:
@@ -223,7 +220,8 @@ static void launch_gated_delta_net(
     }
 }
 
-void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+static void ggml_cuda_op_gated_delta_net_impl(
+        ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_cuda_gated_delta_net_fused_cache * cache) {
     ggml_tensor * src_q     = dst->src[0];
     ggml_tensor * src_k     = dst->src[1];
     ggml_tensor * src_v     = dst->src[2];
@@ -288,25 +286,42 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor *
     const int K = ggml_get_op_params_i32(dst, 0);
     const bool keep_rs = K > 1;
 
+    // recurrent state -> gdn_out tail (after attention scores), or the cache when fusing
+    float * state_d           = dst_d + S_v * H * n_tokens * n_seqs;
+    int64_t state_slot_stride = S_v * S_v * H * n_seqs;
+    if (cache != nullptr) {
+        state_d           = cache->data;
+        state_slot_stride = cache->slot_stride;
+    }
+
     if (kda) {
         if (keep_rs) {
-            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     } else {
         if (keep_rs) {
-            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     }
 }
+
+void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    ggml_cuda_op_gated_delta_net_impl(ctx, dst, nullptr);
+}
+
+void ggml_cuda_op_gated_delta_net_fused_cache(
+        ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_cuda_gated_delta_net_fused_cache cache) {
+    ggml_cuda_op_gated_delta_net_impl(ctx, dst, &cache);
+}

ggml/src/ggml-cuda/gated_delta_net.cuh

@@ -1,4 +1,14 @@
 #include "common.cuh"
 #include "ggml.h"
 
+// fused-kernel recurrent-state output; strides in elements (per-seq stride is always D, set in-kernel)
+struct ggml_cuda_gated_delta_net_fused_cache {
+    float * data;        // rollback slot 0
+    int64_t slot_stride; // between rollback slots (0 when K==1)
+};
+
 void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+// same op, but writes the snapshot(s) into the cache instead of dst (see ggml_cuda_try_gdn_cache_fusion)
+void ggml_cuda_op_gated_delta_net_fused_cache(ggml_backend_cuda_context & ctx, ggml_tensor * dst,
+                                              ggml_cuda_gated_delta_net_fused_cache cache);

ggml/src/ggml-cuda/getrows.cu

@@ -78,26 +78,29 @@ static __global__ void k_get_rows_float(
 
 template<typename grad_t, typename dst_t>
 static __global__ void k_get_rows_back_float(
-        const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst, const int64_t ncols, const int64_t nrows_grad) {
+        const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst,
+        const int64_t ncols, const int64_t nrows_grad, const int64_t nrows_dst) {
     const int col = blockIdx.x*blockDim.x + threadIdx.x;
 
     if (col >= ncols) {
         return;
     }
 
-    const int dst_row = blockIdx.y*blockDim.y + threadIdx.y;
-
-    float sum = 0.0f;
-
     ggml_cuda_pdl_sync();
-    for (int64_t i = 0; i < nrows_grad; ++i) {
-        if (rows[i] != dst_row) {
-            continue;
-        }
-        sum += grad[i*ncols + col];
-    }
 
-    dst[dst_row*ncols + col] = sum;
+    // grid.y is clamped to the CUDA grid limit, so stride over the destination rows
+    for (int64_t dst_row = blockIdx.y; dst_row < nrows_dst; dst_row += gridDim.y) {
+        float sum = 0.0f;
+
+        for (int64_t i = 0; i < nrows_grad; ++i) {
+            if (rows[i] != dst_row) {
+                continue;
+            }
+            sum += grad[i*ncols + col];
+        }
+
+        dst[dst_row*ncols + col] = sum;
+    }
 }
 
 template<int qk, int qr, dequantize_kernel_t dq, typename dst_t>
@@ -302,7 +305,7 @@ void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * d
 
     const dim3 block_dims(CUDA_GET_ROWS_BACK_BLOCK_SIZE, 1, 1);
     const int block_num_x = (ne00 + CUDA_GET_ROWS_BACK_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BACK_BLOCK_SIZE;
-    const dim3 block_nums(block_num_x, ne1, 1);
+    const dim3 block_nums(block_num_x, MIN(ne1, (int64_t)UINT16_MAX), 1);
 
-    k_get_rows_back_float<<<block_nums, block_dims, 0, stream>>>(src0_d, src1_d, dst_d, ne00, ne10);
+    k_get_rows_back_float<<<block_nums, block_dims, 0, stream>>>(src0_d, src1_d, dst_d, ne00, ne10, ne1);
 }

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -543,12 +543,42 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
             // the memory allocation handle is no longer needed after mapping
             CU_CHECK(cuMemRelease(handle));
 
-            // set access
-            CUmemAccessDesc access = {};
-            access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-            access.location.id = device;
-            access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
-            CU_CHECK(cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1));
+            // VMM Bug fix for P2P access if GGML_CUDA_P2P is set, or if NCCL build
+            bool use_peer_access = getenv("GGML_CUDA_P2P") != nullptr;
+#if defined(GGML_USE_NCCL)
+            use_peer_access = true;
+#endif // defined(GGML_USE_NCCL)
+
+            if (use_peer_access) {
+                // NCCL implicitly enables peer access (cudaDeviceEnablePeerAccess), and
+                // GGML_CUDA_P2P enables it explicitly. Unlike cudaMalloc buffers, VMM
+                // allocations do not become peer-accessible from that alone, so access
+                // must be granted explicitly here.
+                std::vector<CUmemAccessDesc> access_descs;
+                const int device_count = ggml_cuda_info().device_count;
+                for (int id = 0; id < device_count; ++id) {
+                    if (id != device) {
+                        int can_access_peer = 0;
+                        CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, device));
+                        if (!can_access_peer) {
+                            continue;
+                        }
+                    }
+                    CUmemAccessDesc access = {};
+                    access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+                    access.location.id = id;
+                    access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+                    access_descs.push_back(access);
+                }
+                CU_CHECK(cuMemSetAccess(start_ptr, reserve_size, access_descs.data(), access_descs.size()));
+            } else {
+                // set access for non P2P
+                CUmemAccessDesc access = {};
+                access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+                access.location.id = device;
+                access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+                CU_CHECK(cuMemSetAccess(start_ptr, reserve_size, &access, 1));
+            }
 
             // add to the pool
             pool_size += reserve_size;
@@ -3263,6 +3293,11 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) {
     GGML_UNUSED(backend);
 }
 
+static bool ggml_cuda_is_view_or_noop(const ggml_tensor * t) {
+    return ggml_is_empty(t) || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE ||
+           t->op == GGML_OP_VIEW || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_NONE;
+}
+
 #ifdef USE_CUDA_GRAPH
 static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
 
@@ -3272,7 +3307,7 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
     for (int i = 0; i < cgraph->n_nodes; i++) {
         ggml_tensor * node = cgraph->nodes[i];
 
-        if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) {
+        if (ggml_cuda_is_view_or_noop(node)) {
             continue;
         }
 
@@ -3419,6 +3454,70 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
     return true;
 }
 
+// match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
+// (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
+static int ggml_cuda_try_gdn_cache_fusion(
+        const ggml_cgraph * cgraph, int node_idx, ggml_cuda_gated_delta_net_fused_cache & fused_state_cpy) {
+    const ggml_tensor * gdn = cgraph->nodes[node_idx];
+    // the kernel skips the snapshot tail, so the gdn output must not be a graph output
+    if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->type != GGML_TYPE_F32 ||
+        (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return 0;
+    }
+
+    const ggml_tensor * src_v     = gdn->src[2];
+    const int64_t       S_v       = src_v->ne[0];
+    const int64_t       H         = src_v->ne[1];
+    const int64_t       n_tokens  = src_v->ne[2];
+    const int64_t       n_seqs    = src_v->ne[3];
+    const int64_t       D         = S_v * S_v * H;
+    const int64_t       K         = ggml_get_op_params_i32(gdn, 0); // snapshot slot count
+    const int64_t       n_written = std::min<int64_t>(n_tokens, K); // newest n_written slots are written
+
+    // snapshot tail starts right after the attention scores
+    const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
+
+    // snapshot cpy is the first real node after the gdn (skip views/no-ops)
+    const ggml_tensor * cpy  = nullptr;
+    int                 skip = 0;
+    for (int j = node_idx + 1; j < cgraph->n_nodes && cpy == nullptr; ++j) {
+        const ggml_tensor * n = cgraph->nodes[j];
+        if (ggml_cuda_is_view_or_noop(n)) {
+            continue;
+        }
+        if (n->op != GGML_OP_CPY || (n->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+            return 0;
+        }
+        cpy  = n;
+        skip = j - node_idx;
+    }
+    if (cpy == nullptr) {
+        return 0;
+    }
+
+    const ggml_tensor * src = cpy->src[0]; // view of the gdn snapshot tail
+    const ggml_tensor * dst = cpy->src[1]; // cache view the kernel writes to
+
+    // src must be this gdn's snapshot tail (contiguous, at the tail offset)
+    if (src->op != GGML_OP_VIEW || src->view_src != gdn || src->view_offs != tail_off ||
+        !ggml_is_contiguous(src)) {
+        return 0;
+    }
+
+    // dst is the [D, n_seqs, n_written] cache view; require nb[1] == D (the per-seq stride the kernel
+    // assumes). ggml_cpy pins src to the same element count.
+    const std::array<int64_t, GGML_MAX_DIMS> expected_ne = { D, n_seqs, n_written, 1 };
+    if (dst->op != GGML_OP_VIEW || dst->type != GGML_TYPE_F32 || dst->data == nullptr ||
+        !std::equal(expected_ne.begin(), expected_ne.end(), dst->ne) ||
+        dst->nb[0] != ggml_type_size(GGML_TYPE_F32) || dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) {
+        return 0;
+    }
+
+    fused_state_cpy.data        = (float *) dst->data; // rollback group 0 (newest)
+    fused_state_cpy.slot_stride = K > 1 ? (int64_t) (dst->nb[2] / sizeof(float)) : 0;
+    return skip;
+}
+
 static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
     args.sigmoid         = false;
     args.softmax         = false;
@@ -3860,6 +3959,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
 
     ggml_tensor * node = cgraph->nodes[i];
 
+    // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
+    if (node->op == GGML_OP_GATED_DELTA_NET) {
+        ggml_cuda_gated_delta_net_fused_cache fused_state_cpy;
+        const int nodes_to_skip = ggml_cuda_try_gdn_cache_fusion(cgraph, i, fused_state_cpy);
+        if (nodes_to_skip > 0) {
+#ifdef GGML_CUDA_DEBUG
+            GGML_LOG_INFO("%s: fused gated_delta_net snapshot copies for %s (skipped %d nodes)\n",
+                          __func__, node->name, nodes_to_skip);
+#endif
+            ggml_cuda_op_gated_delta_net_fused_cache(*cuda_ctx, node, fused_state_cpy);
+            return nodes_to_skip;
+        }
+    }
+
     //topk-moe
     if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX ||
             cgraph->nodes[i]->op == GGML_OP_ARGSORT) {
@@ -4388,7 +4501,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
 #endif
                 prev_i = i;
 
-                if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) {
+                if (ggml_cuda_is_view_or_noop(node)) {
                     continue;
                 }
 
@@ -5325,12 +5438,24 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 ggml_type src1_type = op->src[1]->type;
                 return src0_type == src1_type &&
                        src0_type == op->type &&
-                       !ggml_is_quantized(src0_type) &&
-                       ggml_blck_size(src0_type) == 1 &&
-                       (ggml_type_size(src0_type) == 1 ||
-                        ggml_type_size(src0_type) == 2 ||
-                        ggml_type_size(src0_type) == 4 ||
-                        ggml_type_size(src0_type) == 8);
+                       (
+                           (
+                               ggml_is_quantized(src0_type) &&
+                               ggml_is_contiguous(op->src[0]) &&
+                               ggml_is_contiguous(op->src[1]) &&
+                               op->src[0]->ne[0] % ggml_blck_size(src0_type) == 0 &&
+                               op->src[1]->ne[0] % ggml_blck_size(src0_type) == 0
+                           ) || (
+                               !ggml_is_quantized(src0_type) &&
+                               ggml_blck_size(src0_type) == 1 &&
+                               (
+                                   ggml_type_size(src0_type) == 1 ||
+                                   ggml_type_size(src0_type) == 2 ||
+                                   ggml_type_size(src0_type) == 4 ||
+                                   ggml_type_size(src0_type) == 8
+                               )
+                           )
+                       );
             } break;
         case GGML_OP_CONV_TRANSPOSE_1D:
             {

ggml/src/ggml-cuda/mmq.cu

@@ -370,5 +370,12 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
         return true;
     }
 
+    // gfx900 (Vega 10) lacks native dp4a, loses to dequant + hipBLAS
+    // for dense matrices; keep MMQ only for MoE, where the
+    // hipBLAS path is much slower.
+    if (cc == GGML_CUDA_CC_VEGA) {
+        return n_experts > 0;
+    }
+
     return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
 }

ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu

@@ -8,3 +8,4 @@ DECL_FATTN_MMA_F16_CASE(96, 96, 16, 2);
 DECL_FATTN_MMA_F16_CASE(112, 112, 16, 2);
 DECL_FATTN_MMA_F16_CASE(128, 128, 16, 2);
 DECL_FATTN_MMA_F16_CASE(256, 256, 16, 2);
+DECL_FATTN_MMA_F16_CASE(512, 512, 16, 2);

ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu

@@ -8,3 +8,4 @@ DECL_FATTN_MMA_F16_CASE(96, 96, 32, 2);
 DECL_FATTN_MMA_F16_CASE(112, 112, 32, 2);
 DECL_FATTN_MMA_F16_CASE(128, 128, 32, 2);
 DECL_FATTN_MMA_F16_CASE(256, 256, 32, 2);
+DECL_FATTN_MMA_F16_CASE(512, 512, 32, 2);

ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu

@@ -8,3 +8,4 @@ DECL_FATTN_MMA_F16_CASE(96, 96, 4, 2);
 DECL_FATTN_MMA_F16_CASE(112, 112, 4, 2);
 DECL_FATTN_MMA_F16_CASE(128, 128, 4, 2);
 DECL_FATTN_MMA_F16_CASE(256, 256, 4, 2);
+DECL_FATTN_MMA_F16_CASE(512, 512, 4, 2);

ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu

@@ -8,3 +8,4 @@ DECL_FATTN_MMA_F16_CASE(96, 96, 8, 2);
 DECL_FATTN_MMA_F16_CASE(112, 112, 8, 2);
 DECL_FATTN_MMA_F16_CASE(128, 128, 8, 2);
 DECL_FATTN_MMA_F16_CASE(256, 256, 8, 2);
+DECL_FATTN_MMA_F16_CASE(512, 512, 8, 2);

ggml/src/ggml-cuda/template-instances/generate_cu_files.py

@@ -92,7 +92,7 @@ for ncols in [8, 16, 32, 64]:
                     continue
                 if head_size_kq == 320 and ncols2 != 32: # Mistral Small 4
                     continue
-                if head_size_kq == 512 and ncols2 not in (4, 8): # Gemma 4
+                if head_size_kq == 512 and ncols2 not in (2, 4, 8): # Gemma 4 (+ MTP)
                     continue
                 if head_size_kq == 576 and ncols2 not in (4, 16, 32): # Deepseek, GLM 4.7 Flash
                     continue

ggml/src/ggml-cuda/topk-moe.cu

@@ -312,6 +312,10 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
             ggml_cuda_kernel_launch(topk_moe_cuda<256, has_bias>, launch_params,
                 logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
             break;
+        case 288: // StepFun 3.7
+            ggml_cuda_kernel_launch(topk_moe_cuda<288, has_bias>, launch_params,
+                logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
+            break;
         case 512:
             ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias>, launch_params,
                 logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
@@ -377,8 +381,10 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
                                    const ggml_tensor * weights,
                                    const ggml_tensor * logits,
                                    const ggml_tensor * ids) {
+    // must match an instantiation of launch_topk_moe_cuda: a power of 2 up to 512,
+    // or one of the non-power-of-2 expert counts of supported models
     const int n_expert = ids->nb[1] / ids->nb[0];
-    if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 576) {
+    if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 288 && n_expert != 576) {
         return false;
     }
 

ggml/src/ggml-hexagon/htp-opnode.h

@@ -11,6 +11,7 @@
 #include <stdio.h>
 #include "htp-ops.h"
 #include "htp/matmul-ops.h"
+#include "htp/flash-attn-ops.h"
 
 struct htp_opnode {
     ggml_tensor * node = nullptr;
@@ -335,7 +336,8 @@ struct htp_opformat {
     }
     void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
         if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN) {
+            node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN ||
+            node.opcode == HTP_OP_MUL_MAT_ADD) {
             const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
             const char * path = "unknown";
             int32_t type = kparams->kernel_type;
@@ -350,6 +352,16 @@ struct htp_opformat {
                 path = "hvx-flat";
             }
             snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+            const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
+            const char * path = "unknown";
+            int32_t type = kparams->kernel_type;
+            if (type == HTP_FA_KERNEL_HMX) {
+                path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+            } else if (type == HTP_FA_KERNEL_HVX) {
+                path = "hvx";
+            }
+            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
         } else {
             snprintf(str, max_size, "----");
         }

ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h

@@ -712,7 +712,17 @@ static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
 
 // output : fp16 -> f32p
 
-static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 *restrict vtcm_src, uint32_t start_row, uint32_t n_rows, uint32_t n_cols, uint32_t dst_stride, uint32_t dst_cols) {
+static void transfer_output_chunk_fp16_to_fp32(
+    float *restrict dst,
+    const float *restrict src2,
+    const __fp16 *restrict vtcm_src,
+    uint32_t start_row,
+    uint32_t n_rows,
+    uint32_t n_cols,
+    uint32_t dst_stride,
+    uint32_t src2_stride,
+    uint32_t dst_cols
+) {
     assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
     const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
 
@@ -727,6 +737,7 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
         const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2;  // index of the row pair within the tile
         const __fp16 *row_base = vtcm_src + r0 * tile_row_stride;
         float *output_row_base = dst + r * dst_stride;  // global memory row base for row r (and r+1)
+        const float *src2_row_base = src2 ? (src2 + r * src2_stride) : NULL;
 
         #pragma unroll(4)
         for (size_t c = 0; c < limit_c_aligned; c += HTP_MM_HMX_TILE_N_COLS) {
@@ -738,9 +749,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
             HVX_Vector *pv_out0 = (HVX_Vector *) (output_row_base + c + 0);
             HVX_Vector *pv_out1 = (HVX_Vector *) (output_row_base + c + dst_stride);
 
-            *pv_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+            HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+            if (src2_row_base) {
+                HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
+                v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
+            }
+            *pv_out0 = v_out0;
+
             if (r + 1 < n_rows) {
-                *pv_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+                HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+                if (src2_row_base) {
+                    HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
+                    v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
+                }
+                *pv_out1 = v_out1;
             }
         }
 
@@ -752,9 +774,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
             HVX_Vector v = ((const HVX_Vector *) tile)[r1];
             HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one);
 
-            hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)));
+            HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
+            if (src2_row_base) {
+                HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
+                v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
+            }
+            hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), v_out0);
+
             if (r + 1 < n_rows) {
-                hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)));
+                HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
+                if (src2_row_base) {
+                    HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
+                    v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
+                }
+                hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), v_out1);
             }
         }
     }
@@ -763,11 +796,13 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
 typedef struct {
     const __fp16  *vtcm_src;
     float         *dst;
+    const float   *src2;
     uint32_t       n_tasks;
     uint32_t       n_tot_chunks;
     uint32_t       n_chunks_per_task;
     uint32_t       n_cols;
     uint32_t       dst_stride;  // DDR row stride
+    uint32_t       src2_stride; // DDR row stride for residual
     uint32_t       dst_cols;    // Actual output columns
     struct htp_thread_trace * traces;
 } output_transfer_task_state_t;

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h

@@ -256,7 +256,7 @@ static inline void quantize_f16_f16_flat_kernel(
 
 // Dot kernels that consume flat (non-tiled) activations
 
-static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -312,10 +312,14 @@ static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const v
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -397,11 +401,19 @@ static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -464,10 +476,14 @@ static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const v
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -561,11 +577,19 @@ static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -620,10 +644,14 @@ static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const v
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -704,11 +732,19 @@ static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -765,10 +801,14 @@ static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -851,11 +891,19 @@ static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -921,10 +969,14 @@ static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const
 
     v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -1019,6 +1071,441 @@ static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float
     v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
     v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
+
+#if __HVX_ARCH__ < 79
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
+#else
+#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
+#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
+#endif
+
+static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+    uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
+    uint32_t nloe = n % VLEN_FP32; // leftover elements
+
+    HVX_Vector rsum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(4)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
+        rsum = HVX_OP_ADD_F32(rsum, prod);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
+        HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
+        rsum = HVX_OP_ADD_F32(rsum, prod);
+    }
+
+    *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
+
+    uint32_t nvec = n / VLEN_FP32;
+    uint32_t nloe = n % VLEN_FP32;
+
+    HVX_Vector rsum0 = Q6_V_vzero();
+    HVX_Vector rsum1 = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector y_sf = y[i];
+        HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
+        HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
+        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector y_sf  = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
+        HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
+        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
+        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
+    }
+
+    HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+    hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0, const void * restrict vy1) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+    uint32_t nvec = n / VLEN_FP32;
+    uint32_t nloe = n % VLEN_FP32;
+
+    HVX_Vector r0_c0_sum = Q6_V_vzero();
+    HVX_Vector r0_c1_sum = Q6_V_vzero();
+    HVX_Vector r1_c0_sum = Q6_V_vzero();
+    HVX_Vector r1_c1_sum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector r0_sf = x0[i];
+        HVX_Vector r1_sf = x1[i];
+        HVX_Vector c0_sf = y0[i];
+        HVX_Vector c1_sf = y1[i];
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+
+        HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
+        HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
+
+        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
+        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
+        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
+        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
+    }
+
+    // Reduce and store results
+    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+    hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
+    hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
+}
+
+static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
+    const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
+    const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
+
+    uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
+    uint32_t nloe = n % VLEN_FP32; // leftover elements
+
+    HVX_Vector       rsum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector x_sf = vx[i];
+        HVX_Vector y_sf = vy[i];
+
+        rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
+    }
+
+    if (nloe) {
+        HVX_Vector x_sf = vx[i];
+        HVX_Vector y_sf = vy[i];
+
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        x_sf = Q6_V_vand_QV(bmask, x_sf);
+        y_sf = Q6_V_vand_QV(bmask, y_sf);
+
+        rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
+    }
+
+    rsum = hvx_vec_reduce_sum_f32(rsum);
+    hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+#undef HVX_OP_ADD_F32
+#undef HVX_OP_MUL_F32
+
+static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
+    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
+
+    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+    HVX_VectorPair rsum_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(4)
+    for (i = 0; i < nvec; i++) {
+        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
+        HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
+    }
+
+    HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
+    hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
+}
+
+static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
+
+    uint32_t nvec = n / VLEN_FP16;
+    uint32_t nloe = n % VLEN_FP16;
+
+    HVX_VectorPair rsum0_p = Q6_W_vzero();
+    HVX_VectorPair rsum1_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector y_hf = y[i];
+        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
+        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        HVX_Vector y_hf  = Q6_V_vand_QV(bmask, y[i]);
+        HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
+        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
+        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
+    }
+
+    HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
+    HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
+    HVX_Vector rsum  = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
+    hvx_vec_store_u(s0, 8, rsum);
+}
+
+static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
+                                const void * restrict vx0, const void * restrict vx1,
+                                const void * restrict vy0, const void * restrict vy1) {
+    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
+    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
+    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
+    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
+
+    uint32_t nvec = n / VLEN_FP16;
+    uint32_t nloe = n % VLEN_FP16;
+
+    // Row sums (sf) - 4 accumulators for 2x2 tile
+    HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
+    HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
+    HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
+    HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        HVX_Vector r0_hf = x0[i];
+        HVX_Vector r1_hf = x1[i];
+        HVX_Vector c0_hf = y0[i];
+        HVX_Vector c1_hf = y1[i];
+
+        // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1
+        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+
+        HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
+        HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
+        HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
+        HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
+
+        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
+        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
+        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
+        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
+    }
+
+    HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
+    HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
+    HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
+    HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
+
+    // Reduce and store results
+    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
+    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
+
+    hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum);  // row0,col0 row1,col0
+    hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum);  // row0,col1 row1,col1
+}
+
+static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
+    const HVX_UVector * restrict x = (const HVX_UVector *) vx;
+    const HVX_UVector * restrict y = (const HVX_UVector *) vy;
+
+    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+    HVX_Vector rsum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(4)
+    for (i = 0; i < nvec; i++) {
+        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]);
+        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
+    }
+
+    if (nloe) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
+        HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
+
+        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
+    }
+
+    rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
+    hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
+    const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
+    const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
+
+    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
+    uint32_t nloe = n % VLEN_FP16; // leftover elements
+
+    const HVX_Vector zero = Q6_V_vzero();
+
+    HVX_Vector       rsum = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    #pragma unroll(2)
+    for (i = 0; i < nvec; i++) {
+        // Load y (fp32) and convert into fp16
+        HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero);  // 32 elements
+        HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero);  // 32 elements
+        HVX_Vector y_hf  = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
+
+        // Load x (fp16)
+        HVX_Vector x_hf  = vx[i];
+
+        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+
+        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
+    }
+
+    if (nloe) {
+        // Load y (fp32) and convert into fp16
+        HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero);  // 32 elements
+        HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero);  // 32 elements
+        HVX_Vector y_hf  = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
+
+        // Load x (fp16)
+        HVX_Vector x_hf  = vx[i];
+
+        // Zero-out unused elements
+        // Note that we need to clear both x and y because they may contain NANs
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
+        x_hf = Q6_V_vand_QV(bmask, x_hf);
+        y_hf = Q6_V_vand_QV(bmask, y_hf);
+
+        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
+
+        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
+    }
+
+    // Convert into fp32 and reduce
+    rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
+    hvx_vec_store_u(&s[0], 4, rsum);
+}
+
+static inline void hvx_tensor_add_f32_grid(
+    const struct htp_tensor * restrict dst,
+    const struct htp_tensor * restrict src2,
+    uint32_t start_row,
+    uint32_t end_row,
+    uint32_t start_col,
+    uint32_t end_col,
+    const struct fastdiv_values * div_ne11_12,
+    const struct fastdiv_values * div_ne11
+) {
+    if (start_row >= end_row || start_col >= end_col) return;
+    const uint32_t nb1 = dst->nb[1]; // row stride in bytes
+
+    const uint32_t ne11 = dst->ne[1];
+    const uint32_t ne12 = dst->ne[2];
+    const uint32_t ne11_12 = ne11 * ne12;
+
+    const bool is_broadcast1 = (src2->ne[1] == 1);
+    const bool is_broadcast2 = (src2->ne[2] == 1);
+    const bool is_broadcast3 = (src2->ne[3] == 1);
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1);
+
+        uint32_t i13 = fastdiv(r, div_ne11_12);
+        uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11);
+        uint32_t i11 = r - i13 * ne11_12 - i12 * ne11;
+
+        uint32_t i23 = is_broadcast3 ? 0 : i13;
+        uint32_t i22 = is_broadcast2 ? 0 : i12;
+        uint32_t i21 = is_broadcast1 ? 0 : i11;
+
+        const float * src2_row = (const float *) ((const uint8_t *) src2->data +
+                                  i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]);
+
+        float * dst_ptr = &dst_row[start_col];
+        const float * src2_ptr = &src2_row[start_col];
+        int remaining = end_col - start_col;
+        while (remaining >= 32) {
+            HVX_Vector v_out = hvx_vmemu(dst_ptr);
+            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
+            hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z);
+            dst_ptr += 32;
+            src2_ptr += 32;
+            remaining -= 32;
+        }
+        if (remaining > 0) {
+            HVX_Vector v_out = hvx_vmemu(dst_ptr);
+            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
+            hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
+        }
+    }
+}
+

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h

@@ -378,7 +378,7 @@ static inline HVX_VectorPair accum_q8_0_32x2(
     return Q6_W_vcombine_VV(v_sum1, v_sum0);
 }
 
-static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -401,10 +401,14 @@ static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -484,11 +488,19 @@ static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -519,10 +531,14 @@ static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -637,11 +653,19 @@ static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -663,10 +687,14 @@ static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -745,11 +773,19 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -773,10 +809,14 @@ static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const
         v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
     }
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -857,11 +897,19 @@ static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, floa
         v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
     }
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
-static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
+static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y_q = vy;
 
@@ -896,10 +944,14 @@ static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const
 
     v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
 
-    hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    if (sz) {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
+    } else {
+        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
+    }
 }
 
-static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
+static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
     const uint8_t * restrict tile_ptr = vx;
     const uint8_t * restrict y0_q = vy0;
     const uint8_t * restrict y1_q = vy1;
@@ -1013,8 +1065,16 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa
     v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
     v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
 
-    hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
-    hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    if (sz0) {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
+    } else {
+        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
+    }
+    if (sz1) {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
+    } else {
+        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
+    }
 }
 
 static inline void quantize_f32_q8_0_tiled_kernel(

ggml/src/ggml-hexagon/htp/matmul-ops.h

@@ -392,56 +392,49 @@ static inline size_t htp_mm_hvx_get_vtcm_sizes(
         case HTP_MM_KERNEL_HVX_QUANT_ROW: {
             size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
-            vtcm_dst_size  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
             vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
             vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
 
-            // src0 spad is also used in dynamic quantizer to store padded src1 rows
-            size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, QK_Q8_0_TILED * sizeof(float));
-            if (vtcm_src0_size < src1_row_size_padded) {
-                vtcm_src0_size = src1_row_size_padded;
-            }
-
             vtcm_src0_size = vtcm_src0_size * n_threads;
-            vtcm_dst_size  = vtcm_dst_size * n_threads;
 
             if (is_repack) {
                 uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
                 uint32_t n_k_tiles = ne10 / 32;
                 uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
                 size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                if (repacked_vtcm_size < src1_row_size_padded) {
-                    repacked_vtcm_size = src1_row_size_padded;
-                }
                 vtcm_src0_size = repacked_vtcm_size * n_threads;
             }
+
+            size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+            size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+            if (dst_size_per_thread < quant_scratch_size_per_thread) {
+                dst_size_per_thread = quant_scratch_size_per_thread;
+            }
+            vtcm_dst_size = dst_size_per_thread * n_threads;
             break;
         }
         case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
             size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
 
-            vtcm_dst_size  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
             vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
             vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
 
-            size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, 256);
-            if (vtcm_src0_size < src1_row_size_padded) {
-                vtcm_src0_size = src1_row_size_padded;
-            }
-
             vtcm_src0_size = vtcm_src0_size * n_threads;
-            vtcm_dst_size  = vtcm_dst_size * n_threads;
 
             if (is_repack) {
                 uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
                 uint32_t n_k_tiles = ne10 / 32;
                 uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
                 size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                if (repacked_vtcm_size < src1_row_size_padded) {
-                    repacked_vtcm_size = src1_row_size_padded;
-                }
                 vtcm_src0_size = repacked_vtcm_size * n_threads;
             }
+
+            size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+            size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+            if (dst_size_per_thread < quant_scratch_size_per_thread) {
+                dst_size_per_thread = quant_scratch_size_per_thread;
+            }
+            vtcm_dst_size = dst_size_per_thread * n_threads;
             break;
         }
         default:
@@ -463,7 +456,8 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
     size_t src0_row_size,    // nb01
     uint32_t n_prefetch,
     size_t * vtcm_src0_size_out,
-    size_t * vtcm_src1_size_out
+    size_t * vtcm_src1_size_out,
+    size_t * vtcm_dst_size_out
 ) {
     const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
                             wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
@@ -476,29 +470,22 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
     size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
     size_t src1_sz            = htp_mm_round_up(src1_row_size * src1_nrows, 256);
 
-    // src0 spad also holds temporary transposed src1 columns during dynamic quantization.
-    const size_t src1_row_size_padded = htp_mm_round_up(src1_row_size, QK_Q8_0_TILED * sizeof(float));
-    if (src0_sz_per_thread < src1_row_size_padded) {
-        src0_sz_per_thread = src1_row_size_padded;
-    }
-
     if (is_repack) {
         const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
         const uint32_t n_k_tiles    = ne10 / 32;
         const uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
         size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-        if (repacked_vtcm_size < src1_row_size_padded) {
-            repacked_vtcm_size = src1_row_size_padded;
-        }
         src0_sz_per_thread = repacked_vtcm_size;
     }
 
     const size_t vtcm_src0_size = src0_sz_per_thread * n_threads;
+    const size_t vtcm_dst_size  = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
 
     *vtcm_src0_size_out = vtcm_src0_size;
     *vtcm_src1_size_out = src1_sz;
+    *vtcm_dst_size_out  = vtcm_dst_size;
 
-    return vtcm_src0_size + src1_sz;
+    return vtcm_src0_size + src1_sz + vtcm_dst_size;
 }
 
 #ifdef __cplusplus

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -135,7 +135,7 @@ typedef struct VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE {
 #endif
 
 #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
-#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+#define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
 static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
 
 #define VK_VENDOR_ID_AMD 0x1002
@@ -1913,6 +1913,38 @@ static bool vk_enable_sync_logger = false;
 static uint32_t vk_perf_logger_frequency = 1;
 static std::string vk_pipeline_stats_filter;
 
+static uint64_t ggml_vk_get_node_flops(const ggml_tensor * node) {
+    if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) {
+        const uint64_t m     = node->ne[0];
+        const uint64_t n     = node->ne[1];
+        const uint64_t k     = node->src[1]->ne[0];
+        const uint64_t batch = node->ne[2] * node->ne[3];
+        return m * n * (k + (k - 1)) * batch;
+    }
+    if (node->op == GGML_OP_CONV_2D || node->op == GGML_OP_CONV_TRANSPOSE_2D) {
+        const ggml_tensor * knl = node->src[0];
+        const uint64_t Cout  = node->ne[2];
+        const uint64_t size_K = node->src[1]->ne[2] * knl->ne[0] * knl->ne[1];
+        const uint64_t size_N = node->ne[3] * node->ne[0] * node->ne[1];
+        return Cout * size_N * (size_K + (size_K - 1));
+    }
+    if (node->op == GGML_OP_CONV_3D) {
+        const ggml_tensor * knl = node->src[0];
+        const uint64_t OC     = ggml_get_op_params_i32(node, 11);
+        const uint64_t IC     = ggml_get_op_params_i32(node, 9);
+        const uint64_t size_K = IC * knl->ne[0] * knl->ne[1] * knl->ne[2];
+        const uint64_t size_N = node->ne[3] / OC * node->ne[0] * node->ne[1] * node->ne[2];
+        return OC * size_N * (size_K + (size_K - 1));
+    }
+    if (node->op == GGML_OP_FLASH_ATTN_EXT) {
+        const ggml_tensor * q = node->src[0];
+        const ggml_tensor * k = node->src[1];
+        const ggml_tensor * v = node->src[2];
+        return 2ull * q->ne[1] * q->ne[2] * (k->ne[0] + v->ne[0]) * k->ne[1] * q->ne[3];
+    }
+    return 0;
+}
+
 class vk_perf_logger {
   public:
     void print_timings(bool force = false) {
@@ -1961,7 +1993,7 @@ class vk_perf_logger {
     }
 
     std::string get_node_fusion_name(const ggml_tensor * node, const char *fusion_name, uint64_t *n_flops) {
-        *n_flops = 0;
+        *n_flops = ggml_vk_get_node_flops(node);
         std::string fusion_str;
         if (fusion_name) {
             fusion_str = fusion_name + std::string(" ");
@@ -1988,35 +2020,22 @@ class vk_perf_logger {
             if (batch > 1) {
                 name += " batch=" + std::to_string(batch);
             }
-            name = fusion_str + name;
-            *n_flops = m * n * (k + (k - 1)) * batch;
-            return name;
+            return fusion_str + name;
         }
         if (node->op == GGML_OP_CONV_2D || node->op == GGML_OP_CONV_TRANSPOSE_2D) {
             std::string   name    = ggml_op_name(node->op);
-            ggml_tensor * knl     = node->src[0];
-            uint64_t      OW      = node->ne[0];
-            uint64_t      OH      = node->ne[1];
-            uint64_t      N       = node->ne[3];
+            const ggml_tensor * knl = node->src[0];
             uint64_t      Cout    = node->ne[2];
-            uint64_t      KW      = knl->ne[0];
-            uint64_t      KH      = knl->ne[1];
-            uint64_t      Cin     = node->src[1]->ne[2];
-            // KxCRS @ CRSxNPQ = KxNPQ -> M=K, K=CRS, N=NPQ
-            uint64_t      size_M  = Cout;
-            uint64_t      size_K  = Cin * KW * KH;
-            uint64_t      size_N  = N * OW * OH;
-            *n_flops = size_M * size_N * (size_K + (size_K - 1));
-            name += " M=Cout=" + std::to_string(size_M) + ", K=Cin*KW*KH=" + std::to_string(size_K) +
+            uint64_t      size_K  = node->src[1]->ne[2] * knl->ne[0] * knl->ne[1];
+            uint64_t      size_N  = node->ne[3] * node->ne[0] * node->ne[1];
+            name += " M=Cout=" + std::to_string(Cout) + ", K=Cin*KW*KH=" + std::to_string(size_K) +
                     ", N=N*OW*OH=" + std::to_string(size_N);
-            name = fusion_str + name;
-            return name;
+            return fusion_str + name;
         }
         if (node->op == GGML_OP_RMS_NORM) {
             std::string   name    = ggml_op_name(node->op);
             name += "(" + std::to_string(node->ne[0]) + "," + std::to_string(node->ne[1]) + "," + std::to_string(node->ne[2]) + "," + std::to_string(node->ne[3]) + ")";
-            name = fusion_str + name;
-            return name;
+            return fusion_str + name;
         }
         if (node->op == GGML_OP_FLASH_ATTN_EXT) {
             const ggml_tensor * dst = node;
@@ -2032,7 +2051,6 @@ class vk_perf_logger {
                 " k(" << k->ne[0] << "," << k->ne[1] << "," << k->ne[2] << "," << k->ne[3] << "), " <<
                 " v(" << v->ne[0] << "," << v->ne[1] << "," << v->ne[2] << "," << v->ne[3] << "), " <<
                 " m(" << (m?m->ne[0]:0) << "," << (m?m->ne[1]:0) << "," << (m?m->ne[2]:0) << "," << (m?m->ne[3]:0) << ")";
-            *n_flops = 2ull * q->ne[1] * q->ne[2] * (k->ne[0] + v->ne[0]) * k->ne[1] * q->ne[3];
             return name.str();
         }
         if (node->op == GGML_OP_TOP_K) {
@@ -2096,7 +2114,7 @@ struct ggml_backend_vk_context {
     bool do_add_rms_partials_offset_calculation;
     bool do_add_rms_partials;
 
-    uint64_t last_total_mul_mat_bytes {};
+    uint64_t last_total_flops {UINT64_MAX};
 
     // Cache most recent tensor that was converted into prealloc_y, and what pipeline it used to convert.
     vk_pipeline_struct * prealloc_y_last_pipeline_used {};
@@ -2463,6 +2481,85 @@ static bool ggml_vk_strip_decode_vector(const uint32_t * code, size_t word_count
     return true;
 }
 
+// Remove the loop unrolling hint of the matmul shader's BK loop
+// and replace it with the dont_unroll hint for better performance on
+// hardware like Apple M1/M2.
+// Assumes 1. code comes from mul_mm.comp 2. the K-tile loop has no loop
+// control hint and 3. the BK loop is the last loop nested directly inside
+// the K-tile loop.
+// Returns true when the input was modified; returns false otherwise
+// without touching `out`.
+static bool ggml_vk_roll_bk_loop(const uint32_t * code, size_t word_count, std::vector<uint32_t> & out) {
+    if (word_count < 5) {
+        return false;
+    }
+
+    struct vk_spv_loop {
+        size_t   header;
+        size_t   end;
+        uint32_t control;
+    };
+
+    std::vector<vk_spv_loop> loops;
+
+    // Collect a list of all loops in the module.
+    for (size_t pos = 5; pos < word_count; ) {
+        const uint32_t wc = code[pos] >> spv::WordCountShift;
+        const uint32_t op = code[pos] & spv::OpCodeMask;
+        if (wc == 0 || pos + wc > word_count) {
+            return false;
+        }
+
+        if (op == spv::OpLoopMerge && wc >= 4) { loops.push_back({ pos, 0, code[pos + 3] }); }
+
+        if (op == spv::OpLabel && wc >= 2) {
+            for (auto & l : loops) {
+                if (l.end == 0 && code[l.header + 1] == code[pos + 1]) { l.end = pos; }
+            }
+        }
+
+        pos += wc;
+    }
+
+    auto encloses = [](const vk_spv_loop & a, const vk_spv_loop & b) {
+        return a.header < b.header && b.header < a.end;
+    };
+
+    // Find the BK loop.
+    const vk_spv_loop * bk = nullptr;
+    for (const auto & h : loops) {
+        if (h.control != spv::LoopControlUnrollMask) {
+            continue;
+        }
+        const vk_spv_loop * parent = nullptr;
+        bool has_child = false;
+        for (const auto & g : loops) {
+            if (encloses(g, h) && (!parent || g.header > parent->header)) {
+                parent = &g;
+            }
+            if (encloses(h, g)) {
+                has_child = true;
+            }
+        }
+        // BK loop should be the last loop nested inside the loop with no hint
+        // and have at least one child loop.
+        if (parent &&
+            parent->control == spv::LoopControlMaskNone &&
+            has_child &&
+            (!bk || h.header > bk->header)) {
+            bk = &h;
+        }
+    }
+    if (!bk) {
+        return false;
+    }
+
+    // set DontUnroll instead of Unroll
+    out.assign(code, code + word_count);
+    out[bk->header + 3] = spv::LoopControlDontUnrollMask;
+    return true;
+}
+
 static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipeline, size_t spv_size, const void* spv_data, const std::string entrypoint,
                                          uint32_t parameter_count, std::array<uint32_t, 3> wg_denoms, std::vector<uint32_t> specialization_constants,
                                          bool disable_robustness, bool require_full_subgroups, uint32_t required_subgroup_size) {
@@ -2546,6 +2643,22 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin
     }
 #endif
 
+#if VK_HEADER_VERSION >= 287
+    // Roll the mul_mm BK loop on Asahi Linux. Skip bf16 and the mul_mmq pipelines.
+    if (device->driver_id == vk::DriverId::eMesaHoneykrisp &&
+        pipeline->name.rfind("matmul", 0) == 0 &&
+        pipeline->name.find("bf16") == std::string::npos &&
+        pipeline->name.find("q8_1") == std::string::npos) {
+        const uint32_t * src   = spirv.empty() ? reinterpret_cast<const uint32_t *>(spv_data) : spirv.data();
+        size_t           src_n = spirv.empty() ? spv_size / sizeof(uint32_t) : spirv.size();
+        std::vector<uint32_t> rolled;
+        if (ggml_vk_roll_bk_loop(src, src_n, rolled)) {
+            spirv = std::move(rolled);
+            shader_module_create_info = vk::ShaderModuleCreateInfo({}, spirv.size() * sizeof(uint32_t), spirv.data());
+        }
+    }
+#endif
+
     pipeline->shader_module = device->device.createShaderModule(shader_module_create_info);
 
     vk::PushConstantRange pcr(
@@ -5908,7 +6021,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         // Submit at least every 100 nodes, in case there are workloads without as much matmul.
         // device->max_nodes_per_submit = 100;
-        device->max_nodes_per_submit = device->uma ? 8 : 64; //kcpp fix for https://github.com/ggml-org/llama.cpp/issues/21724
+        device->max_nodes_per_submit = device->uma ? 4 : 16; //kcpp fix for https://github.com/ggml-org/llama.cpp/issues/21724
         const char* GGML_VK_MAX_NODES_PER_SUBMIT = getenv("GGML_VK_MAX_NODES_PER_SUBMIT");
         if (GGML_VK_MAX_NODES_PER_SUBMIT != nullptr) {
             uint32_t max_nodes_per_submit = std::stoul(GGML_VK_MAX_NODES_PER_SUBMIT);
@@ -16221,22 +16334,23 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
     }
 
     // Submit after enough work has accumulated, to overlap CPU cmdbuffer generation with GPU execution.
-    // Estimate the amount of matmul work by looking at the weight matrix size, and submit every 100MB
-    // (and scaled down based on model size, so smaller models submit earlier).
-    int submitted_nodes = 0;
-    int submit_count = 0;
-    uint64_t mul_mat_bytes = 0;
-    uint64_t total_mul_mat_bytes = 0;
-    uint64_t mul_mat_bytes_per_submit = std::min(uint64_t(100*1000*1000), ctx->last_total_mul_mat_bytes / 40u);
+    // Estimate the amount of compute work using flops, and submit every 200 GFLOP
+    // (and scaled down based on total graph flops, so smaller models submit earlier).
+    // Also submit at least every 100 nodes, in case there are workloads without heavy compute.
+    uint32_t submitted_nodes = 0;
+    uint32_t submit_count = 0;
+    uint64_t batch_flops = 0;
+    uint64_t total_flops = 0;
+    uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u);
     for (int i = 0; i < cgraph->n_nodes; i++) {
         if (first_node_in_batch) {
             submit_node_idx = i;
         }
 
-        if (cgraph->nodes[i]->op == GGML_OP_MUL_MAT || cgraph->nodes[i]->op == GGML_OP_MUL_MAT_ID) {
-            auto bytes = ggml_nbytes(cgraph->nodes[i]->src[0]);
-            mul_mat_bytes += bytes;
-            total_mul_mat_bytes += bytes;
+        {
+            auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]);
+            batch_flops += node_flops;
+            total_flops += node_flops;
         }
 
         // op_srcs_fused_elementwise indicates whether an op's srcs all contribute to
@@ -16448,8 +16562,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
 
         // Signal the almost_ready fence when the graph is mostly complete (< 20% remaining)
         bool almost_ready = (cgraph->n_nodes - i) < cgraph->n_nodes / 5;
-        bool submit = ((uint32_t)submitted_nodes >= ctx->device->max_nodes_per_submit) ||
-                      (mul_mat_bytes_per_submit != 0 && mul_mat_bytes >= mul_mat_bytes_per_submit) ||
+        bool submit = (submitted_nodes >= ctx->device->max_nodes_per_submit) ||
+                      (flops_per_submit != 0 && batch_flops >= flops_per_submit) ||
                       (i + ctx->num_additional_fused_ops >= last_node) ||
                       (almost_ready && !ctx->almost_ready_fence_pending);
 
@@ -16483,9 +16597,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         if (submit && enqueued) {
             first_node_in_batch = true;
             submitted_nodes = 0;
-            mul_mat_bytes = 0;
+            batch_flops = 0;
             if (submit_count < 3) {
-                mul_mat_bytes_per_submit *= 2;
+                flops_per_submit *= 2;
             }
             submit_count++;
         }
@@ -16494,7 +16608,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         ctx->fused_ops_write_mask = 0;
     }
 
-    ctx->last_total_mul_mat_bytes = total_mul_mat_bytes;
+    ctx->last_total_flops = total_flops;
 
     if (vk_perf_logger_enabled) {
         // End the command buffer and submit/wait

ggml/src/ggml-vulkan/vulkan-shaders/unary.comp

@@ -42,7 +42,7 @@ float op_leaky_relu(float x) {
 }
 
 float op_step(float x) {
-    return x >= 0.0f ? 1.0f : 0.0f;
+    return x > 0.0f ? 1.0f : 0.0f;
 }
 
 float op_tanh(float x) {

gguf-py/gguf/constants.py

@@ -145,6 +145,7 @@ class Keys:
         TOKEN_SHIFT_COUNT                 = "{arch}.token_shift_count"
         INTERLEAVE_MOE_LAYER_STEP         = "{arch}.interleave_moe_layer_step"
         FULL_ATTENTION_INTERVAL           = "{arch}.full_attention_interval"
+        HASH_LAYER_COUNT                  = "{arch}.hash_layer_count"
         ACTIVATION_SPARSITY_SCALE         = "{arch}.activation_sparsity_scale"
         ALTUP_ACTIVE_IDX                  = "{arch}.altup.active_idx"
         ALTUP_NUM_INPUTS                  = "{arch}.altup.num_inputs"
@@ -156,6 +157,7 @@ class Keys:
         DENSE_FEAT_OUT_SIZE               = "{arch}.{dense}_feat_out"
         TARGET_LAYERS                     = "{arch}.target_layers"
         TARGET_HIDDEN_SIZE                = "{arch}.target_hidden_size"
+        BLOCK_SIZE                        = "{arch}.block_size"
         NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
 
     class Attention:
@@ -179,8 +181,12 @@ class Keys:
         REL_BUCKETS_COUNT            = "{arch}.attention.relative_buckets_count"
         SLIDING_WINDOW               = "{arch}.attention.sliding_window"
         SCALE                        = "{arch}.attention.scale"
+        OUTPUT_GROUP_COUNT           = "{arch}.attention.output_group_count"
+        OUTPUT_LORA_RANK             = "{arch}.attention.output_lora_rank"
         OUTPUT_SCALE                 = "{arch}.attention.output_scale"
         VALUE_SCALE                  = "{arch}.attention.value_scale"
+        COMPRESS_RATIOS              = "{arch}.attention.compress_ratios"
+        COMPRESS_ROPE_FREQ_BASE      = "{arch}.attention.compress_rope_freq_base"
         TEMPERATURE_LENGTH           = "{arch}.attention.temperature_length"
         KEY_LENGTH_MLA               = "{arch}.attention.key_length_mla"
         VALUE_LENGTH_MLA             = "{arch}.attention.value_length_mla"
@@ -195,6 +201,11 @@ class Keys:
             KEY_LENGTH = "{arch}.attention.indexer.key_length"
             TOP_K      = "{arch}.attention.indexer.top_k"
 
+    class HyperConnection:
+        COUNT                = "{arch}.hyper_connection.count"
+        SINKHORN_ITERATIONS  = "{arch}.hyper_connection.sinkhorn_iterations"
+        EPSILON              = "{arch}.hyper_connection.epsilon"
+
     class Rope:
         DIMENSION_COUNT           = "{arch}.rope.dimension_count"
         DIMENSION_COUNT_SWA       = "{arch}.rope.dimension_count_swa"
@@ -469,6 +480,7 @@ class MODEL_ARCH(IntEnum):
     DEEPSEEK2        = auto()
     DEEPSEEK2OCR     = auto()
     DEEPSEEK32       = auto()
+    DEEPSEEK4        = auto()
     CHATGLM          = auto()
     GLM4             = auto()
     GLM4_MOE         = auto()
@@ -517,6 +529,7 @@ class MODEL_ARCH(IntEnum):
     PANGU_EMBED      = auto()
     MISTRAL3         = auto()
     EAGLE3           = auto()
+    DFLASH           = auto()
     MISTRAL4         = auto()
     PADDLEOCR        = auto()
     MIMO2            = auto()
@@ -553,6 +566,9 @@ class MODEL_TENSOR(IntEnum):
     DENSE_2_OUT          = auto() # embeddinggemma 2_Dense
     DENSE_3_OUT          = auto() # embeddinggemma 3_Dense
     OUTPUT_NORM          = auto()
+    HC_HEAD_FN           = auto()
+    HC_HEAD_BASE         = auto()
+    HC_HEAD_SCALE        = auto()
     ROPE_FREQS           = auto()
     ROPE_FACTORS_LONG    = auto()
     ROPE_FACTORS_SHORT   = auto()
@@ -592,6 +608,7 @@ class MODEL_TENSOR(IntEnum):
     FFN_DOWN_CHEXP       = auto()
     FFN_UP_CHEXP         = auto()
     FFN_EXP_PROBS_B      = auto()
+    FFN_GATE_TID2EID     = auto()
     MOE_LATENT_DOWN      = auto() # nemotron 3 super
     MOE_LATENT_UP        = auto() # nemotron 3 super
     ATTN_Q_NORM          = auto()
@@ -679,6 +696,20 @@ class MODEL_TENSOR(IntEnum):
     ATTN_V_B             = auto()
     ATTN_Q_A_NORM        = auto()
     ATTN_KV_A_NORM       = auto()
+    ATTN_KV              = auto()
+    ATTN_KV_NORM         = auto()
+    ATTN_OUT_A           = auto()
+    ATTN_OUT_B           = auto()
+    HC_ATTN_FN           = auto()
+    HC_ATTN_BASE         = auto()
+    HC_ATTN_SCALE        = auto()
+    HC_FFN_FN            = auto()
+    HC_FFN_BASE          = auto()
+    HC_FFN_SCALE         = auto()
+    ATTN_COMPRESSOR_WKV  = auto()
+    ATTN_COMPRESSOR_WGATE = auto()
+    ATTN_COMPRESSOR_APE  = auto()
+    ATTN_COMPRESSOR_NORM = auto()
     FFN_SUB_NORM         = auto()
     ATTN_SUB_NORM        = auto()
     DEC_ATTN_NORM        = auto()
@@ -740,6 +771,10 @@ class MODEL_TENSOR(IntEnum):
     INDEXER_PROJ         = auto()
     INDEXER_ATTN_K       = auto()
     INDEXER_ATTN_Q_B     = auto()
+    INDEXER_COMPRESSOR_WKV = auto()
+    INDEXER_COMPRESSOR_WGATE = auto()
+    INDEXER_COMPRESSOR_APE = auto()
+    INDEXER_COMPRESSOR_NORM = auto()
     # vision
     V_MMPROJ             = auto()
     V_MMPROJ_FC          = auto()
@@ -1025,6 +1060,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.DEEPSEEK2:        "deepseek2",
     MODEL_ARCH.DEEPSEEK2OCR:     "deepseek2-ocr",
     MODEL_ARCH.DEEPSEEK32:       "deepseek32",
+    MODEL_ARCH.DEEPSEEK4:        "deepseek4",
     MODEL_ARCH.CHATGLM:          "chatglm",
     MODEL_ARCH.GLM4:             "glm4",
     MODEL_ARCH.GLM4_MOE:         "glm4moe",
@@ -1074,6 +1110,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.PANGU_EMBED:      "pangu-embedded",
     MODEL_ARCH.MISTRAL3:         "mistral3",
     MODEL_ARCH.EAGLE3:           "eagle3",
+    MODEL_ARCH.DFLASH:           "dflash",
     MODEL_ARCH.MISTRAL4:         "mistral4",
     MODEL_ARCH.PADDLEOCR:        "paddleocr",
     MODEL_ARCH.MIMO2:            "mimo2",
@@ -1108,6 +1145,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.OUTPUT:                    "output",
     MODEL_TENSOR.DENSE_2_OUT:               "dense_2", # embeddinggemma 2_Dense
     MODEL_TENSOR.DENSE_3_OUT:               "dense_3", # embeddinggemma 2_Dense
+    MODEL_TENSOR.HC_HEAD_FN:                "output_hc_fn",
+    MODEL_TENSOR.HC_HEAD_BASE:              "output_hc_base",
+    MODEL_TENSOR.HC_HEAD_SCALE:             "output_hc_scale",
     MODEL_TENSOR.ROPE_FREQS:                "rope_freqs",
     MODEL_TENSOR.ROPE_FACTORS_LONG:         "rope_factors_long",
     MODEL_TENSOR.ROPE_FACTORS_SHORT:        "rope_factors_short",
@@ -1149,6 +1189,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.FFN_UP_EXP:                "blk.{bid}.ffn_up_exps",
     MODEL_TENSOR.FFN_GATE_UP_EXP:           "blk.{bid}.ffn_gate_up_exps",
     MODEL_TENSOR.FFN_EXP_PROBS_B:           "blk.{bid}.exp_probs_b",
+    MODEL_TENSOR.FFN_GATE_TID2EID:          "blk.{bid}.ffn_gate_tid2eid",
     MODEL_TENSOR.MOE_LATENT_DOWN:           "blk.{bid}.ffn_latent_down",      # nemotron 3 super
     MODEL_TENSOR.MOE_LATENT_UP:             "blk.{bid}.ffn_latent_up",        # nemotron 3 super
     MODEL_TENSOR.LAYER_OUT_NORM:            "blk.{bid}.layer_output_norm",
@@ -1234,6 +1275,20 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.ATTN_V_B:                  "blk.{bid}.attn_v_b",
     MODEL_TENSOR.ATTN_Q_A_NORM:             "blk.{bid}.attn_q_a_norm",
     MODEL_TENSOR.ATTN_KV_A_NORM:            "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_KV:                   "blk.{bid}.attn_kv",
+    MODEL_TENSOR.ATTN_KV_NORM:              "blk.{bid}.attn_kv_a_norm",
+    MODEL_TENSOR.ATTN_OUT_A:                "blk.{bid}.attn_output_a",
+    MODEL_TENSOR.ATTN_OUT_B:                "blk.{bid}.attn_output_b",
+    MODEL_TENSOR.HC_ATTN_FN:                "blk.{bid}.hc_attn_fn",
+    MODEL_TENSOR.HC_ATTN_BASE:              "blk.{bid}.hc_attn_base",
+    MODEL_TENSOR.HC_ATTN_SCALE:             "blk.{bid}.hc_attn_scale",
+    MODEL_TENSOR.HC_FFN_FN:                 "blk.{bid}.hc_ffn_fn",
+    MODEL_TENSOR.HC_FFN_BASE:               "blk.{bid}.hc_ffn_base",
+    MODEL_TENSOR.HC_FFN_SCALE:              "blk.{bid}.hc_ffn_scale",
+    MODEL_TENSOR.ATTN_COMPRESSOR_WKV:       "blk.{bid}.attn_compressor_kv",
+    MODEL_TENSOR.ATTN_COMPRESSOR_WGATE:     "blk.{bid}.attn_compressor_gate",
+    MODEL_TENSOR.ATTN_COMPRESSOR_APE:       "blk.{bid}.attn_compressor_ape",
+    MODEL_TENSOR.ATTN_COMPRESSOR_NORM:      "blk.{bid}.attn_compressor_norm",
     MODEL_TENSOR.ATTN_SUB_NORM:             "blk.{bid}.attn_sub_norm",
     MODEL_TENSOR.FFN_SUB_NORM:              "blk.{bid}.ffn_sub_norm",
     MODEL_TENSOR.DEC_ATTN_NORM:             "dec.blk.{bid}.attn_norm",
@@ -1295,6 +1350,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.INDEXER_PROJ:              "blk.{bid}.indexer.proj",
     MODEL_TENSOR.INDEXER_ATTN_K:            "blk.{bid}.indexer.attn_k",
     MODEL_TENSOR.INDEXER_ATTN_Q_B:          "blk.{bid}.indexer.attn_q_b",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_WKV:    "blk.{bid}.indexer_compressor_kv",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE:  "blk.{bid}.indexer_compressor_gate",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_APE:    "blk.{bid}.indexer_compressor_ape",
+    MODEL_TENSOR.INDEXER_COMPRESSOR_NORM:   "blk.{bid}.indexer_compressor_norm",
     # vision
     MODEL_TENSOR.V_MMPROJ:                  "mm.{bid}",
     MODEL_TENSOR.V_MMPROJ_FC:               "mm.model.fc",
@@ -3135,6 +3194,49 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
         MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
     ],
+    MODEL_ARCH.DEEPSEEK4: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.HC_HEAD_FN,
+        MODEL_TENSOR.HC_HEAD_BASE,
+        MODEL_TENSOR.HC_HEAD_SCALE,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_SINKS,
+        MODEL_TENSOR.ATTN_Q_A,
+        MODEL_TENSOR.ATTN_Q_B,
+        MODEL_TENSOR.ATTN_Q_A_NORM,
+        MODEL_TENSOR.ATTN_KV,
+        MODEL_TENSOR.ATTN_KV_NORM,
+        MODEL_TENSOR.ATTN_OUT_A,
+        MODEL_TENSOR.ATTN_OUT_B,
+        MODEL_TENSOR.HC_ATTN_FN,
+        MODEL_TENSOR.HC_ATTN_BASE,
+        MODEL_TENSOR.HC_ATTN_SCALE,
+        MODEL_TENSOR.HC_FFN_FN,
+        MODEL_TENSOR.HC_FFN_BASE,
+        MODEL_TENSOR.HC_FFN_SCALE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WKV,
+        MODEL_TENSOR.ATTN_COMPRESSOR_WGATE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_APE,
+        MODEL_TENSOR.ATTN_COMPRESSOR_NORM,
+        MODEL_TENSOR.INDEXER_PROJ,
+        MODEL_TENSOR.INDEXER_ATTN_Q_B,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_WKV,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_APE,
+        MODEL_TENSOR.INDEXER_COMPRESSOR_NORM,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_TID2EID,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
     MODEL_ARCH.ERNIE4_5_MOE: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -4086,6 +4188,22 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FC,
         MODEL_TENSOR.D2T,
     ],
+    MODEL_ARCH.DFLASH: [
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FC,
+        MODEL_TENSOR.ENC_OUTPUT_NORM,
+    ],
     MODEL_ARCH.MISTRAL4: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -4418,8 +4536,9 @@ class GGMLQuantizationType(IntEnum):
 
 
 class ExpertGatingFuncType(IntEnum):
-    SOFTMAX  = 1
-    SIGMOID  = 2
+    SOFTMAX       = 1
+    SIGMOID       = 2
+    SQRTSOFTPLUS  = 4
 
 
 # TODO: add GGMLFileType from ggml_ftype in ggml.h

gguf-py/gguf/gguf_writer.py

@@ -715,6 +715,9 @@ class GGUFWriter:
     def add_full_attention_interval(self, interval: int) -> None:
         self.add_uint32(Keys.LLM.FULL_ATTENTION_INTERVAL.format(arch=self.arch), interval)
 
+    def add_hash_layer_count(self, count: int) -> None:
+        self.add_uint32(Keys.LLM.HASH_LAYER_COUNT.format(arch=self.arch), count)
+
     def add_feed_forward_length(self, length: int | Sequence[int]) -> None:
         if isinstance(length, int):
             self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
@@ -940,6 +943,39 @@ class GGUFWriter:
     def add_sliding_window(self, value: int) -> None:
         self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value)
 
+    def add_block_size(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value)
+
+    def add_target_layers(self, value: Sequence[int]) -> None:
+        self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value)
+
+    def add_target_hidden_size(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.TARGET_HIDDEN_SIZE.format(arch=self.arch), value)
+
+    def add_norm_before_residual(self, value: bool) -> None:
+        self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value)
+
+    def add_attention_output_group_count(self, count: int) -> None:
+        self.add_uint32(Keys.Attention.OUTPUT_GROUP_COUNT.format(arch=self.arch), count)
+
+    def add_attention_output_lora_rank(self, length: int) -> None:
+        self.add_uint32(Keys.Attention.OUTPUT_LORA_RANK.format(arch=self.arch), length)
+
+    def add_attention_compress_ratios(self, values: Sequence[int]) -> None:
+        self.add_array(Keys.Attention.COMPRESS_RATIOS.format(arch=self.arch), values)
+
+    def add_attention_compress_rope_freq_base(self, value: float) -> None:
+        self.add_float32(Keys.Attention.COMPRESS_ROPE_FREQ_BASE.format(arch=self.arch), value)
+
+    def add_hyper_connection_count(self, count: int) -> None:
+        self.add_uint32(Keys.HyperConnection.COUNT.format(arch=self.arch), count)
+
+    def add_hyper_connection_sinkhorn_iterations(self, count: int) -> None:
+        self.add_uint32(Keys.HyperConnection.SINKHORN_ITERATIONS.format(arch=self.arch), count)
+
+    def add_hyper_connection_epsilon(self, value: float) -> None:
+        self.add_float32(Keys.HyperConnection.EPSILON.format(arch=self.arch), value)
+
     def add_attention_scale(self, value: float) -> None:
         self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value)
 

gguf-py/gguf/tensor_mapping.py

@@ -1283,6 +1283,11 @@ class TensorNameMap:
         MODEL_TENSOR.ENC_OUTPUT_NORM: (
             "encoder.final_layer_norm", # t5
             "layer_norm",               # neobert
+            "model.hidden_norm",        # dflash
+        ),
+
+        MODEL_TENSOR.FC: (
+            "model.fc", # dflash
         ),
 
         MODEL_TENSOR.CLS: (

gpttype_adapter.cpp

@@ -20,6 +20,8 @@
 #include <cstdint>
 #include <string>
 #include <cctype>
+#include <cstdlib>
+#include <cfloat>
 #include <locale>
 #include <chrono>
 #include <algorithm>
@@ -31,12 +33,12 @@
 #include "utils.h"
 #include "llmutils.h"
 
-//for easier compilation
-//concat source files into one file for compilation purposes
+#ifdef GGML_USE_CUDA
+#  include "ggml-cuda.h"
+#endif
+
 #include "llama_v2.cpp"
 #include "llama_v3.cpp"
-#include "src/llama.cpp"
-#include "common/chat.cpp"
 #include "gptj_v1.cpp"
 #include "gptj_v2.cpp"
 #include "gptj_v3.cpp"
@@ -48,14 +50,23 @@
 #include "neox_v2.cpp"
 #include "neox_v3.cpp"
 #include "mpt_v3.cpp"
+
 #include "tools/mtmd/mtmd.h"
 #include "tools/mtmd/mtmd-helper.h"
 #include "common/speculative.h"
+#include "common/chat.h"
+#include "common/log.h"
+#include "llama-grammar.h"
 #include "vendor/stb/stb_image.h"
 #include "otherarch/sdcpp/thirdparty/stb_image_resize.h"
 #include "common/common.h"
 #include "common/fit.h"
 #include "ggml-rpc.h"
+#include "llama-impl.h"
+#include "llama-ext.h"
+#include "llama-model.h"
+#include "llama-vocab.h"
+#include "nlohmann/json.hpp"
 
 #if defined(GGML_USE_HIP)
 // for rocblas_initialize()
@@ -128,6 +139,7 @@ static llama_context * guidance_ctx = nullptr; //for classifier free guidance, w
 static mtmd_context * mtmd_ctx = nullptr; //for multimodal media
 static std::vector<media_object> media_objects;
 static std::vector<int> last_media_mem; //for storing dummy tokens that will be consumed by mtmd
+static std::vector<int> media_object_token_counts; //per media object dummy token counts for inline placeholders
 static std::string media_composite_image_signature = ""; //for identifying when the media changes, we need to invalidate the cache
 static int current_media_identifier = MEDIA_TOKEN_IDENTIFIER_A;
 static int vision_max_res = 2048;
@@ -184,6 +196,10 @@ static const int smartcache_rnn_lifeboat_percent = 65;
 static const int smartcache_rnn_lifeboat_extra_slot_min_user_slots = 4;
 
 extern bool kcpp_permit_any_repack;
+extern bool kcpp_pipeline_parallelism;
+extern bool OldBPETokenizerMode;
+extern int kcpp_extra_swa_padding;
+extern int kcpp_active_swa_size;
 
 inline int kcpp_cpu_has_blas(void) {
 #if defined(GGML_USE_BLAS) || defined(GGML_USE_CUDA) || defined(GGML_USE_VULKAN) || defined(GGML_USE_SYCL)
@@ -597,8 +613,9 @@ static size_t estimate_draft_autofit_tax_mb(
     if(draft_is_mtp_estimate)
     {
         draft_ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+        draft_ctx_params.n_seq_max = base_ctx_params.n_seq_max;
         draft_ctx_params.n_rs_seq = speculative_chunk_amt;
-        draft_ctx_params.n_outputs_max = base_ctx_params.n_seq_max; //match the real MTP draft context (see speculative_decoding_setup) so the autofit tax doesn't over-reserve the draft compute buffer at n_batch*n_vocab (~2GB on large-vocab models like Gemma)
+        draft_ctx_params.n_outputs_max = std::max<uint32_t>(1, base_ctx_params.n_seq_max); //match the real MTP draft context so the autofit tax doesn't over-reserve the draft compute buffer at n_batch*n_vocab (~2GB on large-vocab models like Gemma)
         measure_model_bytes = has_draft_model;
     }
 
@@ -883,7 +900,7 @@ static void mtp_decoding_setup(llama_model * main_model, llama_context * main_ct
     mtp_ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
     mtp_ctx_params.ctx_other = main_ctx;
     mtp_ctx_params.n_rs_seq = 0;
-    mtp_ctx_params.n_outputs_max = 1;
+    mtp_ctx_params.n_outputs_max = std::max<uint32_t>(1, mtp_ctx_params.n_seq_max);
 
     printf("\nAttempting to create built-in MTP context from the main model.\n");
     draft_ctx = llama_init_from_model(main_model, mtp_ctx_params);
@@ -2050,6 +2067,77 @@ void sample_temperature(llama_token_data_array * candidates_p, float temp, float
     }
 }
 
+static std::pair<std::vector<uint32_t>, llama_partial_utf8> kcpp_decode_utf8(
+        const std::string & src,
+        llama_partial_utf8 partial_start) {
+    static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 3, 4 };
+    const char * pos = src.c_str();
+    std::vector<uint32_t> code_points;
+
+    code_points.reserve(src.size() + 1);
+    uint32_t value    = partial_start.value;
+    int      n_remain = partial_start.n_remain;
+
+    while (*pos != 0 && n_remain > 0) {
+        uint8_t next_byte = static_cast<uint8_t>(*pos);
+        if ((next_byte >> 6) != 2) {
+            code_points.push_back(0);
+            return std::make_pair(std::move(code_points), llama_partial_utf8{ 0, -1 });
+        }
+        value = (value << 6) + (next_byte & 0x3F);
+        ++pos;
+        --n_remain;
+    }
+
+    if (partial_start.n_remain > 0 && n_remain == 0) {
+        code_points.push_back(value);
+    }
+
+    while (*pos != 0) {
+        uint8_t first_byte = static_cast<uint8_t>(*pos);
+        uint8_t highbits   = first_byte >> 4;
+        n_remain = lookup[highbits] - 1;
+
+        if (n_remain < 0) {
+            code_points.clear();
+            code_points.push_back(0);
+            return std::make_pair(std::move(code_points), llama_partial_utf8{ 0, n_remain });
+        }
+
+        uint8_t mask = (1 << (7 - n_remain)) - 1;
+        value = first_byte & mask;
+
+        ++pos;
+        while (*pos != 0 && n_remain > 0) {
+            value = (value << 6) + (static_cast<uint8_t>(*pos) & 0x3F);
+            ++pos;
+            --n_remain;
+        }
+        if (n_remain == 0) {
+            code_points.push_back(value);
+        }
+    }
+    code_points.push_back(0);
+
+    return std::make_pair(std::move(code_points), llama_partial_utf8{ value, n_remain });
+}
+
+static llama_grammar_candidates kcpp_llama_grammar_reject_candidates(
+        const llama_grammar_rules      & rules,
+        const llama_grammar_stacks     & stacks,
+        const llama_grammar_candidates & candidates) {
+    if (stacks.empty()) {
+        return {};
+    }
+
+    auto rejects = llama_grammar_reject_candidates_for_stack(rules, stacks.front(), candidates);
+    for (size_t i = 1, size = stacks.size(); i < size; ++i) {
+        rejects = llama_grammar_reject_candidates_for_stack(rules, stacks[i], rejects);
+    }
+
+    return rejects;
+}
+
 void sample_grammar(FileFormat file_format, int32_t n_vocab, llama_token_data_array * candidates, const struct llama_grammar * grammar) {
 
     const int64_t t_start_sample_us = ggml_time_us();
@@ -2082,12 +2170,12 @@ void sample_grammar(FileFormat file_format, int32_t n_vocab, llama_token_data_ar
         } else if (piece.empty() || piece[0] == 0) {
             rejects[i] = true;
         } else {
-            candidates_decoded.push_back(decode_utf8(piece.c_str(), grammar->partial_utf8));
+            candidates_decoded.push_back(kcpp_decode_utf8(piece, grammar->partial_utf8));
             candidates_grammar.push_back({ i, candidates_decoded.back().first.data(), candidates_decoded.back().second });
         }
     }
 
-    for (auto reject: llama_grammar_reject_candidates(grammar->rules, grammar->stacks, candidates_grammar)) {
+    for (auto reject: kcpp_llama_grammar_reject_candidates(grammar->rules, grammar->stacks, candidates_grammar)) {
         rejects[reject.index] = true;
     }
 
@@ -2321,7 +2409,7 @@ static void grammar_accept_token(FileFormat file_format, int32_t n_vocab, struct
     const std::string piece = FileFormatTokenizeID(token,file_format);
 
     // Note terminating 0 in decoded string
-    const auto   decoded     = decode_utf8(piece.c_str(), grammar->partial_utf8);
+    const auto   decoded     = kcpp_decode_utf8(piece, grammar->partial_utf8);
     const auto & code_points = decoded.first;
     for (auto it = code_points.begin(), end = code_points.end() - 1; it != end; ++it) {
         llama_grammar_accept(grammar, *it);
@@ -2395,6 +2483,201 @@ static bool mtmd_text_chunk_has_invalid_tokens(const mtmd_input_chunk * mtmdchun
     return false;
 }
 
+static bool kcpp_is_media_token(int token)
+{
+    return token <= MEDIA_TOKEN_IDENTIFIER_A && token > MEDIA_TOKEN_IDENTIFIER_A - 4096;
+}
+
+static int kcpp_media_token_for_index(int media_index)
+{
+    return current_media_identifier - (media_index * 2);
+}
+
+static int kcpp_media_index_from_token(int token)
+{
+    if(!kcpp_is_media_token(token))
+    {
+        return -1;
+    }
+    const int diff = current_media_identifier - token;
+    if(diff < 0 || (diff % 2) != 0)
+    {
+        return -1;
+    }
+    return diff / 2;
+}
+
+static bool kcpp_parse_attached_media_placeholder(
+    const std::string & prompt,
+    size_t pos,
+    int image_count,
+    int audio_count,
+    int & media_index,
+    size_t & placeholder_len)
+{
+    const std::string image_prefix = "(Attached Image ";
+    const std::string audio_prefix = "(Attached Audio ";
+    bool is_audio = false;
+    size_t prefix_len = 0;
+
+    if(prompt.compare(pos, image_prefix.size(), image_prefix) == 0)
+    {
+        prefix_len = image_prefix.size();
+    }
+    else if(prompt.compare(pos, audio_prefix.size(), audio_prefix) == 0)
+    {
+        is_audio = true;
+        prefix_len = audio_prefix.size();
+    }
+    else
+    {
+        return false;
+    }
+
+    size_t number_pos = pos + prefix_len;
+    size_t end_pos = number_pos;
+    while(end_pos < prompt.size() && std::isdigit((unsigned char) prompt[end_pos]))
+    {
+        ++end_pos;
+    }
+    if(end_pos == number_pos || end_pos >= prompt.size() || prompt[end_pos] != ')')
+    {
+        return false;
+    }
+
+    int number = std::atoi(prompt.substr(number_pos, end_pos - number_pos).c_str());
+    if(number <= 0)
+    {
+        return false;
+    }
+
+    int candidate = -1;
+    if(is_audio)
+    {
+        if(number <= audio_count)
+        {
+            candidate = image_count + number - 1;
+        }
+        else if(number <= (int) media_objects.size() && media_objects[number - 1].is_audio)
+        {
+            candidate = number - 1;
+        }
+    }
+    else
+    {
+        if(number <= image_count)
+        {
+            candidate = number - 1;
+        }
+        else if(number <= (int) media_objects.size() && !media_objects[number - 1].is_audio)
+        {
+            candidate = number - 1;
+        }
+    }
+
+    if(candidate < 0 || candidate >= (int) media_objects.size())
+    {
+        return false;
+    }
+    media_index = candidate;
+    placeholder_len = end_pos - pos + 1;
+    return true;
+}
+
+static void kcpp_append_media_placeholder_tokens(std::vector<int> & tokens, int media_index)
+{
+    if(media_index < 0 || media_index >= (int) media_object_token_counts.size())
+    {
+        return;
+    }
+    const int token_count = media_object_token_counts[media_index];
+    const int media_token = kcpp_media_token_for_index(media_index);
+    for(int i = 0; i < token_count; ++i)
+    {
+        tokens.push_back(media_token);
+    }
+}
+
+static bool kcpp_tokenize_prompt_with_inline_media(
+    const std::string & prompt,
+    std::vector<int> & output_tokens,
+    FileFormat file_format,
+    bool add_bos,
+    int image_count,
+    int audio_count)
+{
+    output_tokens.clear();
+    bool inserted_media = false;
+    bool emitted_anything = false;
+    size_t text_start = 0;
+
+    auto append_text = [&](size_t start, size_t end)
+    {
+        if(end <= start)
+        {
+            return;
+        }
+        std::vector<int> text_tokens;
+        TokenizeString(prompt.substr(start, end - start), text_tokens, file_format, add_bos && !emitted_anything);
+        output_tokens.insert(output_tokens.end(), text_tokens.begin(), text_tokens.end());
+        emitted_anything = true;
+    };
+
+    for(size_t pos = 0; pos < prompt.size(); ++pos)
+    {
+        int media_index = -1;
+        size_t placeholder_len = 0;
+        if(kcpp_parse_attached_media_placeholder(prompt, pos, image_count, audio_count, media_index, placeholder_len))
+        {
+            append_text(text_start, pos);
+            if(add_bos && !emitted_anything)
+            {
+                std::vector<int> bos;
+                TokenizeString("", bos, file_format, true);
+                output_tokens.insert(output_tokens.end(), bos.begin(), bos.end());
+            }
+            kcpp_append_media_placeholder_tokens(output_tokens, media_index);
+            emitted_anything = true;
+            inserted_media = true;
+            pos += placeholder_len - 1;
+            text_start = pos + 1;
+        }
+    }
+
+    append_text(text_start, prompt.size());
+    if(!inserted_media)
+    {
+        output_tokens.clear();
+    }
+    return inserted_media;
+}
+
+static int kcpp_adjust_media_truncation_start(const std::vector<int> & tokens, int offset)
+{
+    if(offset <= 0 || offset >= (int) tokens.size())
+    {
+        return offset;
+    }
+    const int token = tokens[offset];
+    if(kcpp_is_media_token(token) && tokens[offset - 1] == token)
+    {
+        while(offset < (int) tokens.size() && tokens[offset] == token)
+        {
+            ++offset;
+        }
+    }
+    return offset;
+}
+
+static bool kcpp_media_span_boundary_ok(const std::vector<int> & tokens, int pos)
+{
+    if(pos <= 0 || pos >= (int) tokens.size())
+    {
+        return true;
+    }
+    return !(kcpp_is_media_token(tokens[pos]) && tokens[pos - 1] == tokens[pos]);
+}
+
 //given an old GGUF context and a new context that has some middle portion removed,
 //find and remove the middle portion from the old context from the KV. Does not fast forward after this destructive action
 //returns true if contextshift is doable, executes it if dryrun is false
@@ -2415,6 +2698,11 @@ bool DoContextShifting(llama_context * ctx, llama_context * draft_ctx, std::vect
 
     for (int i = 0; i < current_context_tokens.size(); ++i)
     {
+        if(i >= new_tokens_len)
+        {
+            purgeneeded = false;
+            break;
+        }
         if (current_context_tokens[i] == new_context_tokens[i])
         {
             trimstart += 1;
@@ -2449,6 +2737,14 @@ bool DoContextShifting(llama_context * ctx, llama_context * draft_ctx, std::vect
         int found = ArrFindIndexOf(current_context_tokens,shared);
         if(found>=0 && found > trimstart)
         {
+            if(!kcpp_media_span_boundary_ok(current_context_tokens, trimstart) || !kcpp_media_span_boundary_ok(current_context_tokens, found))
+            {
+                if(debugmode==1 && !is_quiet)
+                {
+                    printf("\n[Context Shifting skipped: refusing to split a multimodal span]");
+                }
+                return false;
+            }
             bool ok = true;
             if(!dryrun)
             {
@@ -4056,7 +4352,7 @@ public:
         std::unique_lock<std::mutex> lock(batch_mutex);
         batch_legacy_waiting++;
         batch_cv.notify_all();
-        batch_cv.wait(lock, [](){ return !batch_has_live_locked(); });
+        batch_cv.wait(lock, [](){ return !batch_legacy_active && !batch_has_live_locked(); });
         batch_legacy_waiting--;
         batch_invalidate_legacy_context_locked();
         batch_legacy_active = true;
@@ -4552,7 +4848,7 @@ static void batch_start_worker_locked()
 
 bool gpttype_batch_generate_enabled()
 {
-    return continuous_batching_slots > 1 && file_format == FileFormat::GGUF_GENERIC && llama_ctx_v4 && kcpp_data;
+    return continuous_batching_slots > 1 && file_format == FileFormat::GGUF_GENERIC && llama_ctx_v4 && kcpp_data && !draft_ctx && !guidance_ctx;
 }
 
 int gpttype_batch_generate_submit(const generation_inputs inputs)
@@ -4994,6 +5290,7 @@ static void PrepareMediaEmbds(const int nctx, const std::vector<int> & media_int
         int introsize = media_intro.size();
         int outrosize = media_outro.size();
         last_media_mem.clear();
+        media_object_token_counts.clear();
 
         for(int i=0;i<media_objects.size();++i)
         {
@@ -5004,6 +5301,7 @@ static void PrepareMediaEmbds(const int nctx, const std::vector<int> & media_int
                 : kcpp_mtmd_bitmap_init_image_from_buf(media_data_buffer.data(), media_data_buffer.size(), vision_max_res));
             if(!bitmap.ptr)
             {
+                media_object_token_counts.push_back(0);
                 printf("\nError: MTMD media %d failed to load!",i);
                 continue;
             }
@@ -5017,6 +5315,7 @@ static void PrepareMediaEmbds(const int nctx, const std::vector<int> & media_int
             int32_t tokenized = mtmd_tokenize(mtmd_ctx, chunks.ptr.get(), &inp_txt, bitmaps.data(), bitmaps.size());
             if(tokenized != 0)
             {
+                media_object_token_counts.push_back(0);
                 media_composite_image_signature = ""; //force invalidate
                 printf("\nError: MTMD media %d failed to tokenize! (status %d)",i, tokenized);
                 continue;
@@ -5079,18 +5378,21 @@ static void PrepareMediaEmbds(const int nctx, const std::vector<int> & media_int
             }
             if(mediatokensneeded>0 && mediatokensneeded < nctx)
             {
+                media_object_token_counts.push_back(mediatokensneeded);
                 int tokcnt = mediatokensneeded;
                 if(i==0)
                 {
                     tokcnt += introsize + outrosize;
                 }
+                const int media_token = kcpp_media_token_for_index(i);
                 for(int n=0;n<tokcnt;++n)
                 {
-                    last_media_mem.push_back(current_media_identifier);
+                    last_media_mem.push_back(media_token);
                 }
             }
             else
             {
+                media_object_token_counts.push_back(0);
                 media_composite_image_signature = ""; //force invalidate
                 printf("\nWarning: Media excluded - Context size too low or not enough mtmd tokens! (needed %d)\nMedia will be IGNORED! You probably want to relaunch with a larger context size!\n",mediatokensneeded);
             }
@@ -5343,11 +5645,11 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
     std::vector<int> media_intro; //added before media list
     std::vector<int> media_outro; //added before media list
     std::string intro = "\nAttached Media:\n";
-    if(mtmd_ctx && kcpp_mtmd_is_gemma4uv(mtmd_ctx)) //ugly fix for gemma4uv vision coherency
-    {
-        intro = "\n<|channel><channel|>" + intro;
-    }
-    TokenizeString(intro, media_intro, file_format, add_bos_token);
+    // if(mtmd_ctx && kcpp_mtmd_is_gemma4uv(mtmd_ctx)) //ugly fix for gemma4uv vision coherency
+    // {
+    //     intro = "\n<|channel><channel|>" + intro;
+    // }
+    TokenizeString(intro, media_intro, file_format, false);
 
     //clear previous run media memory, just-in-time free
     for(int i=0;i<media_objects.size();++i)
@@ -5600,6 +5902,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
     if(media_composite_image_signature=="")
     {
         last_media_mem.clear();
+        media_object_token_counts.clear();
     }
     if(media_data_changed)
     {
@@ -5607,7 +5910,21 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
         media_embds_built = true;
     }
 
-    TokenizeString(kcpp_data->prompt, embd_inp, file_format, add_bos_token);
+    bool media_inserted_inline = false;
+    if(last_media_mem.size()>0)
+    {
+        media_inserted_inline = kcpp_tokenize_prompt_with_inline_media(
+            kcpp_data->prompt,
+            embd_inp,
+            file_format,
+            add_bos_token,
+            inputs.images_len,
+            inputs.audio_len);
+    }
+    if(!media_inserted_inline)
+    {
+        TokenizeString(kcpp_data->prompt, embd_inp, file_format, add_bos_token);
+    }
     if(addedmemory!="")
     {
         TokenizeString(addedmemory, embd_inp_mem, file_format, add_bos_token);
@@ -5620,6 +5937,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
         std::vector<int> bos;
         TokenizeString("", bos, file_format, add_bos_token);
         int offset = embd_inp.size() - nctx + kcpp_data->n_predict;
+        offset = kcpp_adjust_media_truncation_start(embd_inp, offset);
         embd_inp = std::vector<int>(embd_inp.begin() + offset, embd_inp.end());
         //replace bos into front if exists
         if(bos.size()>0 && embd_inp.size()>0)
@@ -5628,7 +5946,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
         }
     }
 
-    if(last_media_mem.size()>0) //stick the media placeholders before the added mem
+    if(last_media_mem.size()>0 && !media_inserted_inline) //stick the media placeholders before the added mem if no inline placeholders were found
     {
         if(last_media_mem.size() + kcpp_data->n_predict + 4 > nctx)
         {
@@ -6757,7 +7075,8 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
             while ((int)embd_inp.size() > input_consumed)
             {
                 int currtoken = embd_inp[input_consumed];
-                if(currtoken==MEDIA_TOKEN_IDENTIFIER_A || currtoken==MEDIA_TOKEN_IDENTIFIER_B) //special media token hit
+                int curr_media_index = kcpp_media_index_from_token(currtoken);
+                if(curr_media_index >= 0) //special media token hit
                 {
                     if(!media_embds_built) //this should never happen! however, handle it anyway
                     {
@@ -6778,7 +7097,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                         int mediatokensevaled = 0;
                         int introsize = media_intro.size();
                         int outrosize = media_outro.size();
-                        while(input_consumed < embd_inp.size() && (embd_inp[input_consumed]==MEDIA_TOKEN_IDENTIFIER_A || embd_inp[input_consumed]==MEDIA_TOKEN_IDENTIFIER_B))
+                        while(input_consumed < embd_inp.size() && embd_inp[input_consumed]==currtoken)
                         {
                             if (!last_n_tokens.empty())
                             {
@@ -6789,10 +7108,15 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                             ++input_consumed;
                             ++mediatokenscounted;
                         }
-                        for(int i=0;i<media_objects.size();++i)
+                        bool include_media_header = false;
+                        if(curr_media_index == 0 && curr_media_index < (int) media_object_token_counts.size())
+                        {
+                            include_media_header = (mediatokenscounted == media_object_token_counts[curr_media_index] + introsize + outrosize);
+                        }
+                        if(curr_media_index < (int) media_objects.size())
                         {
                             //note: no handling for draft_ctx as we don't support vision for it
-                            if(introsize>0 && i==0)
+                            if(include_media_header && introsize>0)
                             {
                                 //added at the start of everything
                                 kcpp_embd_batch batch = kcpp_embd_batch(media_intro, n_past, use_mrope, false);
@@ -6809,10 +7133,10 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                                 mediatokensevaled += introsize;
                             }
 
-                            int start_size = media_objects[i].chunk_start_seq.size();
+                            int start_size = media_objects[curr_media_index].chunk_start_seq.size();
                             if (start_size > 0) {
                                 //add a separator between each image
-                                kcpp_embd_batch batch = kcpp_embd_batch(media_objects[i].chunk_start_seq, n_past, use_mrope, false);
+                                kcpp_embd_batch batch = kcpp_embd_batch(media_objects[curr_media_index].chunk_start_seq, n_past, use_mrope, false);
                                 auto evr = llama_decode(llama_ctx_v4, batch.batch);
                                 if(evr!=0)
                                 {
@@ -6826,12 +7150,12 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                                 mediatokensevaled += start_size;
                             }
 
-                            for(int j=0;j<media_objects[i].mediachunks.size();++j)
+                            for(int j=0;j<media_objects[curr_media_index].mediachunks.size();++j)
                             {
-                                media_chunk chunk = media_objects[i].mediachunks[j];
+                                media_chunk chunk = media_objects[curr_media_index].mediachunks[j];
                                 if(allow_regular_prints)
                                 {
-                                    printf("\rProcessing Media Embedding %d (%d tokens)",(i+1), chunk.clp_image_tokens);
+                                    printf("\rProcessing Media Embedding %d (%d tokens)",(curr_media_index+1), chunk.clp_image_tokens);
                                 }
                                 bool err = kcpp_eval_media(llama_ctx_v4,chunk,kcpp_data->n_batch,&n_past);
                                 mediatokensevaled += chunk.clp_image_tokens;
@@ -6849,10 +7173,10 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                                 }
                             }
 
-                            int end_size = media_objects[i].chunk_end_seq.size();
+                            int end_size = media_objects[curr_media_index].chunk_end_seq.size();
                             if (end_size > 0) {
                                 //add a separator between each image
-                                kcpp_embd_batch batch = kcpp_embd_batch(media_objects[i].chunk_end_seq, n_past, use_mrope, false);
+                                kcpp_embd_batch batch = kcpp_embd_batch(media_objects[curr_media_index].chunk_end_seq, n_past, use_mrope, false);
                                 auto evr = llama_decode(llama_ctx_v4, batch.batch);
                                 if(evr!=0)
                                 {
@@ -6866,7 +7190,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs)
                                 mediatokensevaled += end_size;
                             }
                         }
-                        if(media_objects.size()>0 && outrosize>0)
+                        if(include_media_header && media_objects.size()>0 && outrosize>0)
                         {
                             //added after all media but before prompt
                             kcpp_embd_batch batch = kcpp_embd_batch(media_outro, n_past, use_mrope, false);

include/llama.h

@@ -162,6 +162,9 @@ extern "C" {
         LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file
     };
 
+    // Get the model file type (quantization) as a string, e.g. "Q8_0" or "Q4_K - Medium"
+    LLAMA_API const char * llama_ftype_name(enum llama_ftype ftype);
+
     enum llama_rope_scaling_type {
         LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1,
         LLAMA_ROPE_SCALING_TYPE_NONE        = 0,
@@ -609,6 +612,9 @@ extern "C" {
     // Get a string describing the model type
     LLAMA_API int32_t llama_model_desc(const struct llama_model * model, char * buf, size_t buf_size);
 
+    // Get the model file type (quantization), e.g. LLAMA_FTYPE_MOSTLY_Q8_0
+    LLAMA_API enum llama_ftype llama_model_ftype(const struct llama_model * model);
+
     // Returns the total size of all the tensors in the model in bytes
     LLAMA_API uint64_t llama_model_size(const struct llama_model * model);
 

koboldcpp.py

@@ -78,7 +78,7 @@ dry_seq_break_max = 128
 extra_images_max = 4 # for kontext/qwen img
 
 # global vars
-KcppVersion = "1.116"
+KcppVersion = "1.117.1"
 showdebug = True
 kcpp_instance = None #global running instance
 global_memory = {"tunnel_url": "", "restart_target":"", "input_to_exit":False, "load_complete":False, "restart_override_base_config":"", "last_active_timestamp":datetime.now(), "triggered_sleeping":False, "current_model":"initial_model", "base_config":"", "swapReqType": None, "autoswapmode": False}
@@ -97,6 +97,8 @@ musicName = None
 imageName = None
 mmprojName = None
 lastgeneratedcomfyimg = b''
+lastgeneratedcachedimg = b''
+lastgeneratedcachedimgkey = b''
 lastuploadedcomfyimg = b''
 fullsdmodelpath = ""  #if empty, it's not initialized
 password = "" #if empty, no auth key required
@@ -376,14 +378,12 @@ class generation_outputs(ctypes.Structure):
 class sd_load_model_inputs(ctypes.Structure):
     _fields_ = [("model_filename", ctypes.c_char_p),
                 ("executable_path", ctypes.c_char_p),
-                ("kcpp_main_device", ctypes.c_int),
+                ("backend", ctypes.c_char_p),
                 ("threads", ctypes.c_int),
                 ("quant", ctypes.c_int),
                 ("flash_attention", ctypes.c_bool),
-                ("offload_cpu", ctypes.c_bool),
+                ("params_backend", ctypes.c_char_p),
                 ("use_mmap", ctypes.c_bool),
-                ("kcpp_vae_device", ctypes.c_int),
-                ("kcpp_clip_device", ctypes.c_int),
                 ("diffusion_conv_direct", ctypes.c_bool),
                 ("vae_conv_direct", ctypes.c_bool),
                 ("taesd", ctypes.c_bool),
@@ -401,8 +401,10 @@ class sd_load_model_inputs(ctypes.Structure):
                 ("upscaler_filename", ctypes.c_char_p),
                 ("img_hard_limit", ctypes.c_int),
                 ("img_soft_limit", ctypes.c_int),
-                ("max_vram", ctypes.c_float),
+                ("max_vram", ctypes.c_char_p),
+                ("split_mode", ctypes.c_char_p),
                 ("stream_layers", ctypes.c_bool),
+                ("auto_fit", ctypes.c_bool),
                 ("devices_override", ctypes.c_char_p),
                 ("quiet", ctypes.c_bool),
                 ("debugmode", ctypes.c_int)]
@@ -2351,7 +2353,7 @@ def continuous_batching_python_eligible(genparams, api_format):
     if model_path and not model_path.endswith(".gguf"):
         utfprint("Batching disabled due to file format",2)
         return False
-    if not getattr(args, "noshift", False) or getattr(args, "smartcontext", False) or getattr(args, "draftmodel", "") or getattr(args, "enableguidance", False):
+    if not getattr(args, "noshift", False) or getattr(args, "smartcontext", False) or getattr(args, "draftmodel", "") or getattr(args, "usemtp", False) or getattr(args, "enableguidance", False):
         utfprint("Batching disabled due to loaded settings",2)
         return False
     if genparams.get("negative_prompt") or genparams.get("images") or genparams.get("audio"):
@@ -2485,6 +2487,21 @@ def sd_resolve_device(name, default_=-1):
         name = str(max(name, -2))
     return sd_get_device_number(name)
 
+def sd_get_device_override(deviceid, module=''):
+    '''formats a device id and a module name in sd.cpp --backend syntax'''
+    global cached_sd_info
+    devices = cached_sd_info.get('devices', [])
+    device_name = ''
+    if deviceid <= -2:
+        device_name = "CPU"
+    elif deviceid >= 0 and deviceid < len(devices):
+        device_name = devices[deviceid]['name']
+    if device_name and module:
+        result = module + '=' + device_name
+    else:
+        result = device_name
+    return result
+
 def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip2_filename,photomaker_filename,upscaler_filename,audio_vae_filename):
     global args
     inputs = sd_load_model_inputs()
@@ -2499,10 +2516,14 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
     inputs.threads = thds
     inputs.quant = args.sdquant
     inputs.flash_attention = args.sdflashattention
-    inputs.offload_cpu = args.sdoffloadcpu
+    inputs.params_backend = b'CPU' if args.sdoffloadcpu else b''
     inputs.use_mmap = args.usemmap
-    inputs.kcpp_vae_device = sd_resolve_device(args.sdvaedevice, default_sdvaedevice)
-    inputs.kcpp_clip_device = sd_resolve_device(args.sdclipdevice, default_sdclipdevice)
+    backends = [
+        sd_get_device_override(sd_resolve_device(args.sdmaingpu, 'main')),
+        sd_get_device_override(sd_resolve_device(args.sdclipdevice, default_sdclipdevice), 'CLIP'),
+        sd_get_device_override(sd_resolve_device(args.sdvaedevice, default_sdvaedevice), 'VAE'),
+    ]
+    inputs.backend = ','.join([b for b in backends if b]).encode("UTF-8")
     sdconvdirect = sd_convdirect_option(args.sdconvdirect)
     inputs.diffusion_conv_direct = sdconvdirect == 'full'
     inputs.vae_conv_direct = sdconvdirect in ['vaeonly', 'full']
@@ -2515,7 +2536,7 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
     inputs.clip2_filename = clip2_filename.encode("UTF-8")
     inputs.photomaker_filename = photomaker_filename.encode("UTF-8")
     inputs.upscaler_filename = upscaler_filename.encode("UTF-8")
-    inputs.max_vram = (args.sdvramlimit/1024.0) if args.sdvramlimit > 0 else 0
+    inputs.max_vram = str((args.sdvramlimit/1024.0) if args.sdvramlimit > 0 else '').encode('UTF-8')
     inputs.stream_layers = False
 
     lora_filenames, lora_multipliers = prepare_initial_lora_multipliers()
@@ -2533,7 +2554,6 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
     inputs.img_hard_limit = args.sdclamped
     inputs.img_soft_limit = args.sdclampedsoft
     inputs = set_backend_props(inputs)
-    inputs.kcpp_main_device = sd_resolve_device(args.sdmaingpu, 'main')
     ret = handle.sd_load_model(inputs)
     return ret
 
@@ -3772,22 +3792,17 @@ def format_jinja(messages_orig, tools, chat_template_kwargs=None):
             if isinstance(m.get("content"), list):
                 normalized = []
                 turn_text = ""
-                media_text = ""
                 for item in m["content"]:
                     if item.get("type")=="text":
                         turn_text += item.get("text","")
-                for item in m["content"]:
-                    if item.get("type")=="text":
-                        pass
                     elif item.get("type")=="image_url" or item.get("type")=="image":
-                        media_text += f"\n(Attached Image {mediacount})\n"
+                        turn_text += f"\n(Attached Image {mediacount})\n"
                         mediacount += 1
                     elif item.get("type")=="input_audio":
-                        media_text += f"\n(Attached Audio {mediacount})\n"
+                        turn_text += f"\n(Attached Audio {mediacount})\n"
                         mediacount += 1
                     else:
                         normalized.append(item)
-                turn_text = media_text + turn_text
                 if turn_text:
                     normalized.append({"type": "text","text": turn_text})
                 m["content"] = normalized
@@ -3882,6 +3897,29 @@ def normalize_tool_call_resp(obj): # Normalize various tool call formats to Open
 
     return obj
 
+def convert_tool_calls_to_ollama(tool_calls):
+    ollama_tool_calls = []
+    for idx, tool_call in enumerate(tool_calls or []):
+        try:
+            func = tool_call.get("function", {})
+            args = func.get("arguments", {})
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args)
+                except Exception:
+                    args = {}
+            ollama_tool_calls.append({
+                "type": "function",
+                "function": {
+                    "index": idx,
+                    "name": func.get("name", ""),
+                    "arguments": args
+                }
+            })
+        except Exception:
+            pass
+    return ollama_tool_calls
+
 # Used to parse json for openai tool calls
 def extract_json_from_string(input_string, check_strict=False):
     parsed_json = None
@@ -4847,7 +4885,7 @@ class KcppProxyHandler(http.server.BaseHTTPRequestHandler):
         is_chat_completions_path = (clean_path.endswith('/v1/chat/completions') or clean_path=='/chat/completions')
 
         #any requests to the following endpoints is capable of waking the server
-        wake_requests = ["/api/extra/generate/stream","/api/extra/tokencount","/api/v1/generate","/sdapi/v1/interrogate","/v1/completions","/v1/chat/completions","/v1/responses","/completions","/chat/completions","/responses","/api/extra/transcribe","/v1/audio/transcriptions","/api/extra/tts","/v1/audio/speech","/api/extra/embeddings","/v1/embeddings","/api/extra/music/prepare","/api/extra/music/generate","/sdapi/v1/txt2img","/sdapi/v1/img2img","/sdapi/v1/upscale"]
+        wake_requests = ["/api/extra/generate/stream","/api/extra/tokencount","/api/v1/generate","/sdapi/v1/interrogate","/v1/completions","/v1/chat/completions","/v1/responses","/completions","/chat/completions","/responses","/api/extra/transcribe","/v1/audio/transcriptions","/api/extra/tts","/v1/audio/speech","/api/extra/embeddings","/v1/embeddings","/api/embed","/api/extra/music/prepare","/api/extra/music/generate","/sdapi/v1/txt2img","/sdapi/v1/img2img","/sdapi/v1/upscale"]
         is_wake_request = clean_path in wake_requests
 
         autoswapEnabled = global_memory["autoswapmode"] is not None and global_memory["autoswapmode"]
@@ -4899,7 +4937,7 @@ class KcppProxyHandler(http.server.BaseHTTPRequestHandler):
                 textReqs = ["/api/extra/generate/stream","/api/extra/tokencount","/api/v1/generate","/sdapi/v1/interrogate","/v1/completions","/v1/chat/completions","/v1/responses","/completions","/chat/completions","/responses"]
                 sttReqs = ["/api/extra/transcribe","/v1/audio/transcriptions"]
                 ttsReqs = ["/api/extra/tts", "/v1/audio/speech"]
-                embedReqs = ["/api/extra/embeddings", "/v1/embeddings"]
+                embedReqs = ["/api/extra/embeddings", "/v1/embeddings", "/api/embed"]
                 musicReqs = ["/api/extra/music/prepare","/api/extra/music/generate"]
                 imageReqs = ["/sdapi/v1/txt2img", "/sdapi/v1/img2img", "/sdapi/v1/upscale"] # "/sdapi/v1/sd-models", "/sdapi/v1/options", "/sdapi/v1/samplers"
 
@@ -5207,9 +5245,28 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
 
         utfprint("\nOutput: " + recvtxt,1)
 
+        #handle potential think tags, but only chat completions will return them. the others just drop them
+        reasoningtxt = ""
+        if api_format==4 or api_format==8 or api_format==9: #chat completions, responses and anthropic messages, but only chat has reasoning returned
+            if recvtxt:
+                for pair in thinkformats:
+                    starter = pair['start']
+                    ender = pair['end']
+                    start_idx = recvtxt.find(starter)
+                    end_idx = recvtxt.find(ender, start_idx + len(starter))
+                    if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
+                        reasoningtxt = recvtxt[start_idx + len(starter):end_idx]
+                        recvtxt = recvtxt[:start_idx] + recvtxt[end_idx + len(ender):]
+                        break
+                    elif starter not in recvtxt and ender in recvtxt:
+                        parts = recvtxt.split(ender, 1)
+                        reasoningtxt = parts[0]
+                        recvtxt = parts[1]
+                        break
+
         #tool calls resolution
         tool_calls = []
-        if api_format == 4 or api_format == 2 or api_format == 8 or api_format == 9:
+        if api_format == 4 or api_format == 2 or api_format == 7 or api_format == 8 or api_format == 9:
             using_openai_tools = genparams.get('using_openai_tools', False)
             if using_openai_tools:
                 # first, let llama.cpp's chat parser handle known template-specific tool formats
@@ -5237,25 +5294,6 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
         modelNameToReturn = friendlymodelname
         if autoswapmode and textName is not None:
             modelNameToReturn = textName
-
-        #handle potential think tags, but only chat completions will return them. the others just drop them
-        reasoningtxt = ""
-        if api_format==4 or api_format==8 or api_format==9: #chat completions, responses and anthropic messages, but only chat has reasoning returned
-            if recvtxt:
-                for pair in thinkformats:
-                    starter = pair['start']
-                    ender = pair['end']
-                    start_idx = recvtxt.find(starter)
-                    end_idx = recvtxt.find(ender, start_idx + len(starter))
-                    if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
-                        reasoningtxt = recvtxt[start_idx + len(starter):end_idx]
-                        recvtxt = recvtxt[:start_idx] + recvtxt[end_idx + len(ender):]
-                        break
-                    elif starter not in recvtxt and ender in recvtxt:
-                        parts = recvtxt.split(ender, 1)
-                        reasoningtxt = parts[0]
-                        recvtxt = parts[1]
-                        break
         if api_format == 1:
             res = {"data": {"seqs": [recvtxt]}}
         elif api_format == 3:
@@ -5278,7 +5316,11 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
             tokarr = tokenize_ids(oldprompt+recvtxt,False)
             res = {"model": modelNameToReturn,"created_at": str(datetime.now(timezone.utc).isoformat()),"response":recvtxt,"done": True,"done_reason":currfinishreason,"context": tokarr,"total_duration": 1,"load_duration": 1,"prompt_eval_count": prompttokens,"prompt_eval_duration": 1,"eval_count": comptokens,"eval_duration": 1}
         elif api_format == 7:
-            res = {"model": modelNameToReturn,"created_at": str(datetime.now(timezone.utc).isoformat()),"message":{"role":"assistant","content":recvtxt},"done": True,"done_reason":currfinishreason,"total_duration": 1,"load_duration": 1,"prompt_eval_count": prompttokens,"prompt_eval_duration": 1,"eval_count": comptokens,"eval_duration": 1}
+            ccmsg = {"role":"assistant","content":recvtxt or ""}
+            ollama_tool_calls = convert_tool_calls_to_ollama(tool_calls)
+            if ollama_tool_calls:
+                ccmsg["tool_calls"] = ollama_tool_calls
+            res = {"model": modelNameToReturn,"created_at": str(datetime.now(timezone.utc).isoformat()),"message":ccmsg,"done": True,"done_reason":currfinishreason,"total_duration": 1,"load_duration": 1,"prompt_eval_count": prompttokens,"prompt_eval_duration": 1,"eval_count": comptokens,"eval_duration": 1}
         elif api_format == 8: #oai-responses
             resp_id = f"resp-A{genparams.get('oai_uniqueid', 1)}"
             output_item_id = f"msg_0{genparams.get('oai_uniqueid', 1)}"
@@ -5365,6 +5407,10 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
         self.wfile.write(f'data: {data}\n\n'.encode())
         self.wfile.flush()
 
+    async def send_ollama_stream_event(self, data):
+        self.wfile.write(f'{data}\n'.encode())
+        self.wfile.flush()
+
     async def handle_sse_stream(self, genparams, api_format):
         global friendlymodelname, currfinishreason, thinkformats, tool_call_pairs, cached_chat_template
         global autoswapmode, textName, sttName, ttsName, embedName, musicName, imageName, mmprojName
@@ -5381,7 +5427,8 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
         self.send_header("X-Accel-Buffering", "no")
         self.send_header("cache-control", "no-cache")
         self.send_header("connection", "keep-alive")
-        self.end_headers(content_type='text/event-stream')
+        stream_content_type = 'application/x-ndjson' if api_format == 6 or api_format == 7 else 'text/event-stream'
+        self.end_headers(content_type=stream_content_type)
 
         # if tools, do not send anything else - OAI tool calls will be handled with fakestreaming!
         # only exception is if we know the exact toolcall tag to segment!
@@ -5393,7 +5440,7 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
                 tool_segment_tag = start
                 break
         jinjatools = (args.jinja and args.jinja_tools)
-        if (api_format == 4 or api_format == 9) and using_openai_tools:
+        if (api_format == 4 or api_format == 7 or api_format == 9) and using_openai_tools:
             if not jinjatools or not tool_segment_tag:
                 genparams['sync_toolcall_stream_ineligible'] = True
                 return
@@ -5487,9 +5534,22 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
 
                         sync_potential_toolcall_splitmatch = ""
                         if tokenStr!="" or streamDone:
+                            if (api_format == 4 or api_format == 7 or api_format == 9) and using_openai_tools and tool_segment_tag and not streamDone and not genparams.get("sync_toolcall_potential_triggered", False) and tool_segment_tag not in tokenStr:
+                                tail = ""
+                                for n in range(1, len(tool_segment_tag)):
+                                    prefix = tool_segment_tag[:n]
+                                    if tokenStr.endswith(prefix) and len(prefix) > len(tail):
+                                        tail = prefix
+                                if tail:
+                                    tokenReserve += tail
+                                    tokenStr = tokenStr[:-len(tail)]
+                                    if tokenStr == "":
+                                        await asyncio.sleep(async_sleep_short)
+                                        continue
+
                             # Tool boundary detection for tool-capable chat completions.
                             # if triggered, stop real streaming, and let the buffered fakestreaming take over
-                            if (api_format == 4 or api_format == 9) and using_openai_tools:
+                            if (api_format == 4 or api_format == 7 or api_format == 9) and using_openai_tools:
                                 tokenStr = tokenReserve + tokenStr
                                 tokenReserve = ""
                                 if tool_segment_tag in tokenStr:
@@ -5599,6 +5659,16 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
                                 elif api_format == 3:  # non chat completions
                                     event_str = json.dumps({"id":cmpl_id,"object":"text_completion","created":int(time.time()),"model":modelNameToReturn,"choices":[{"index":0,"finish_reason":None,"text":tokenStr}]})
                                     await self.send_oai_sse_event(event_str)
+                                elif api_format == 6 or api_format == 7:
+                                    created_at = str(datetime.now(timezone.utc).isoformat())
+                                    ollama_content = ""
+                                    if api_format == 6:
+                                        event_str = json.dumps({"model":modelNameToReturn,"created_at":created_at,"response":tokenStr,"done":False})
+                                    else:
+                                        ollama_content = delta.get("content", tokenStr) if delta else tokenStr
+                                        event_str = json.dumps({"model":modelNameToReturn,"created_at":created_at,"message":{"role":"assistant","content":ollama_content},"done":False})
+                                    if api_format == 6 or ollama_content:
+                                        await self.send_ollama_stream_event(event_str)
                                 elif api_format == 9:
                                     if anthropic_first_loop:
                                         await self.send_anthropic_sse_event("message_start", json.dumps({"type":"message_start","message":{"type":"message","id":f"msg_A{req_id_suffix}","role":"assistant","model":modelNameToReturn,"usage":{"input_tokens":prompttokens,"output_tokens":0}}}))
@@ -5642,6 +5712,28 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler):
                                     await self.send_oai_sse_event(addonstr)
                                 event_str = json.dumps({"id":cmpl_id,"object":"text_completion","created":int(time.time()),"model":modelNameToReturn,"choices":[{"index":0,"finish_reason":currfinishreason,"text":tokenStr}]})
                                 await self.send_oai_sse_event(event_str)
+                            elif api_format == 6 or api_format == 7: # Ollama newline-delimited JSON streaming
+                                created_at = str(datetime.now(timezone.utc).isoformat())
+                                ollama_content = delta.get("content", tokenStr) if api_format == 7 and delta else tokenStr
+                                if tokenStr and (api_format == 6 or ollama_content):
+                                    if api_format == 6:
+                                        event_str = json.dumps({"model":modelNameToReturn,"created_at":created_at,"response":tokenStr,"done":False})
+                                    else:
+                                        event_str = json.dumps({"model":modelNameToReturn,"created_at":created_at,"message":{"role":"assistant","content":ollama_content},"done":False})
+                                    await self.send_ollama_stream_event(event_str)
+                                if streamDone:
+                                    prompttokens = batch_final_result.prompt_tokens if using_batch_stream else handle.get_last_input_count()
+                                    completion_tokens = current_token
+                                    final_chunk = {"model":modelNameToReturn,"created_at":created_at,"done":True,"done_reason":currfinishreason,"total_duration":1,"load_duration":1,"prompt_eval_count":prompttokens,"prompt_eval_duration":1,"eval_count":completion_tokens,"eval_duration":1}
+                                    if api_format == 6:
+                                        finalraw = handle.batch_generate_pending_output(batch_request_id) if using_batch_stream else handle.get_pending_output()
+                                        finaltxt = finalraw.decode("UTF-8", "ignore")
+                                        oldprompt = genparams.get('ollamabodyprompt', "")
+                                        final_chunk["response"] = ""
+                                        final_chunk["context"] = tokenize_ids(oldprompt+finaltxt,False)
+                                    else:
+                                        final_chunk["message"] = {"role":"assistant","content":""}
+                                    await self.send_ollama_stream_event(json.dumps(final_chunk))
                             elif api_format == 8: #oai-responses
                                 resp_id = f"resp-A{genparams.get('oai_uniqueid', 1)}"
                                 item_id = f"msg_0{genparams.get('oai_uniqueid', 1)}"
@@ -6026,7 +6118,7 @@ Change Mode<br>
     def do_GET(self):
         global embedded_kailite, embedded_kcpp_docs, embedded_kcpp_sdui, embedded_kailite_gz, embedded_kcpp_docs_gz, embedded_kcpp_sdui_gz, embedded_lcpp_ui_gz, embedded_musicui, embedded_musicui_gz
         global last_req_time, start_time, cached_chat_template, cached_sd_info, has_vision_support, has_audio_support, has_whisper, friendlymodelname
-        global savedata_obj, has_multiplayer, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, maxctx, maxhordelen, friendlymodelname, lastuploadedcomfyimg, lastgeneratedcomfyimg, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey, friendlysdmodelname, fullsdmodelpath, password, friendlyembeddingsmodelname, voicelist
+        global savedata_obj, has_multiplayer, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, maxctx, maxhordelen, friendlymodelname, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey, friendlysdmodelname, fullsdmodelpath, password, friendlyembeddingsmodelname, voicelist
         global autoswapmode, textName, sttName, ttsName, embedName, musicName, imageName, mmprojName
 
         clean_path = self.path.split("?")[0] #for cases where we do not want query params
@@ -6219,8 +6311,6 @@ Change Mode<br>
                 response_body = (json.dumps([{"name":name,"label":name} for name in cached_sd_info.get('available_schedulers', [])]).encode())
         elif clean_path.endswith('/sdapi/v1/latent-upscale-modes'):
            response_body = (json.dumps([]).encode())
-        elif clean_path.endswith('/sdapi/v1/upscalers'):
-           response_body = (json.dumps([]).encode())
 
         #vits compatible
         elif clean_path=='/voice/check':
@@ -6253,7 +6343,7 @@ Change Mode<br>
             modelNameToReturn = friendlymodelname
             if autoswapmode and textName is not None:
                 modelNameToReturn = textName
-            response_body = (json.dumps({"models":[{"name":"koboldcpp","model":f"{modelNameToReturn}:latest","modified_at":"2024-07-19T15:26:55.6122841+08:00","expires_at": "2055-06-04T19:06:25.5433636+08:00","size":394998579,"size_vram":394998579,"digest":"b5dc5e784f2a3ee1582373093acf69a2f4e2ac1710b253a001712b86a61f88bb","details":{"parent_model":"","format":"gguf","family":"koboldcpp","families":["koboldcpp"],"parameter_size":"128M","quantization_level":"Q4_0"}},{"name":"koboldcpp","model":modelNameToReturn,"modified_at":"2025-01-01T01:00:00.0000000+00:00","expires_at": "2069-01-01T01:00:00.0000000+00:00","size":394998579,"size_vram":394998579,"digest":"b5dc5e784f2a3ee1582373093acf69a2f4e2ac1710b253a001712b86a61f88bb","details":{"parent_model":"","format":"gguf","family":"koboldcpp","families":["koboldcpp"],"parameter_size":"128M","quantization_level":"Q4_0"}}]}).encode())
+            response_body = (json.dumps({"models":[{"name":f"{modelNameToReturn}","model":f"{modelNameToReturn}:latest","modified_at":"2024-07-19T15:26:55.6122841+08:00","expires_at": "2055-06-04T19:06:25.5433636+08:00","size":394998579,"size_vram":394998579,"digest":"b5dc5e784f2a3ee1582373093acf69a2f4e2ac1710b253a001712b86a61f88bb","details":{"parent_model":"","format":"gguf","family":"koboldcpp","families":["koboldcpp"],"parameter_size":"128M","quantization_level":"Q4_0"}},{"name":modelNameToReturn,"model":modelNameToReturn,"modified_at":"2025-01-01T01:00:00.0000000+00:00","expires_at": "2069-01-01T01:00:00.0000000+00:00","size":394998579,"size_vram":394998579,"digest":"b5dc5e784f2a3ee1582373093acf69a2f4e2ac1710b253a001712b86a61f88bb","details":{"parent_model":"","format":"gguf","family":"koboldcpp","families":["koboldcpp"],"parameter_size":"128M","quantization_level":"Q4_0"}}]}).encode())
         elif clean_path.endswith('/api/version'): #ollama compatible, NOT the kcpp version
             response_body = (json.dumps({"version":"0.7.0"}).encode())
         elif clean_path=='/ping':
@@ -6279,6 +6369,15 @@ Change Mode<br>
         elif clean_path=='/view' or clean_path=='/view.png' or clean_path=='/api/view' or clean_path.startswith('/view_image'): #emulate comfyui
             content_type = 'image/png'
             response_body = lastgeneratedcomfyimg
+        elif clean_path.startswith('/sdapi/v1/get_last.png'):
+            parsed_url = urllib.parse.urlparse(self.path)
+            parsed_dict = urllib.parse.parse_qs(parsed_url.query)
+            genkey = parsed_dict.get('genkey', [''])[0]
+            if genkey and genkey==lastgeneratedcachedimgkey and lastgeneratedcachedimg:
+                content_type = 'image/png'
+                response_body = lastgeneratedcachedimg
+            else:
+                response_body = None
         elif clean_path=='/history' or clean_path=='/api/history' or clean_path.startswith('/api/history/') or clean_path.startswith('/history/'): #emulate comfyui
             modelNameToReturn = friendlysdmodelname
             if autoswapmode and imageName is not None:
@@ -6414,7 +6513,7 @@ Change Mode<br>
 
     def do_POST(self):
         global thinkformats
-        global modelbusy, batched_request_runner_count, requestsinqueue, currentusergenkey, totalgens, pendingabortkey, lastuploadedcomfyimg, lastgeneratedcomfyimg, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, net_save_slots, has_vision_support, savestate_limit, mcp_lock
+        global modelbusy, batched_request_runner_count, requestsinqueue, currentusergenkey, totalgens, pendingabortkey, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, net_save_slots, has_vision_support, savestate_limit, mcp_lock
         global autoswapmode, textName, sttName, ttsName, embedName, musicName, imageName, mmprojName
         contlenstr = self.headers['content-length']
         content_length = 0
@@ -6910,6 +7009,7 @@ Change Mode<br>
             is_transcribe = False
             is_tts = False
             is_embeddings = False
+            is_ollama_embeddings = False
             is_music_codes = False
             is_music_audio = False
             response_body = None
@@ -7014,8 +7114,9 @@ Change Mode<br>
                 is_transcribe = True
             elif clean_path.endswith('/api/extra/tts') or clean_path.endswith('/v1/audio/speech') or clean_path=="/audio/speech" or clean_path.endswith('/tts_to_audio'):
                 is_tts = True
-            elif clean_path.endswith('/api/extra/embeddings') or clean_path.endswith('/v1/embeddings'):
+            elif clean_path.endswith('/api/extra/embeddings') or clean_path.endswith('/v1/embeddings') or clean_path=="/api/embed":
                 is_embeddings = True
+                is_ollama_embeddings = (clean_path=="/api/embed")
             elif clean_path.endswith('/api/extra/music/prepare'):
                 is_music_codes = True
             elif clean_path.endswith('/api/extra/music/generate'):
@@ -7114,6 +7215,8 @@ Change Mode<br>
                     # Check if streaming chat completions, if so, set stream mode to true
                     if (api_format == 4 or api_format == 3 or api_format == 8 or api_format == 9) and "stream" in genparams and genparams["stream"]:
                         sse_stream_flag = True
+                    if (api_format == 6 or api_format == 7) and genparams.get('stream', True):
+                        sse_stream_flag = True
                     if continuous_batching_python_eligible(genparams, api_format):
                         genparams['_batch_expected'] = True
                         modelbusy.release()
@@ -7128,41 +7231,13 @@ Change Mode<br>
                         if autoswapmode and textName is not None:
                             modelNameToReturn = textName
                         # Headers are already sent when streaming
-                        if (api_format == 6 or api_format == 7) and genparams.get('stream', True):
-                            #ollama fake streaming
-                            self.send_response(200)
-                            self.send_header("X-Accel-Buffering", "no")
-                            self.send_header("cache-control", "no-cache")
-                            self.send_header("connection", "keep-alive")
-                            self.end_headers(content_type='text/event-stream')
-                            if api_format == 6:
-                                bodytxt = gendat.get("response","") # extract and erase the AI response from the sync payload.
-                                gendat["response"] = ""
-                                pl = {"model":modelNameToReturn,"created_at":str(datetime.now(timezone.utc).isoformat()),"response":bodytxt,"done":False}
-                                self.wfile.write(f'{json.dumps(pl)}\n'.encode())
-                                self.wfile.flush()
-                                time.sleep(0.05) #short delay
-                                self.wfile.write(f'{json.dumps(gendat)}\n'.encode()) # note: gendat already contains done=true and empty response
-                                self.wfile.flush()
-                                time.sleep(0.05) #short delay
-                            else:
-                                bodytxt = gendat.get("message",{}).get("content","") # extract and erase the AI response from the sync payload.
-                                gendat["message"] = {"role":"assistant","content":""}
-                                pl = {"model":modelNameToReturn,"created_at":str(datetime.now(timezone.utc).isoformat()),"message":{"role":"assistant","content":bodytxt},"done":False}
-                                self.wfile.write(f'{json.dumps(pl)}\n'.encode())
-                                self.wfile.flush()
-                                time.sleep(0.05) #short delay
-                                self.wfile.write(f'{json.dumps(gendat)}\n'.encode()) # note: gendat already contains done=true and empty response
-                                self.wfile.flush()
-                                time.sleep(0.05) #short delay
-                            self.close_connection = True
-                        elif not sse_stream_flag:
+                        if not sse_stream_flag:
                             self.send_response(200)
                             genresp = (json.dumps(gendat).encode())
                             self.send_header('content-length', str(len(genresp)))
                             self.end_headers(content_type='application/json')
                             self.wfile.write(genresp)
-                        elif (api_format == 4 or api_format == 9) and genparams.get('using_openai_tools', False): #special case, fake streaming for openai tool calls
+                        elif (api_format == 4 or api_format == 7 or api_format == 9) and genparams.get('using_openai_tools', False): #special case, fake streaming for tool calls
                             # we only send content_text and reasoning_text if tools aren't used. they contain the balance of the output after sync_toolcall_potential_triggered was triggered
                             content_text = genparams.get('sync_toolcall_extra_content', "") #populated by the sse call, we don't use gendat['choices'][0]['message'].get('content', None)
                             reasoning_text = genparams.get('sync_toolcall_extra_reasoning_content', "")
@@ -7172,6 +7247,8 @@ Change Mode<br>
                                     toolsdata_res = gendat['choices'][0]['message']['tool_calls']
                                     if toolsdata_res and len(toolsdata_res)>0:
                                         toolsdata_res[0]["index"] = 0 # need to add an index for OWUI
+                                elif api_format == 7:
+                                    toolsdata_res = gendat.get("message", {}).get("tool_calls", [])
                                 elif api_format == 9:
                                     # gendat["content"] is a list of Anthropic content blocks; pull out the tool_use ones and reformat to OAI shape for the shared emission code
                                     for block in gendat.get("content", []):
@@ -7187,7 +7264,23 @@ Change Mode<br>
                             except Exception:
                                 toolsdata_res = []
 
-                            if api_format == 9: # Anthropic fake-stream for tool calls
+                            if api_format == 7: # Ollama fake-stream for tool calls
+                                created_at = str(datetime.now(timezone.utc).isoformat())
+                                if not content_text and genparams.get('sync_toolcall_stream_ineligible', False):
+                                    content_text = gendat.get("message", {}).get("content", "")
+                                if content_text or toolsdata_res:
+                                    chunk_msg = {"role":"assistant","content":"" if toolsdata_res else (content_text or "")}
+                                    if toolsdata_res:
+                                        chunk_msg["tool_calls"] = toolsdata_res
+                                    chunk = {"model":modelNameToReturn,"created_at":created_at,"message":chunk_msg,"done":False}
+                                    self.wfile.write(f'{json.dumps(chunk)}\n'.encode())
+                                    self.wfile.flush()
+                                final_msg = {"role":"assistant","content":""}
+                                final_chunk = {"model":modelNameToReturn,"created_at":created_at,"message":final_msg,"done":True,"done_reason":gendat.get("done_reason", currfinishreason),"total_duration":gendat.get("total_duration", 1),"load_duration":gendat.get("load_duration", 1),"prompt_eval_count":gendat.get("prompt_eval_count", handle.get_last_input_count()),"prompt_eval_duration":gendat.get("prompt_eval_duration", 1),"eval_count":gendat.get("eval_count", handle.get_last_token_count()),"eval_duration":gendat.get("eval_duration", 1)}
+                                self.wfile.write(f'{json.dumps(final_chunk)}\n'.encode())
+                                self.wfile.flush()
+
+                            elif api_format == 9: # Anthropic fake-stream for tool calls
                                 req_id_suffix = genparams.get('oai_uniqueid', 1)
                                 start_msg = {"type": "message", "id": f"msg_A{req_id_suffix}", "role": "assistant", "model": modelNameToReturn, "usage": {"input_tokens": 0, "output_tokens": 0}}
                                 self.wfile.write(f'event: message_start\ndata: {json.dumps({"type":"message_start","message":start_msg})}\n\n'.encode())
@@ -7381,6 +7474,8 @@ Change Mode<br>
                     return
                 elif is_imggen: #image gen
                     try:
+                        lastgeneratedcachedimg = b''
+                        lastgeneratedcachedimgkey = ''
                         if is_comfyui_imggen:
                             lastgeneratedcomfyimg = b''
                             genparams = sd_comfyui_tranform_params(genparams)
@@ -7399,6 +7494,11 @@ Change Mode<br>
                         genfinalframe = gen["final_frame"]
                         geninfo = json.dumps(gen["info"]) # sdapi really expects a stringified JSON
                         genresp = None
+                        if gendat:
+                            lastgeneratedcachedimg = base64.b64decode(gendat)
+                            lastgeneratedcachedimgkey = genparams.get('genkey', '')
+                        else:
+                            lastgeneratedcachedimg = b''
                         if is_comfyui_imggen:
                             if gendat:
                                 lastgeneratedcomfyimg = base64.b64decode(gendat)
@@ -7475,17 +7575,20 @@ Change Mode<br>
                         if autoswapmode and embedName is not None:
                             modelNameToReturn = embedName
                         gendat = embeddings_generate(genparams)
-                        outdatas = []
-                        odidx = 0
-                        for od in gendat["data"]:
-                            if genparams.get("encoding_format", "")=="base64":
-                                binary_data = struct.pack('<' + 'f' * len(od), *od)
-                                b64_string = base64.b64encode(binary_data).decode('utf-8')
-                                outdatas.append({"object":"embedding","index":odidx,"embedding":b64_string})
-                            else:
-                                outdatas.append({"object":"embedding","index":odidx,"embedding":od})
-                            odidx += 1
-                        genresp = (json.dumps({"object":"list","data":outdatas,"model":modelNameToReturn,"usage":{"prompt_tokens":gendat["count"],"total_tokens":gendat["count"]}}).encode())
+                        if is_ollama_embeddings:
+                            genresp = (json.dumps({"model":modelNameToReturn,"embeddings":gendat["data"],"total_duration":1,"load_duration":1,"prompt_eval_count":gendat["count"]}).encode())
+                        else:
+                            outdatas = []
+                            odidx = 0
+                            for od in gendat["data"]:
+                                if genparams.get("encoding_format", "")=="base64":
+                                    binary_data = struct.pack('<' + 'f' * len(od), *od)
+                                    b64_string = base64.b64encode(binary_data).decode('utf-8')
+                                    outdatas.append({"object":"embedding","index":odidx,"embedding":b64_string})
+                                else:
+                                    outdatas.append({"object":"embedding","index":odidx,"embedding":od})
+                                odidx += 1
+                            genresp = (json.dumps({"object":"list","data":outdatas,"model":modelNameToReturn,"usage":{"prompt_tokens":gendat["count"],"total_tokens":gendat["count"]}}).encode())
                         self.send_response(200)
                         self.send_header('content-length', str(len(genresp)))
                         self.end_headers(content_type='application/json')
@@ -7550,9 +7653,16 @@ Change Mode<br>
         self.end_headers(content_type='text/html')
 
     def end_headers(self, content_type=None):
-        self.send_header('access-control-allow-origin', '*')
+        origin = self.headers.get('Origin')
+        if origin:
+            self.send_header('access-control-allow-origin', origin)
+            self.send_header('access-control-allow-credentials', 'true')
+            self.send_header('vary', 'Origin')
+        else:
+            self.send_header('access-control-allow-origin', '*')
         self.send_header('access-control-allow-methods', '*')
         self.send_header('access-control-allow-headers', '*, Accept, Content-Type, Content-Length, Cache-Control, Accept-Encoding, X-CSRF-Token, Client-Agent, X-Fields, Content-Type, Authorization, X-Requested-With, X-HTTP-Method-Override, apikey, genkey')
+        self.send_header('access-control-allow-private-network', 'true')
         self.send_header("cache-control", "no-store")
         if content_type is not None:
             self.send_header('content-type', content_type)
@@ -9169,7 +9279,7 @@ def show_gui():
     makecheckbox(admin_tab, "Enable Model Administration", admin_var, 1, 0, command=toggleadmin,tooltiptxt="Enable a admin server, allowing you to remotely relaunch and swap models and configs.")
     makelabelentry(admin_tab, "Admin Password:" , admin_password_var, 3, 150,padx=(120),singleline=True,tooltip="Require a password to access admin functions. You are strongly advised to use one for publically accessible instances!")
     makefileentry(admin_tab, "Config Directory (Required):", "Select directory containing .gguf or .kcpps files to relaunch from", admin_dir_var, 5, width=280, dialog_type=2, tooltiptxt="Specify a directory to look for .kcpps configs in, which can be used to swap models.")
-    makefileentry(admin_tab, "Base config .kcpps (For reloading):", "", baseconfig_var, 7, width=280, dialog_type=0, tooltiptxt="Specify a base .kcpps config to apply, if no custom base config is selected during a model swap.")
+    makefileentry(admin_tab, "Base config .kcpps (Optional, for reloading):", "", baseconfig_var, 7, width=280, dialog_type=0, tooltiptxt="Specify a base .kcpps config to apply, if no custom base config is selected during a model swap.")
     makelabelentry(admin_tab, "Auto Unload Timeout:" , admin_unload_timeout_var, 17, 70,padx=(150),singleline=True,tooltip="Set an idle timeout in seconds after which KoboldCpp will automatically unload the current model.")
     makecheckbox(admin_tab, "SingleInstance Mode", singleinstance_var, 19, 0,tooltiptxt="Allows this server to be shut down by another KoboldCpp instance with singleinstance starting on the same port.")
     router_mode_box = makecheckbox(admin_tab, "Router Mode", router_mode_var, 21, 0, command=togglerouter, tooltiptxt="Router mode uses a reverse proxy router, allowing you to easily hotswap models and configs within a single request. Requires admin mode.")
@@ -10502,19 +10612,32 @@ def sanitize_string(input_string):
     return sanitized_string
 
 def resolve_huggingface_xet_url(input_url):
+    global nocertify
     if "https://huggingface.co/" not in input_url or "/resolve/" not in input_url:
         return input_url
-    try:
+
+    ssl_cert_dir = os.environ.get('SSL_CERT_DIR')
+    if not ssl_cert_dir and not nocertify and os.name != 'nt':
+        os.environ['SSL_CERT_DIR'] = '/etc/ssl/certs'
+
+    def resolve_with_context(ssl_context=None):
         req = urllib.request.Request(input_url, headers={'User-Agent': 'Mozilla/5.0'}, method="HEAD")
-        with urllib.request.urlopen(req, timeout=10) as response:
-            resolved_url = response.geturl()
-    except Exception:
         try:
+            with urllib.request.urlopen(req, timeout=10, context=ssl_context) as response:
+                return response.geturl()
+        except Exception:
             req = urllib.request.Request(input_url, headers={'User-Agent': 'Mozilla/5.0'})
-            with urllib.request.urlopen(req, timeout=10) as response:
-                resolved_url = response.geturl()
+            with urllib.request.urlopen(req, timeout=10, context=ssl_context) as response:
+                return response.geturl()
+
+    try:
+        resolved_url = resolve_with_context()
+    except Exception as first_error:
+        try:
+            import ssl
+            resolved_url = resolve_with_context(ssl._create_unverified_context())
         except Exception as e:
-            print(f"Could not pre-resolve Hugging Face URL, using original URL: {e}")
+            print(f"Could not pre-resolve Hugging Face URL, using original URL: {first_error}; {e}")
             return input_url
     # resolved_host = urllib.parse.urlparse(resolved_url).netloc.lower()
     if ("xet-bridge" in resolved_url) and resolved_url != input_url:
@@ -11511,7 +11634,7 @@ def kcpp_main_process(launch_args, g_memory=None, gui_launcher=False):
         global maxctx
         maxctx = args.contextsize
 
-    args.defaultgenamt = max(64, min(args.defaultgenamt, 16384))
+    args.defaultgenamt = max(64, min(args.defaultgenamt, 32768))
     args.defaultgenamt = min(args.defaultgenamt, maxctx / 2)
 
     #this uses the true port instead of the displayport, because we dont want to shut down a router
@@ -11760,8 +11883,8 @@ def kcpp_main_process(launch_args, g_memory=None, gui_launcher=False):
             friendlysdmodelname = os.path.basename(imgmodel)
             friendlysdmodelname = os.path.splitext(friendlysdmodelname)[0]
             friendlysdmodelname = sanitize_string(friendlysdmodelname)
-            loadok = sd_load_model(imgmodel,imgvae,imgt5xxl,imgclip1,imgclip2,imgphotomaker,imgupscaler,imgaudiovae)
             cached_sd_info = sd_get_info()
+            loadok = sd_load_model(imgmodel,imgvae,imgt5xxl,imgclip1,imgclip2,imgphotomaker,imgupscaler,imgaudiovae)
             print("Load Image Model OK: " + str(loadok))
             if not loadok:
                 exitcounter = 999
@@ -12217,7 +12340,7 @@ if __name__ == '__main__':
     modelgroup.add_argument("--model","-m", metavar=('[filenames]'), help="Model file to load. Accepts multiple values if they are URLs.", type=str, nargs='+', default=[])
     modelgroup.add_argument("model_param", help="Model file to load (positional)", nargs="?")
     parser.add_argument("--config", metavar=('[filename]'), help="Load settings from a .kcpps file. Other arguments will be ignored", type=str, nargs=1)
-    parser.add_argument("--contextsize","--ctx-size", "-c", help=f"Controls the memory allocated for maximum context size, only change if you need more RAM for big contexts. (default {default_maxctx}).",metavar=('[256 to 262144]'), type=check_range(int,256,262144), default=default_maxctx)
+    parser.add_argument("--contextsize","--ctx-size", "-c", help=f"Controls the memory allocated for maximum context size, only change if you need more RAM for big contexts. (default {default_maxctx}).",metavar=('[256 to 524288]'), type=check_range(int,256,524288), default=default_maxctx)
     parser.add_argument("--gpulayers","--gpu-layers","--n-gpu-layers","-ngl", help="Set number of layers to offload to GPU (when using GPU). Set to -1 to enable autofit (default), set to 0 to disable GPU offload.",metavar=('[GPU layers]'), nargs='?', const=1, type=int, default=-1)
     parser.add_argument("--host", metavar=('[ipaddr]'), help="Host IP to listen on. If this flag is not set, all routable interfaces are accepted.", default="")
     parser.add_argument("--launch", help="Launches a web browser when load is completed.", action='store_true')
@@ -12239,7 +12362,7 @@ if __name__ == '__main__':
     advparser.add_argument("--chatcompletionsadapter", metavar=('[filename]'), help="Select an optional ChatCompletions Adapter JSON file to force custom instruct tags.", default="AutoGuess")
     advparser.add_argument("--cli", help="Does not launch KoboldCpp HTTP server. Instead, enables KoboldCpp from the command line, accepting interactive console input and displaying responses to the terminal.", action='store_true')
     advparser.add_argument("--debugmode", help="Shows additional debug info in the terminal. Levels: -1 (Horde-quiet, suppresses non-essential prints; auto-applied when Horde args are set), 0 (default, normal output), 1 (verbose: extra slot/cache info, larger print buffers, retains horde-debug prefix). Passing the flag without a value implies 1.", nargs='?', const=1, type=int, default=0)
-    advparser.add_argument("--defaultgenamt", help="How many tokens to generate by default, if not specified. Must be smaller than context size. Usually, your frontend GUI will override this.", type=check_range(int,64,16384), default=default_genlen)
+    advparser.add_argument("--defaultgenamt", help="How many tokens to generate by default, if not specified. Must be smaller than context size. Usually, your frontend GUI will override this.", type=check_range(int,64,32768), default=default_genlen)
     advparser.add_argument("--device", "-dev", metavar=('<dev1,dev2,..>'), help="Set llama.cpp compatible device selection override. Comma separated. Overrides normal device choices.", default="")
     advparser.add_argument("--downloaddir", metavar=('[directory]'), help="Specify a directory that models will be downloaded to or searched from, if unset uses the working directory.", default="")
     advparser.add_argument("--draftamount","--draft-max","--draft-n","--spec-draft-n-max", metavar=('[tokens]'), help="How many tokens to draft per chunk before verifying results", type=int, default=default_draft_amount)

model_adapter.cpp

@@ -486,6 +486,16 @@ bool useSmartContext, const bool requireFullSubset, const int minimum_to_proceed
 
     for (int i = 0; i < cur_ctx_len; ++i)
     {
+        if (i >= embd_inp_len)
+        {
+            if(requireFullSubset)
+            {
+                last_n_tokens.erase(last_n_tokens.end() - n_past, last_n_tokens.end());
+                n_past = 0;
+                fastforwardok = false;
+            }
+            break;
+        }
         if (current_context_tokens[i] == embd_inp[i])
         {
             n_past += 1;

otherarch/sdcpp/examples/cli/main.cpp

@@ -1,6 +1,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <time.h>
+#include <algorithm>
 #include <cctype>
 #include <filesystem>
 #include <functional>
@@ -53,6 +54,9 @@ struct SDCliParams {
     bool metadata_brief      = false;
     bool metadata_all        = false;
 
+    std::string imatrix_out;
+    std::vector<std::string> imatrix_in;
+
     bool normal_exit = false;
 
     ArgOptions get_options() {
@@ -79,6 +83,11 @@ struct SDCliParams {
              "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
              0,
              &preview_path},
+            {"",
+             "--imat-out",
+             "compute the imatrix for this run and save it to the provided path",
+             0,
+             &imatrix_out},
         };
 
         options.int_options = {
@@ -179,6 +188,14 @@ struct SDCliParams {
             return -1;
         };
 
+        auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) {
+            if (++index >= argc) {
+                return -1;
+            }
+            imatrix_in.push_back(argv[index]);
+            return 1;
+        };
+
         options.manual_options = {
             {"-M",
              "--mode",
@@ -192,6 +209,10 @@ struct SDCliParams {
              "--help",
              "show this help message and exit",
              on_help_arg},
+            {"",
+             "--imat-in",
+             "load an imatrix file for quantization or continued collection; can be specified multiple times",
+             on_imatrix_in_arg},
         };
 
         return options;
@@ -253,6 +274,7 @@ struct SDCliParams {
             << "  preview_fps: " << preview_fps << ",\n"
             << "  taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
             << "  preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
+            << "  imatrix_out: \"" << imatrix_out << "\",\n"
             << "  metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
             << "  metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
             << "  metadata_all: " << (metadata_all ? "true" : "false") << "\n"
@@ -459,7 +481,8 @@ bool save_results(const SDCliParams& cli_params,
         if (!img.data)
             return false;
 
-        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
+        int images_per_batch        = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1;
+        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch;
         std::string params          = gen_params.embed_image_metadata
                                           ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
                                           : "";
@@ -605,13 +628,33 @@ int main(int argc, const char* argv[]) {
     LOG_DEBUG("%s", ctx_params.to_string().c_str());
     LOG_DEBUG("%s", gen_params.to_string().c_str());
 
+    if (!cli_params.imatrix_out.empty()) {
+        if (fs::exists(cli_params.imatrix_out) &&
+            std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) {
+            LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str());
+        }
+        enable_imatrix_collection();
+    }
+
+    for (const auto& in_file : cli_params.imatrix_in) {
+        LOG_INFO("loading imatrix from '%s'", in_file.c_str());
+        if (!load_imatrix(in_file.c_str())) {
+            LOG_WARN("failed to load imatrix from '%s'", in_file.c_str());
+        }
+    }
+
     if (cli_params.mode == CONVERT) {
-        bool success = convert(ctx_params.model_path.c_str(),
-                               ctx_params.vae_path.c_str(),
-                               cli_params.output_path.c_str(),
-                               ctx_params.wtype,
-                               ctx_params.tensor_type_rules.c_str(),
-                               cli_params.convert_name);
+        bool success = convert_with_components(ctx_params.model_path.c_str(),
+                                               ctx_params.clip_l_path.c_str(),
+                                               ctx_params.clip_g_path.c_str(),
+                                               ctx_params.t5xxl_path.c_str(),
+                                               ctx_params.diffusion_model_path.c_str(),
+                                               ctx_params.vae_path.c_str(),
+                                               cli_params.output_path.c_str(),
+                                               ctx_params.wtype,
+                                               ctx_params.tensor_type_rules.c_str(),
+                                               cli_params.convert_name,
+                                               ctx_params.n_threads);
         if (!success) {
             LOG_ERROR("convert '%s'/'%s' to '%s' failed",
                       ctx_params.model_path.c_str(),
@@ -766,8 +809,12 @@ int main(int argc, const char* argv[]) {
         if (cli_params.mode == IMG_GEN) {
             sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
 
-            num_results = gen_params.batch_count;
-            results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
+            sd_image_t* generated_images = nullptr;
+            if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
+                generated_images = nullptr;
+                num_results      = 0;
+            }
+            results.adopt(generated_images, num_results);
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
             sd_image_t* generated_video        = nullptr;
@@ -802,12 +849,22 @@ int main(int argc, const char* argv[]) {
                 SDImageOwner current_image(results[i]);
                 results[i] = {0, 0, 0, nullptr};
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
-                    SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
-                    if (upscaled_image.get().data == nullptr) {
+                    sd_image_t* upscaled_images = nullptr;
+                    int upscaled_count          = 0;
+                    bool upscale_ok             = upscale(upscaler_ctx.get(),
+                                                          current_image.get(),
+                                                          upscale_factor,
+                                                          &upscaled_images,
+                                                          &upscaled_count);
+                    if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) {
+                        free_sd_images(upscaled_images, upscaled_count);
                         LOG_ERROR("upscale failed");
                         break;
                     }
-                    current_image = std::move(upscaled_image);
+                    sd_image_t upscaled_image = upscaled_images[0];
+                    upscaled_images[0]        = {0, 0, 0, nullptr};
+                    free_sd_images(upscaled_images, upscaled_count);
+                    current_image.reset(upscaled_image);
                 }
                 results[i] = current_image.release();  // Set the final upscaled image as the result
             }
@@ -819,6 +876,11 @@ int main(int argc, const char* argv[]) {
         return 1;
     }
 
+    if (!cli_params.imatrix_out.empty()) {
+        LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str());
+        save_imatrix(cli_params.imatrix_out.c_str());
+    }
+
     free_sd_audio(generated_audio);
 
     return 0;

otherarch/sdcpp/examples/common/common.cpp

@@ -443,6 +443,12 @@ ArgOptions SDContextParams::get_options() {
          "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
          (int)',',
          &tensor_type_rules},
+        {"",
+         "--model-args",
+         "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
+         "chroma_t5_mask_pad, qwen_image_zero_cond_t",
+         (int)',',
+         &model_args},
         {"",
          "--photo-maker",
          "path to PHOTOMAKER model",
@@ -468,6 +474,13 @@ ArgOptions SDContextParams::get_options() {
          "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
          (int)',',
          &params_backend},
+        {"",
+         "--split-mode",
+         "weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): "
+         "layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). "
+         "Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer",
+         (int)',',
+         &split_mode},
         {"",
          "--rpc-servers",
          "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
@@ -486,10 +499,6 @@ ArgOptions SDContextParams::get_options() {
          "number of threads to use during computation (default: -1). "
          "If threads <= 0, then threads will be set to the number of CPU physical cores",
          &n_threads},
-        {"",
-         "--chroma-t5-mask-pad",
-         "t5 mask pad size of chroma",
-         &chroma_t5_mask_pad},
     };
 
     options.bool_options = {
@@ -501,6 +510,12 @@ ArgOptions SDContextParams::get_options() {
          "--eager-load",
          "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
          true, &eager_load},
+        {"",
+         "--auto-fit",
+         "pick the diffusion/te/vae device placements automatically from the model size and the per-device "
+         "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and "
+         "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)",
+         true, &auto_fit},
         {"",
          "--force-sdxl-vae-conv-scale",
          "force use of conv scale on sdxl vae",
@@ -541,30 +556,6 @@ ArgOptions SDContextParams::get_options() {
          "--vae-conv-direct",
          "use ggml_conv2d_direct in the vae model",
          true, &vae_conv_direct},
-        {"",
-         "--circular",
-         "enable circular padding for convolutions",
-         true, &circular},
-        {"",
-         "--circularx",
-         "enable circular RoPE wrapping on x-axis (width) only",
-         true, &circular_x},
-        {"",
-         "--circulary",
-         "enable circular RoPE wrapping on y-axis (height) only",
-         true, &circular_y},
-        {"",
-         "--chroma-disable-dit-mask",
-         "disable dit mask for chroma",
-         false, &chroma_use_dit_mask},
-        {"",
-         "--qwen-image-zero-cond-t",
-         "enable zero_cond_t for qwen image",
-         true, &qwen_image_zero_cond_t},
-        {"",
-         "--chroma-enable-t5-mask",
-         "enable t5 mask for chroma",
-         true, &chroma_use_t5_mask},
     };
 
     auto on_type_arg = [&](int argc, const char** argv, int index) {
@@ -653,7 +644,7 @@ ArgOptions SDContextParams::get_options() {
          on_sampler_rng_arg},
         {"",
          "--prediction",
-         "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]",
+         "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]",
          on_prediction_arg},
         {"",
          "--lora-apply-mode",
@@ -663,6 +654,18 @@ ArgOptions SDContextParams::get_options() {
          "but it usually offers faster inference speed and, in some cases, lower memory usage. "
          "The at_runtime mode, on the other hand, is exactly the opposite.",
          on_lora_apply_mode_arg},
+        {"",
+         "--list-devices",
+         "list available ggml backend devices (one 'name<TAB>description' per line) and exit; "
+         "the names are the device names accepted by --backend and --params-backend",
+         [](int /*argc*/, const char** /*argv*/, int /*index*/) {
+             size_t device_list_size = sd_list_devices(nullptr, 0);
+             std::vector<char> devices(device_list_size + 1);
+             sd_list_devices(devices.data(), devices.size());
+             fputs(devices.data(), stdout);
+             std::exit(0);
+             return 0;
+         }},
     };
 
     return options;
@@ -710,7 +713,18 @@ bool SDContextParams::resolve(SDMode mode) {
 }
 
 bool SDContextParams::validate(SDMode mode) {
-    if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
+    if (mode == CONVERT) {
+        const bool has_convert_input = model_path.length() != 0 ||
+                                       clip_l_path.length() != 0 ||
+                                       clip_g_path.length() != 0 ||
+                                       t5xxl_path.length() != 0 ||
+                                       diffusion_model_path.length() != 0 ||
+                                       vae_path.length() != 0;
+        if (!has_convert_input) {
+            LOG_ERROR("error: convert mode needs at least one model input path\n");
+            return false;
+        }
+    } else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
         LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n");
         return false;
     }
@@ -807,6 +821,9 @@ std::string SDContextParams::to_string() const {
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
+        << "  split_mode: \"" << split_mode << "\",\n"
+        << "  model_args: \"" << model_args << "\",\n"
+        << "  auto_fit: " << (auto_fit ? "true" : "false") << ",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
         << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
         << "  clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
@@ -815,13 +832,6 @@ std::string SDContextParams::to_string() const {
         << "  diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
         << "  diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
         << "  vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
-        << "  circular: " << (circular ? "true" : "false") << ",\n"
-        << "  circular_x: " << (circular_x ? "true" : "false") << ",\n"
-        << "  circular_y: " << (circular_y ? "true" : "false") << ",\n"
-        << "  chroma_use_dit_mask: " << (chroma_use_dit_mask ? "true" : "false") << ",\n"
-        << "  qwen_image_zero_cond_t: " << (qwen_image_zero_cond_t ? "true" : "false") << ",\n"
-        << "  chroma_use_t5_mask: " << (chroma_use_t5_mask ? "true" : "false") << ",\n"
-        << "  chroma_t5_mask_pad: " << chroma_t5_mask_pad << ",\n"
         << "  prediction: " << sd_prediction_name(prediction) << ",\n"
         << "  lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n"
         << "  force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n"
@@ -874,20 +884,17 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.tae_preview_only                = taesd_preview;
     sd_ctx_params.diffusion_conv_direct           = diffusion_conv_direct;
     sd_ctx_params.vae_conv_direct                 = vae_conv_direct;
-    sd_ctx_params.circular_x                      = circular || circular_x;
-    sd_ctx_params.circular_y                      = circular || circular_y;
     sd_ctx_params.force_sdxl_vae_conv_scale       = force_sdxl_vae_conv_scale;
-    sd_ctx_params.chroma_use_dit_mask             = chroma_use_dit_mask;
-    sd_ctx_params.chroma_use_t5_mask              = chroma_use_t5_mask;
-    sd_ctx_params.chroma_t5_mask_pad              = chroma_t5_mask_pad;
-    sd_ctx_params.qwen_image_zero_cond_t          = qwen_image_zero_cond_t;
     sd_ctx_params.vae_format                      = str_to_vae_format(vae_format);
     sd_ctx_params.max_vram                        = max_vram.c_str();
     sd_ctx_params.stream_layers                   = stream_layers;
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
+    sd_ctx_params.split_mode                      = split_mode.c_str();
+    sd_ctx_params.auto_fit                        = auto_fit;
     sd_ctx_params.rpc_servers                     = rpc_servers.c_str();
+    sd_ctx_params.model_args                      = model_args.empty() ? nullptr : model_args.c_str();
     return sd_ctx_params;
 }
 
@@ -960,7 +967,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--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; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
+         "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;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
          (int)',',
          &extra_sample_args},
         {"",
@@ -996,6 +1003,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--batch-count",
          "batch count",
          &batch_count},
+        {"",
+         "--qwen-image-layers",
+         "number of Qwen Image Layered layers; latent/output count is layers + 1 (default: 3)",
+         &qwen_image_layers},
         {"",
          "--video-frames",
          "video frames (default: 1)",
@@ -1062,7 +1073,7 @@ ArgOptions SDGenerationParams::get_options() {
          &sample_params.guidance.slg.layer_end},
         {"",
          "--eta",
-         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
          &sample_params.eta},
         {"",
          "--flow-shift",
@@ -1094,7 +1105,7 @@ ArgOptions SDGenerationParams::get_options() {
          &high_noise_sample_params.guidance.slg.layer_end},
         {"",
          "--high-noise-eta",
-         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
          &high_noise_sample_params.eta},
         {"",
          "--strength",
@@ -1145,6 +1156,18 @@ ArgOptions SDGenerationParams::get_options() {
          "disable auto resize of ref images",
          false,
          &auto_resize_ref_image},
+        {"",
+         "--circular",
+         "enable circular padding on both axes for tileable output",
+         true, &circular},
+        {"",
+         "--circularx",
+         "enable circular padding on x-axis (width) only",
+         true, &circular_x},
+        {"",
+         "--circulary",
+         "enable circular padding on y-axis (height) only",
+         true, &circular_y},
         {"",
          "--disable-image-metadata",
          "do not embed generation metadata on image files",
@@ -1465,17 +1488,17 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
          "--scheduler",
-         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal], default: model-specific",
+         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",
@@ -1816,6 +1839,7 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("width", width);
     load_if_exists("height", height);
     load_if_exists("batch_count", batch_count);
+    load_if_exists("qwen_image_layers", qwen_image_layers);
     load_if_exists("video_frames", video_frames);
     load_if_exists("fps", fps);
     load_if_exists("upscale_repeats", upscale_repeats);
@@ -2240,6 +2264,11 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (qwen_image_layers < 0) {
+        LOG_ERROR("error: qwen_image_layers must be non-negative");
+        return false;
+    }
+
     if (sample_params.sample_steps <= 0) {
         LOG_ERROR("error: the sample_steps must be greater than 0\n");
         return false;
@@ -2406,6 +2435,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.strength              = strength;
     params.seed                  = seed;
     params.batch_count           = batch_count;
+    params.qwen_image_layers     = qwen_image_layers;
     params.control_image         = control_image.get();
     params.control_strength      = control_strength;
     params.pm_params             = pm_params;
@@ -2424,6 +2454,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.hires.upscale_tile_size   = hires_upscale_tile_size;
     params.hires.custom_sigmas       = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
     params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
+    params.circular_x                = circular || circular_x;
+    params.circular_y                = circular || circular_y;
     return params;
 }
 
@@ -2489,6 +2521,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     params.hires.upscale_tile_size   = hires_upscale_tile_size;
     params.hires.custom_sigmas       = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
     params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
+    params.circular_x                = circular || circular_x;
+    params.circular_y                = circular || circular_y;
     return params;
 }
 
@@ -2531,6 +2565,7 @@ std::string SDGenerationParams::to_string() const {
         << "  width: " << width << ",\n"
         << "  height: " << height << ",\n"
         << "  batch_count: " << batch_count << ",\n"
+        << "  qwen_image_layers: " << qwen_image_layers << ",\n"
         << "  init_image_path: \"" << init_image_path << "\",\n"
         << "  end_image_path: \"" << end_image_path << "\",\n"
         << "  mask_image_path: \"" << mask_image_path << "\",\n"

otherarch/sdcpp/examples/common/common.h

@@ -151,6 +151,9 @@ struct SDContextParams {
     bool eager_load             = false;
     std::string backend;
     std::string params_backend;
+    std::string split_mode;
+    std::string model_args;
+    bool auto_fit = false;
     std::string rpc_servers;
     std::string effective_backend;
     std::string effective_params_backend;
@@ -163,16 +166,6 @@ struct SDContextParams {
     bool diffusion_conv_direct = false;
     bool vae_conv_direct       = false;
 
-    bool circular   = false;
-    bool circular_x = false;
-    bool circular_y = false;
-
-    bool chroma_use_dit_mask = true;
-    bool chroma_use_t5_mask  = false;
-    int chroma_t5_mask_pad   = 1;
-
-    bool qwen_image_zero_cond_t = false;
-
     prediction_t prediction           = PREDICTION_COUNT;
     lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
 
@@ -197,6 +190,7 @@ struct SDGenerationParams {
     int width                  = -1;
     int height                 = -1;
     int batch_count            = 1;
+    int qwen_image_layers      = 3;
     int64_t seed               = 42;
     float strength             = 0.75f;
     float control_strength     = 0.9f;
@@ -243,6 +237,10 @@ struct SDGenerationParams {
     int upscale_repeats   = 1;
     int upscale_tile_size = 128;
 
+    bool circular   = false;
+    bool circular_x = false;
+    bool circular_y = false;
+
     bool hires_enabled         = false;
     std::string hires_upscaler = "Latent";
     std::string hires_upscaler_model_path;

otherarch/sdcpp/include/stable-diffusion.h

@@ -54,6 +54,8 @@ enum sample_method_t {
     EULER_CFG_PP_SAMPLE_METHOD,
     EULER_A_CFG_PP_SAMPLE_METHOD,
     EULER_GE_SAMPLE_METHOD,
+    DPMPP2M_SDE_SAMPLE_METHOD,
+    DPMPP2M_SDE_BT_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -71,6 +73,9 @@ enum scheduler_t {
     BONG_TANGENT_SCHEDULER,
     LTX2_SCHEDULER,
     LOGIT_NORMAL_SCHEDULER,
+    FLUX2_SCHEDULER,
+    FLUX_SCHEDULER,
+    BETA_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -80,7 +85,8 @@ enum prediction_t {
     EDM_V_PRED,
     FLOW_PRED,
     FLUX_FLOW_PRED,
-    FLUX2_FLOW_PRED,
+    SEFI_FLOW_PRED,
+    MINIT2I_FLOW_PRED,
     PREDICTION_COUNT
 };
 
@@ -210,20 +216,17 @@ typedef struct {
     bool tae_preview_only;
     bool diffusion_conv_direct;
     bool vae_conv_direct;
-    bool circular_x;
-    bool circular_y;
     bool force_sdxl_vae_conv_scale;
-    bool chroma_use_dit_mask;
-    bool chroma_use_t5_mask;
-    int chroma_t5_mask_pad;
-    bool qwen_image_zero_cond_t;
     enum sd_vae_format_t vae_format;
     const char* max_vram;  // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
     bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
     bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;
+    const char* split_mode;  // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row"
+    bool auto_fit;
     const char* rpc_servers;
+    const char* model_args;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -376,6 +379,9 @@ typedef struct {
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
     sd_hires_params_t hires;
+    int qwen_image_layers;
+    bool circular_x;
+    bool circular_y;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -402,17 +408,22 @@ typedef struct {
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
     sd_hires_params_t hires;
+    bool circular_x;
+    bool circular_y;
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;
+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 void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
+typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data);
 
 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);
 SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
+SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
 SD_API int32_t sd_get_num_physical_cores();
 SD_API const char* sd_get_system_info();
 SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
@@ -453,7 +464,10 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa
 
 SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
 SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
+SD_API bool generate_image(sd_ctx_t* sd_ctx,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           sd_image_t** images_out,
+                           int* num_images_out);
 
 enum sd_cancel_mode_t {
     // Stop the current generation as soon as possible.
@@ -483,9 +497,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         const char* params_backend);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
-SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,
-                          sd_image_t input_image,
-                          uint32_t upscale_factor);
+SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
+                    sd_image_t input_image,
+                    uint32_t upscale_factor,
+                    sd_image_t** images_out,
+                    int* num_images_out);
 
 SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
 
@@ -496,6 +512,18 @@ SD_API bool convert(const char* input_path,
                     const char* tensor_type_rules,
                     bool convert_name);
 
+SD_API bool convert_with_components(const char* model_path,
+                                    const char* clip_l_path,
+                                    const char* clip_g_path,
+                                    const char* t5xxl_path,
+                                    const char* diffusion_model_path,
+                                    const char* vae_path,
+                                    const char* output_path,
+                                    enum sd_type_t output_type,
+                                    const char* tensor_type_rules,
+                                    bool convert_name,
+                                    int n_threads);
+
 SD_API bool preprocess_canny(sd_image_t image,
                              float high_threshold,
                              float low_threshold,
@@ -503,9 +531,20 @@ SD_API bool preprocess_canny(sd_image_t image,
                              float strong,
                              bool inverse);
 
+SD_API bool load_imatrix(const char* imatrix_path);
+SD_API void save_imatrix(const char* imatrix_path);
+SD_API void enable_imatrix_collection(void);
+SD_API void disable_imatrix_collection(void);
+
 SD_API const char* sd_commit(void);
 SD_API const char* sd_version(void);
 
+// List available ggml backend devices, one `name<TAB>description` per line.
+// The names are the device names accepted by the --backend / --params-backend
+// assignment specs. Returns the number of bytes required, excluding the null
+// terminator. Passing nullptr or buffer_size 0 only queries the required size.
+SD_API size_t sd_list_devices(char* buffer, size_t buffer_size);
+
 // for C API, caller needs to call free_sd_images to free the memory after use
 // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT.
 SD_API void free_sd_images(sd_image_t* result_images, int num_images);

otherarch/sdcpp/sdtype_adapter.cpp

@@ -270,35 +270,6 @@ std::string load_gpt_oss_vocab_json()
     return load_embd_file(cache, "embd_res/gpt_oss_vocab_json.embd");
 }
 
-static std::string get_device_override(int value, const char * module = nullptr)
-{
-    std::string device_name;
-    if (value <= -2) {
-        device_name = "CPU";
-    } else if (value >= 0) {
-        size_t gpu_index = static_cast<size_t>(value);
-        if (gpu_index >= ggml_backend_dev_count()) {
-            printf("\nWARNING: device %zu doesn't exist, falling back to default for %s\n",
-                   gpu_index,
-                   module ? module : "the main device");
-        } else {
-            auto dev = ggml_backend_dev_get(gpu_index);
-            device_name = ggml_backend_dev_name(dev);
-        }
-    }
-    std::string result;
-    if (device_name == "") {
-        result = ""; // no override: sdcpp will use the main device
-    } else if (module) {
-        printf("Selecting %s as %s image generation device\n", device_name.c_str(), module);
-        result = std::string{","} + module + "=" + device_name;
-    } else {
-        printf("Selecting %s as the main image generation device\n", device_name.c_str());
-        result = device_name;
-    }
-    return result;
-}
-
 bool sdtype_load_model(const sd_load_model_inputs inputs) {
     sd_is_quiet = inputs.quiet;
     set_sd_quiet(sd_is_quiet);
@@ -323,7 +294,9 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
     cfg_square_limit = inputs.img_soft_limit;
     printf("\nImageGen Init - Load Model: %s\n",inputs.model_filename);
 
-    std::string backends = get_device_override(inputs.kcpp_main_device);
+    std::string backend = inputs.backend ? inputs.backend : "";
+    std::string params_backend = inputs.params_backend ? inputs.params_backend : "";
+    std::string split_mode = inputs.split_mode ? inputs.split_mode : "";
 
     int lora_apply_mode = LORA_APPLY_AT_RUNTIME;
     bool lora_dynamic = false;
@@ -402,20 +375,33 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
     {
         printf("Conv2D Direct for VAE model is enabled\n");
     }
-    if (inputs.use_mmap && inputs.offload_cpu) {
+    if(backend != "")
+    {
+        printf("Backend assignment: \"%s\"\n", backend.c_str());
+    }
+    if (inputs.use_mmap && params_backend == "CPU") {
         printf("Offloading weights to system RAM with mmap\n");
         if (!lora_dynamic && inputs.lora_len > 0) {
             printf("Note: static LoRAs can reduce mmap memory savings!\n");
         }
-    } else if (inputs.offload_cpu) {
+    } else if (inputs.params_backend == "CPU") {
         printf("Offloading weights to system RAM\n");
     } else if (inputs.use_mmap) {
         printf("Using mmap for I/O\n");
     }
+    if(inputs.auto_fit) {
+        printf("Using auto-fit");
+    }
+    if(params_backend != "" && params_backend != "CPU") {
+        printf("Parameters backend assignment: \"%s\"\n", params_backend.c_str());
+    }
+    if(split_mode != "") {
+        printf("Using split mode: \"%s\"\n", split_mode.c_str());
+    }
     std::string max_vram;
-    if(inputs.max_vram != 0.f) {
-        printf("Using max VRAM = %0.2f GB\n", inputs.max_vram);
-        max_vram = std::to_string(inputs.max_vram);
+    if(inputs.max_vram && *inputs.max_vram) {
+        max_vram = inputs.max_vram;
+        printf("Using max VRAM = %s GB\n", max_vram.c_str());
     }
     if(inputs.quant > 0)
     {
@@ -479,21 +465,15 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
     params.diffusion_flash_attn = sd_params->diffusion_flash_attn;
     params.diffusion_conv_direct = sd_params->diffusion_conv_direct;
     params.vae_conv_direct = sd_params->vae_conv_direct;
-    params.chroma_use_dit_mask = true;
+    params.model_args = "chroma_use_dit_mask=true";
     params.max_vram = max_vram.c_str();
     params.stream_layers = inputs.stream_layers;
     params.eager_load = true; //kcpp should preload everything
     params.enable_mmap = inputs.use_mmap;
-    params.params_backend = inputs.offload_cpu ? "CPU" : "";
-    backends += get_device_override(inputs.kcpp_vae_device, "VAE");
-    backends += get_device_override(inputs.kcpp_clip_device, "CLIP");
-    if (backends.rfind(",", 0) == 0) {
-        backends = "auto" + backends;
-    }
-    params.backend = backends.c_str();
-    if (inputs.debugmode==1) {
-        printf("\nSetting sd backend list to \"%s\", params backend list to \"%s\"", params.backend, params.params_backend);
-    }
+    params.backend = backend.c_str();
+    params.params_backend = params_backend.c_str();
+    params.split_mode = split_mode.c_str();
+    params.auto_fit = inputs.auto_fit;
     params.lora_apply_mode = (lora_apply_mode_t)lora_apply_mode;
 
     // also switches flash attn for the vae and conditioner
@@ -977,6 +957,21 @@ bool supports_reference_images(kcpp_sd::model_info info)
     return supported;
 }
 
+static std::string upscale_image_to_png_base64(upscaler_ctx_t* upscaler_ctx, const sd_image_t& input_image, int upscale_factor = 2, const std::string& meta_image_info = "")
+{
+    std::string gen_data;
+    sd_image_t* upscaled = nullptr;
+    int upscaled_count = 0;
+    if (upscale(upscaler_ctx, input_image, upscale_factor, &upscaled, &upscaled_count)) {
+        gen_data = raw_image_to_png_base64(*upscaled, meta_image_info);
+        free_sd_images(upscaled, upscaled_count);
+    } else {
+        printf("Upscaling failed!\n");
+        gen_data = raw_image_to_png_base64(input_image, meta_image_info);
+    }
+    return gen_data;
+}
+
 sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
 {
     if(sd_ctx == nullptr || sd_params == nullptr)
@@ -984,6 +979,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         return sd_generation.error("Warning: KCPP image generation not initialized!");
     }
     sd_image_t * results = nullptr;
+    int generated_num_results = 0;
 
     std::string img2img_data = std::string(inputs.init_images);
     std::string img2img_mask = std::string(inputs.mask);
@@ -1014,8 +1010,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         sd_params->sample_method = sd_get_default_sample_method(sd_ctx);
     }
 
-    SetCircularAxesAll(sd_ctx, inputs.circular_x, inputs.circular_y);
-
     sd_params->cache_mode    = inputs.cache_mode ? inputs.cache_mode : "";
     sd_params->cache_options = inputs.cache_options ? inputs.cache_options : "";
 
@@ -1273,6 +1267,8 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         params.vae_tiling_params.temporal_tiling = true;
     }
     parse_cache_options(params.cache, sd_params->cache_mode, sd_params->cache_options);
+    params.circular_x = inputs.circular_x;
+    params.circular_y = inputs.circular_y;
 
     LoraMap lora_map = sd_params->lora_map;
     if (sd_params->lora_dynamic) {
@@ -1309,7 +1305,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
     //the below params are only used in video models. May move into standalone object in future
     int vid_req_frames = inputs.vid_req_frames;
     int video_output_type = inputs.video_output_type;
-    int generated_num_results = 1;
     int vid_fps = inputs.vid_fps;
     remove_limits = inputs.remove_limits;
 
@@ -1399,7 +1394,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
 
         fflush(stdout);
 
-        results = generate_image(sd_ctx, &params);
+        if (!generate_image(sd_ctx, &params, &results, &generated_num_results)) {
+            results = nullptr;
+            generated_num_results = 0;
+        }
 
     } else {
 
@@ -1466,7 +1464,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         if (is_passthrough) {
             printf("No generation triggered, passthrough mode.\n");
         } else {
-            results = generate_image(sd_ctx, &params);
+            if (!generate_image(sd_ctx, &params, &results, &generated_num_results)) {
+                results = nullptr;
+                generated_num_results = 0;
+            }
         }
     }
 
@@ -1509,113 +1510,105 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         jsoninfo["all_prompts"] = nlohmann::json::array();
         jsoninfo["all_negative_prompts"] = nlohmann::json::array();
         jsoninfo["all_seeds"] = nlohmann::json::array();
+        jsoninfo["lora_meta"] = lora_meta;
         jsoninfo["version"] = "KoboldCpp";
     }
-    sd_image_t upscaled_image;
-    upscaled_image.data = nullptr;
+    sd_image_t* upscaled_image = nullptr;
     std::string gen_data;
     std::string gen_data2;
     std::string final_frame_data;
 
     if (is_passthrough)
     {
-        //either return original image or upscale if needed
-        sd_image_t *result_image = &input_image;
         if(inputs.upscale && upscaler_ctx != nullptr)
         {
             printf("Upscaling original image (passthrough)...\n");
-            upscaled_image = upscale(upscaler_ctx, input_image, 2);
-            result_image = &upscaled_image;
+            gen_data = upscale_image_to_png_base64(upscaler_ctx, input_image, 2);
         }
-        gen_data = raw_image_to_png_base64(*result_image);
+        else {
+            gen_data = raw_image_to_png_base64(input_image);
+        }
+    }
+    else if (isanim)
+    {
+        //if multiframe, make a video
+        if (generated_num_results > 0 && results && results->data)
+        {
+            if(!sd_is_quiet && sddebugmode==1)
+            {
+                printf("\nSaving video buffer, VIDEO_OUTPUT_TYPE=%d...",video_output_type);
+            }
+            uint8_t * out_data = nullptr;
+            uint8_t * out_data2 = nullptr;
+            size_t out_len = 0;
+            size_t out_len2 = 0;
+            int status = 0;
+            int status2 = 0;
+
+            if(video_output_type==0 || video_output_type==2)
+            {
+                status = create_gif_buf_from_sd_images_msf(results, generated_num_results, vid_fps, &out_data,&out_len);
+            }
+            if(video_output_type==1 || video_output_type==2)
+            {
+                status2 = create_mjpg_avi_membuf_from_sd_images(results, generated_num_results, vid_fps, 40, &out_data2,&out_len2, generated_audio);
+            }
+
+            if(generated_num_results>1)
+            {
+                sd_image_t *final_frame_image = &results[generated_num_results-1];
+                final_frame_data = raw_image_to_png_base64(*final_frame_image);
+            }
+
+            if(!sd_is_quiet && sddebugmode==1)
+            {
+                printf("Video Output Sizes: GIF=%zu AVI=%zu\n",out_len,out_len2);
+                if(status==0 && status2==0)
+                {
+                    printf("Video(s) Saved (Len %zu)!\n",out_len);
+                } else {
+                    printf("Save Failed!\n");
+                }
+            }
+            if(status==0 && out_len>0)
+            {
+                gen_data = kcpp_base64_encode(out_data, out_len);
+                free(out_data);
+            }
+            if (status2 == 0 && out_len2 > 0) {
+                if (gen_data == "") {
+                    gen_data = kcpp_base64_encode(out_data2, out_len2);
+                } else {
+                    gen_data2 = kcpp_base64_encode(out_data2, out_len2);
+                }
+                free(out_data2);
+            }
+        }
+        free_sd_images(results, generated_num_results);
     }
     else
     {
-        for (int i = 0; i < params.batch_count; i++)
+        for (int i = 0; i < generated_num_results; i++)
         {
-            if (results[i].data == NULL) {
+            sd_image_t& result_image = results[i];
+            if (result_image.data == NULL) {
                 continue;
             }
-
-            //if multiframe, make a video
-            if(isanim)
+            std::string meta_image_info = get_image_params(params, lora_meta, i);
+            if(inputs.upscale && upscaler_ctx != nullptr)
             {
-                if(!sd_is_quiet && sddebugmode==1)
-                {
-                    printf("\nSaving video buffer, VIDEO_OUTPUT_TYPE=%d...",video_output_type);
-                }
-                uint8_t * out_data = nullptr;
-                uint8_t * out_data2 = nullptr;
-                size_t out_len = 0;
-                size_t out_len2 = 0;
-                int status = 0;
-                int status2 = 0;
-
-                if(video_output_type==0 || video_output_type==2)
-                {
-                    status = create_gif_buf_from_sd_images_msf(results, generated_num_results, vid_fps, &out_data,&out_len);
-                }
-                if(video_output_type==1 || video_output_type==2)
-                {
-                    status2 = create_mjpg_avi_membuf_from_sd_images(results, generated_num_results, vid_fps, 40, &out_data2,&out_len2, generated_audio);
-                }
-
-                if(generated_num_results>1)
-                {
-                    sd_image_t *final_frame_image = &results[generated_num_results-1];
-                    final_frame_data = raw_image_to_png_base64(*final_frame_image);
-                }
-
-                if(!sd_is_quiet && sddebugmode==1)
-                {
-                    printf("Video Output Sizes: GIF=%zu AVI=%zu\n",out_len,out_len2);
-                    if(status==0 && status2==0)
-                    {
-                        printf("Video(s) Saved (Len %zu)!\n",out_len);
-                    } else {
-                        printf("Save Failed!\n");
-                    }
-                }
-                if(status==0 && out_len>0)
-                {
-                    gen_data = kcpp_base64_encode(out_data, out_len);
-                    free(out_data);
-                }
-                if (status2 == 0 && out_len2 > 0) {
-                    if (gen_data == "") {
-                        gen_data = kcpp_base64_encode(out_data2, out_len2);
-                    } else {
-                        gen_data2 = kcpp_base64_encode(out_data2, out_len2);
-                    }
-                    free(out_data2);
-                }
+                printf("Upscaling output image...\n");
+                gen_data = upscale_image_to_png_base64(upscaler_ctx, result_image, 2, meta_image_info);
+            } else {
+                gen_data = raw_image_to_png_base64(result_image, meta_image_info);
             }
-            else
-            {
-                sd_image_t *result_image = &results[i];
-                if(inputs.upscale && upscaler_ctx != nullptr)
-                {
-                    printf("Upscaling output image...\n");
-                    upscaled_image = upscale(upscaler_ctx, results[i], 2);
-                    result_image = &upscaled_image;
-                }
-                std::string meta_image_info = get_image_params(params, lora_meta, i);
-                gen_data = raw_image_to_png_base64(*result_image, meta_image_info);
-                jsoninfo["infotexts"][i] = meta_image_info;
-                jsoninfo["all_seeds"][i] = params.seed + i;
-                jsoninfo["all_prompts"][i] = params.prompt;
-                jsoninfo["all_negative_prompts"][i] = params.negative_prompt;
-            }
-
-            free(results[i].data);
-            results[i].data = NULL;
+            jsoninfo["infotexts"][i] = meta_image_info;
+            jsoninfo["all_seeds"][i] = params.seed + i;
+            jsoninfo["all_prompts"][i] = params.prompt;
+            jsoninfo["all_negative_prompts"][i] = params.negative_prompt;
         }
-    }
-
-    if(upscaled_image.data)
-    {
-        free(upscaled_image.data);
-        upscaled_image.data = nullptr;
+        free_sd_images(results, generated_num_results);
+        results = nullptr;
     }
 
     if (generated_audio) {
@@ -1627,8 +1620,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
         input_audio.data = nullptr;
     }
 
-    free(results);
-
     total_img_gens += 1;
     if(!sd_is_quiet)
     {
@@ -1662,9 +1653,7 @@ sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs)
     }
     upscale_src_buffer = load_image_from_b64(rawb64,nx,ny);
     sd_image_t source_img;
-    sd_image_t upscaled_image;
     source_img.data = nullptr;
-    upscaled_image.data = nullptr;
     std::string result;
     if(upscale_src_buffer)
     {
@@ -1673,10 +1662,7 @@ sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs)
         source_img.channel = 3;
         source_img.data = upscale_src_buffer;
 
-        upscaled_image = upscale(upscaler_ctx, source_img, inputs.upscaling_resize);
-        result = raw_image_to_png_base64(upscaled_image);
-        free(upscaled_image.data);
-
+        result = upscale_image_to_png_base64(upscaler_ctx, source_img, inputs.upscaling_resize);
     }
 
     if (result == "") {

otherarch/sdcpp/src/conditioning/conditioner.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
+#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
 #include <cmath>
@@ -6,6 +6,7 @@
 #include <optional>
 
 #include "core/tensor_ggml.hpp"
+#include "core/util.h"
 #include "model/te/clip.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
@@ -116,6 +117,10 @@ public:
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
     virtual void set_stream_layers_enabled(bool enabled) {}
+    virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
+    virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
+    virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
+    virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
     virtual void runner_done() {}
@@ -178,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        text_model->set_runtime_backends(backends);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        text_model->set_graph_cut_layer_split_enabled(enabled);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         text_model->set_flash_attention_enabled(enabled);
         if (sd_version_is_sdxl(version)) {
@@ -635,6 +661,48 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
+        if (clip_g) {
+            clip_g->set_runtime_backends(backends);
+        }
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -994,6 +1062,39 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1191,8 +1292,27 @@ struct T5CLIPEmbedder : public Conditioner {
                    bool use_mask                                       = false,
                    int mask_pad                                        = 0,
                    bool is_umt5                                        = false,
-                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                   const char* model_args                              = nullptr)
         : use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) {
+        for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+            if (key == "chroma_use_t5_mask") {
+                bool parsed = false;
+                if (parse_strict_bool(value, parsed)) {
+                    this->use_mask = parsed;
+                } else {
+                    LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "chroma_t5_mask_pad") {
+                int parsed = 0;
+                if (parse_strict_int(value, parsed)) {
+                    this->mask_pad = parsed;
+                } else {
+                    LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+
         bool use_t5 = false;
         for (auto pair : tensor_storage_map) {
             if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
@@ -1226,6 +1346,30 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (t5) {
             t5->set_flash_attention_enabled(enabled);
@@ -1378,6 +1522,125 @@ struct T5CLIPEmbedder : public Conditioner {
     }
 };
 
+struct MiniT2IConditioner : public Conditioner {
+    T5UniGramTokenizer tokenizer;
+    std::shared_ptr<T5Runner> t5;
+    size_t prompt_length = 256;
+
+    MiniT2IConditioner(ggml_backend_t backend,
+                       const String2TensorStorage& tensor_storage_map      = {},
+                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
+        bool use_t5 = false;
+        for (const auto& pair : tensor_storage_map) {
+            if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
+                use_t5 = true;
+                break;
+            }
+        }
+        if (!use_t5) {
+            LOG_WARN("IMPORTANT NOTICE: No MiniT2I T5 text encoder provided, cannot process prompts!");
+            return;
+        }
+        t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager);
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        if (t5) {
+            t5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
+    void set_stream_layers_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_stream_layers_enabled(enabled);
+        }
+    }
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_flash_attention_enabled(enabled);
+        }
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        if (t5) {
+            t5->set_weight_adapter(adapter);
+        }
+    }
+
+    void runner_done() override {
+        if (t5) {
+            t5->runner_done();
+        }
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        SDCondition result;
+        if (!t5) {
+            result.c_crossattn = sd::Tensor<float>::zeros({1024, static_cast<int64_t>(prompt_length)});
+            result.c_vector    = sd::Tensor<float>::zeros({static_cast<int64_t>(prompt_length)});
+            return result;
+        }
+
+        std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
+        if (tokens.size() > prompt_length) {
+            tokens.resize(prompt_length);
+        }
+        std::vector<float> mask(tokens.size(), 1.0f);
+        while (tokens.size() < prompt_length) {
+            tokens.push_back(tokenizer.PAD_TOKEN_ID);
+            mask.push_back(0.0f);
+        }
+
+        sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
+        std::vector<float> t5_mask(mask.size(), 0.0f);
+        for (size_t i = 0; i < mask.size(); ++i) {
+            t5_mask[i] = mask[i] > 0.0f ? 0.0f : -HUGE_VALF;
+        }
+        sd::Tensor<float> hidden_states = t5->compute(n_threads,
+                                                      input_ids,
+                                                      sd::Tensor<float>::from_vector(t5_mask),
+                                                      false,
+                                                      true,
+                                                      true);
+        GGML_ASSERT(!hidden_states.empty());
+        result.c_crossattn = std::move(hidden_states);
+        result.c_vector    = sd::Tensor<float>::from_vector(mask);
+        return result;
+    }
+};
+
 struct AnimaConditioner : public Conditioner {
     std::shared_ptr<BPETokenizer> qwen_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
@@ -1407,6 +1670,22 @@ struct AnimaConditioner : public Conditioner {
         llm->set_stream_layers_enabled(enabled);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1518,7 +1797,7 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) {
+        } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -1552,6 +1831,26 @@ struct LLMEmbedder : public Conditioner {
         llm->set_stream_layers_enabled(enabled);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1997,6 +2296,18 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+        } else if (sd_version_is_sefi_image(version)) {
+            prompt_template_encode_start_idx = 0;
+            min_length                       = 1024;
+            out_layers                       = {9, 18, 27};
+
+            prompt = "<|im_start|>user\n";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (version == VERSION_OVIS_IMAGE) {
             prompt_template_encode_start_idx = 28;
             min_length                       = prompt_template_encode_start_idx + 256;
@@ -2209,6 +2520,22 @@ struct LTXAVEmbedder : public Conditioner {
         projector->set_max_graph_vram_bytes(max_vram_bytes);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         llm->set_weight_adapter(adapter);
         projector->set_weight_adapter(adapter);

otherarch/sdcpp/src/convert.cpp

@@ -1,14 +1,33 @@
+#include <algorithm>
+#include <condition_variable>
+#include <cstdint>
 #include <cstring>
+#include <exception>
+#include <fstream>
+#include <memory>
 #include <mutex>
 #include <regex>
+#include <string>
+#include <thread>
 #include <vector>
 
+#include "core/util.h"
 #include "model_io/gguf_io.h"
 #include "model_io/safetensors_io.h"
+#include "model_io/streaming_writer.h"
 #include "model_loader.h"
-#include "util.h"
 
-#include "ggml_extend_backend.h"
+struct TensorExportInfo {
+    TensorStorage storage;
+    ggml_type type;
+};
+
+struct TensorExportJob {
+    TensorExportInfo info;
+    std::vector<uint8_t> data;
+    std::string error;
+    bool success = false;
+};
 
 static ggml_type get_export_tensor_type(ModelLoader& model_loader,
                                         const TensorStorage& tensor_storage,
@@ -33,106 +52,355 @@ static ggml_type get_export_tensor_type(ModelLoader& model_loader,
     return tensor_type;
 }
 
-static bool load_tensors_for_export(ModelLoader& model_loader,
-                                    ggml_context* ggml_ctx,
-                                    ggml_type type,
-                                    const TensorTypeRules& tensor_type_rules,
-                                    std::vector<TensorWriteInfo>& tensors) {
-    std::mutex tensor_mutex;
-    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-        const std::string& name = tensor_storage.name;
-        ggml_type tensor_type   = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+static bool collect_tensors_for_export(ModelLoader& model_loader,
+                                       ggml_type type,
+                                       const TensorTypeRules& tensor_type_rules,
+                                       std::vector<TensorExportInfo>& tensors) {
+    tensors.clear();
+    tensors.reserve(model_loader.get_tensor_storage_map().size());
+    for (const auto& kv : model_loader.get_tensor_storage_map()) {
+        const TensorStorage& tensor_storage = kv.second;
+        TensorExportInfo info;
+        info.storage = tensor_storage;
+        info.type    = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+        tensors.push_back(std::move(info));
+    }
+    LOG_INFO("collected %zu tensors for export", tensors.size());
+    return true;
+}
 
-        std::lock_guard<std::mutex> lock(tensor_mutex);
-        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
-        if (tensor == nullptr) {
-            LOG_ERROR("ggml_new_tensor failed");
+static size_t export_tensor_nbytes(const TensorExportInfo& info) {
+    TensorStorage output_storage = info.storage;
+    output_storage.type          = info.type;
+    return static_cast<size_t>(output_storage.nbytes());
+}
+
+static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info) {
+    TensorWritePlan plan;
+    plan.name   = info.storage.name;
+    plan.type   = info.type;
+    plan.n_dims = info.storage.n_dims;
+    for (int i = 0; i < SD_MAX_DIMS; i++) {
+        plan.ne[i] = info.storage.ne[i];
+    }
+    return plan;
+}
+
+static std::vector<TensorWritePlan> tensor_write_plans_from_export_infos(const std::vector<TensorExportInfo>& tensors) {
+    std::vector<TensorWritePlan> plans;
+    plans.reserve(tensors.size());
+    for (const TensorExportInfo& info : tensors) {
+        plans.push_back(tensor_write_plan_from_export_info(info));
+    }
+    return plans;
+}
+
+static bool preallocate_output_file(const std::string& output_path, uint64_t file_size, std::string* error) {
+    if (file_size == 0) {
+        return true;
+    }
+
+    std::fstream file(output_path, std::ios::binary | std::ios::in | std::ios::out);
+    if (!file.is_open()) {
+        if (error != nullptr) {
+            *error = "failed to open output file '" + output_path + "' for preallocation";
+        }
+        return false;
+    }
+
+    // This portable fallback sets the final file size. A platform-specific
+    // posix_fallocate/ftruncate path can replace it later.
+    file.seekp(static_cast<std::streamoff>(file_size - 1), std::ios::beg);
+    file.put('\0');
+    file.flush();
+    if (!file) {
+        if (error != nullptr) {
+            *error = "failed to preallocate output file '" + output_path + "'";
+        }
+        return false;
+    }
+    return true;
+}
+
+static bool load_tensor_for_export(ModelLoader& model_loader, TensorExportJob& job) {
+    size_t mem_size = 1 * 1024 * 1024;
+    mem_size += ggml_tensor_overhead();
+    TensorStorage output_storage = job.info.storage;
+    output_storage.type          = job.info.type;
+    mem_size += static_cast<size_t>(output_storage.nbytes());
+
+    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
+    if (ggml_ctx == nullptr) {
+        job.error = "ggml_init failed for tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+
+    ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, job.info.type, job.info.storage.n_dims, job.info.storage.ne);
+    if (tensor == nullptr) {
+        ggml_free(ggml_ctx);
+        job.error = "ggml_new_tensor failed for tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+    ggml_set_name(tensor, job.info.storage.name.c_str());
+
+    const size_t tensor_nbytes = ggml_nbytes(tensor);
+    if (tensor_nbytes > 0 && !model_loader.load_tensor(job.info.storage, tensor)) {
+        ggml_free(ggml_ctx);
+        job.error = "failed to load tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+
+    job.data.resize(tensor_nbytes);
+    if (tensor_nbytes > 0) {
+        memcpy(job.data.data(), tensor->data, tensor_nbytes);
+    }
+    ggml_free(ggml_ctx);
+    return true;
+}
+
+static bool stream_tensor_data(ModelLoader& model_loader,
+                               const std::string& output_path,
+                               const std::vector<TensorExportInfo>& tensors,
+                               const StreamingModelWriter& writer,
+                               int n_threads,
+                               std::string* error) {
+    n_threads = n_threads > 0 ? n_threads : sd_get_num_physical_cores();
+    n_threads = std::max(1, n_threads);
+    LOG_INFO("streaming convert with %d threads", n_threads);
+
+    int64_t start_time       = ggml_time_ms();
+    uint64_t bytes_written   = 0;
+    size_t tensors_written   = 0;
+    size_t next_tensor_index = 0;
+    bool failed              = false;
+    std::string failure;
+
+    const size_t memory_budget = 1024ull * 1024ull * 1024ull;
+    size_t reserved_bytes      = 0;
+
+    std::mutex work_mutex;
+    std::mutex progress_mutex;
+    std::condition_variable memory_cv;
+    std::vector<std::thread> workers;
+    workers.reserve(n_threads);
+
+    auto reserve_memory = [&](size_t bytes) -> bool {
+        std::unique_lock<std::mutex> lock(work_mutex);
+        memory_cv.wait(lock, [&]() {
+            return failed || reserved_bytes == 0 || reserved_bytes + bytes <= memory_budget;
+        });
+        if (failed) {
             return false;
         }
-        ggml_set_name(tensor, name.c_str());
-
-        if (!tensor->data) {
-            GGML_ASSERT(ggml_nelements(tensor) == 0);
-            // Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
-            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
-            tensor->data = ggml_get_mem_buffer(ggml_ctx);
-        }
-
-        TensorWriteInfo write_info;
-        write_info.tensor = tensor;
-        write_info.n_dims = tensor_storage.n_dims;
-        for (int i = 0; i < tensor_storage.n_dims; ++i) {
-            write_info.ne[i] = tensor_storage.ne[i];
-        }
-
-        *dst_tensor = tensor;
-        tensors.push_back(std::move(write_info));
-
+        reserved_bytes += bytes;
         return true;
     };
 
-    bool success = model_loader.load_tensors(on_new_tensor_cb);
-    LOG_INFO("load tensors done");
+    auto release_memory = [&](size_t bytes) {
+        {
+            std::lock_guard<std::mutex> lock(work_mutex);
+            reserved_bytes -= std::min(reserved_bytes, bytes);
+        }
+        memory_cv.notify_all();
+    };
+
+    auto fail = [&](const std::string& message) {
+        {
+            std::lock_guard<std::mutex> lock(work_mutex);
+            if (!failed) {
+                failed  = true;
+                failure = message;
+            }
+        }
+        memory_cv.notify_all();
+    };
+
+    for (int worker = 0; worker < n_threads; worker++) {
+        workers.emplace_back([&]() {
+            std::fstream output_file(output_path, std::ios::binary | std::ios::in | std::ios::out);
+            if (!output_file.is_open()) {
+                fail("failed to open output file '" + output_path + "' for tensor writing");
+                return;
+            }
+
+            while (true) {
+                size_t tensor_index = 0;
+                {
+                    std::lock_guard<std::mutex> lock(work_mutex);
+                    if (failed || next_tensor_index >= tensors.size()) {
+                        return;
+                    }
+                    tensor_index = next_tensor_index++;
+                }
+
+                const size_t tensor_bytes = export_tensor_nbytes(tensors[tensor_index]);
+                if (!reserve_memory(tensor_bytes)) {
+                    return;
+                }
+
+                TensorExportJob job;
+                job.info = tensors[tensor_index];
+                try {
+                    job.success = load_tensor_for_export(model_loader, job);
+                } catch (const std::exception& e) {
+                    job.error   = e.what();
+                    job.success = false;
+                }
+
+                if (!job.success) {
+                    release_memory(tensor_bytes);
+                    fail(job.error.empty() ? "streaming conversion failed" : job.error);
+                    return;
+                }
+
+                std::string write_error;
+                if (!writer.write_tensor(output_file,
+                                         tensor_index,
+                                         job.data.empty() ? nullptr : job.data.data(),
+                                         job.data.size(),
+                                         &write_error)) {
+                    release_memory(tensor_bytes);
+                    fail(write_error.empty() ? "streaming conversion write failed" : write_error);
+                    return;
+                }
+
+                {
+                    std::lock_guard<std::mutex> lock(progress_mutex);
+                    bytes_written += job.data.size();
+                    tensors_written++;
+                    float elapsed_seconds = (ggml_time_ms() - start_time) / 1000.0f;
+                    pretty_bytes_progress(static_cast<int>(tensors_written),
+                                          static_cast<int>(tensors.size()),
+                                          bytes_written,
+                                          elapsed_seconds);
+                }
+                release_memory(tensor_bytes);
+            }
+        });
+    }
+
+    for (auto& worker : workers) {
+        worker.join();
+    }
+    printf("\n");
+    if (failed) {
+        if (error != nullptr) {
+            *error = failure;
+        }
+        return false;
+    }
+    LOG_INFO("streaming conversion completed, taking %.2fs", (ggml_time_ms() - start_time) / 1000.f);
+    return true;
+}
+
+static bool write_model_file_streaming(ModelLoader& model_loader,
+                                       const std::string& output_path,
+                                       const std::vector<TensorExportInfo>& tensors,
+                                       StreamingModelWriter& writer,
+                                       int n_threads,
+                                       std::string* error) {
+    std::vector<TensorWritePlan> plans = tensor_write_plans_from_export_infos(tensors);
+    if (!writer.write_metadata(output_path, plans, error)) {
+        return false;
+    }
+    if (!preallocate_output_file(output_path, writer.file_size(), error)) {
+        return false;
+    }
+    model_loader.process_model_files(false, false);
+    return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, error);
+}
+
+static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
+    if (path == nullptr || strlen(path) == 0) {
+        return true;
+    }
+    if (!model_loader.init_from_file(path, prefix)) {
+        LOG_ERROR("init model loader from file failed: '%s'", path);
+        return false;
+    }
+    loaded_any = true;
+    return true;
+}
+
+static bool export_loaded_model(ModelLoader& model_loader,
+                                const char* output_path,
+                                sd_type_t output_type,
+                                const char* tensor_type_rules,
+                                int n_threads) {
+    ggml_type type             = sd_type_to_ggml_type(output_type);
+    bool output_is_safetensors = ends_with(output_path, ".safetensors");
+    TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
+
+    std::vector<TensorExportInfo> tensors;
+    bool success = collect_tensors_for_export(model_loader, type, type_rules, tensors);
+    std::string error;
+    if (success) {
+        std::unique_ptr<StreamingModelWriter> writer;
+        if (output_is_safetensors) {
+            writer = std::make_unique<SafetensorsStreamingWriter>();
+        } else {
+            writer = std::make_unique<GGUFStreamingWriter>();
+        }
+        success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, &error);
+    }
+
+    if (!success && !error.empty()) {
+        LOG_ERROR("%s", error.c_str());
+    }
+
     return success;
 }
 
+bool convert_with_components(const char* model_path,
+                             const char* clip_l_path,
+                             const char* clip_g_path,
+                             const char* t5xxl_path,
+                             const char* diffusion_model_path,
+                             const char* vae_path,
+                             const char* output_path,
+                             sd_type_t output_type,
+                             const char* tensor_type_rules,
+                             bool convert_name,
+                             int n_threads) {
+    ModelLoader model_loader;
+    bool loaded_any = false;
+
+    if (!init_convert_path(model_loader, model_path, "", loaded_any) ||
+        !init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) ||
+        !init_convert_path(model_loader, vae_path, "vae.", loaded_any)) {
+        return false;
+    }
+
+    if (!loaded_any) {
+        LOG_ERROR("no input model path provided for convert");
+        return false;
+    }
+
+    if (convert_name) {
+        model_loader.convert_tensors_name();
+    }
+
+    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads);
+}
+
 bool convert(const char* input_path,
              const char* vae_path,
              const char* output_path,
              sd_type_t output_type,
              const char* tensor_type_rules,
              bool convert_name) {
-    ModelLoader model_loader;
-
-    if (!model_loader.init_from_file(input_path)) {
-        LOG_ERROR("init model loader from file failed: '%s'", input_path);
-        return false;
-    }
-
-    if (vae_path != nullptr && strlen(vae_path) > 0) {
-        if (!model_loader.init_from_file(vae_path, "vae.")) {
-            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
-            return false;
-        }
-    }
-    if (convert_name) {
-        model_loader.convert_tensors_name();
-    }
-
-    ggml_type type             = sd_type_to_ggml_type(output_type);
-    bool output_is_safetensors = ends_with(output_path, ".safetensors");
-    TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
-
-    auto backend    = sd_backend_cpu_init();
-    size_t mem_size = 1 * 1024 * 1024;  // for padding
-    mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
-    mem_size += model_loader.get_params_mem_size(backend, type);
-    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
-    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
-
-    if (ggml_ctx == nullptr) {
-        LOG_ERROR("ggml_init failed for converter");
-        ggml_backend_free(backend);
-        return false;
-    }
-
-    std::vector<TensorWriteInfo> tensors;
-    bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
-    ggml_backend_free(backend);
-
-    std::string error;
-    if (success) {
-        if (output_is_safetensors) {
-            success = write_safetensors_file(output_path, tensors, &error);
-        } else {
-            success = write_gguf_file(output_path, tensors, &error);
-        }
-    }
-
-    if (!success && !error.empty()) {
-        LOG_ERROR("%s", error.c_str());
-    }
-
-    ggml_free(ggml_ctx);
-    return success;
+    return convert_with_components(input_path,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   vae_path,
+                                   output_path,
+                                   output_type,
+                                   tensor_type_rules,
+                                   convert_name,
+                                   0);
 }

otherarch/sdcpp/src/core/backend_fit.cpp

@@ -0,0 +1,390 @@
+#include "backend_fit.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "ggml-backend.h"
+
+namespace sd::backend_fit {
+    namespace {
+
+        constexpr int64_t MiB = 1024ll * 1024;
+
+        enum class ComponentKind {
+            DIT         = 0,
+            VAE         = 1,
+            CONDITIONER = 2,
+        };
+
+        struct Component {
+            ComponentKind kind;
+            const char* name;
+            int64_t params_bytes  = 0;
+            int64_t reserve_bytes = 0;
+            bool splittable       = false;
+        };
+
+        struct Device {
+            ggml_backend_dev_t dev = nullptr;
+            std::string name;
+            std::string description;
+            int64_t free_bytes   = 0;
+            int64_t total_bytes  = 0;
+            int64_t budget_bytes = 0;
+        };
+
+        struct Decision {
+            ComponentKind kind;
+            bool on_cpu = false;
+            std::vector<size_t> device_idxs;
+        };
+
+        struct Plan {
+            bool valid      = false;
+            bool time_share = false;
+            std::vector<Decision> decisions;
+        };
+
+        bool classify_tensor(const std::string& name, ComponentKind& out) {
+            auto contains = [&](const char* s) { return name.find(s) != std::string::npos; };
+
+            if (contains("model.diffusion_model.") || contains("unet.")) {
+                out = ComponentKind::DIT;
+                return true;
+            }
+            if (contains("first_stage_model.") ||
+                name.rfind("vae.", 0) == 0 ||
+                name.rfind("tae.", 0) == 0) {
+                out = ComponentKind::VAE;
+                return true;
+            }
+            if (contains("text_encoders") ||
+                contains("cond_stage_model") ||
+                contains("te.text_model.") ||
+                contains("conditioner") ||
+                name.rfind("text_encoder.", 0) == 0 ||
+                name.rfind("text_embedding_projection.", 0) == 0 ||
+                contains(".aggregate_embed.")) {
+                out = ComponentKind::CONDITIONER;
+                return true;
+            }
+            return false;
+        }
+
+        std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
+            const auto& storage = loader.get_tensor_storage_map();
+
+            int64_t bytes[3] = {0, 0, 0};
+            for (const auto& [name, ts_const] : storage) {
+                TensorStorage ts = ts_const;
+                if (is_unused_tensor(ts.name)) {
+                    continue;
+                }
+                ComponentKind kind;
+                if (!classify_tensor(ts.name, kind)) {
+                    continue;
+                }
+                if (override_wtype != GGML_TYPE_COUNT &&
+                    loader.tensor_should_be_converted(ts, override_wtype)) {
+                    ts.type = override_wtype;
+                } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) {
+                    ts.type = ts.expected_type;
+                }
+                bytes[int(kind)] += (int64_t)ts.nbytes() + 64;
+            }
+
+            std::vector<Component> out;
+            out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true});
+            out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false});
+            out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true});
+            return out;
+        }
+
+        std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
+            std::vector<Device> out;
+            for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
+                ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+                if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
+                    continue;
+                }
+                Device d;
+                d.dev             = dev;
+                d.name            = ggml_backend_dev_name(dev);
+                d.description     = ggml_backend_dev_description(dev);
+                size_t free_bytes = 0, total_bytes = 0;
+                ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+                d.free_bytes  = (int64_t)free_bytes;
+                d.total_bytes = (int64_t)total_bytes;
+
+                std::string budget_key = d.name;
+                std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(),
+                               [](unsigned char c) { return (char)std::tolower(c); });
+                float gib = budgets.default_gib;
+                auto it   = budgets.backend_gib.find(budget_key);
+                if (it != budgets.backend_gib.end()) {
+                    gib = it->second;
+                }
+                if (gib > 0.f) {
+                    d.budget_bytes = std::min<int64_t>((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes);
+                } else if (gib < 0.f) {
+                    d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0);
+                } else {
+                    d.budget_bytes = d.free_bytes - 512 * MiB;
+                }
+                d.budget_bytes = std::max<int64_t>(d.budget_bytes, 0);
+                out.push_back(d);
+            }
+            return out;
+        }
+
+        Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
+            Plan plan;
+            if (devices.empty()) {
+                return plan;
+            }
+
+            std::vector<size_t> order(components.size());
+            for (size_t i = 0; i < order.size(); i++) {
+                order[i] = i;
+            }
+            std::sort(order.begin(), order.end(), [&](size_t a, size_t b) {
+                return components[a].params_bytes > components[b].params_bytes;
+            });
+
+            {
+                std::vector<int64_t> params_sum(devices.size(), 0);
+                std::vector<int64_t> max_reserve(devices.size(), 0);
+                std::vector<Decision> decisions(components.size());
+                bool ok = true;
+                for (size_t ci : order) {
+                    const Component& comp = components[ci];
+                    decisions[ci].kind    = comp.kind;
+                    if (comp.params_bytes == 0) {
+                        continue;
+                    }
+                    int best = -1;
+                    for (size_t di = 0; di < devices.size(); di++) {
+                        int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes);
+                        if (need <= devices[di].budget_bytes &&
+                            (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) {
+                            best = (int)di;
+                        }
+                    }
+                    if (best < 0) {
+                        ok = false;
+                        break;
+                    }
+                    params_sum[best] += comp.params_bytes;
+                    max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes);
+                    decisions[ci].device_idxs.push_back((size_t)best);
+                }
+                if (ok) {
+                    plan.valid      = true;
+                    plan.time_share = false;
+                    plan.decisions  = std::move(decisions);
+                    return plan;
+                }
+            }
+
+            plan.decisions.assign(components.size(), {});
+            for (size_t ci : order) {
+                const Component& comp = components[ci];
+                Decision& decision    = plan.decisions[ci];
+                decision.kind         = comp.kind;
+                if (comp.params_bytes == 0) {
+                    continue;
+                }
+                int best = -1;
+                for (size_t di = 0; di < devices.size(); di++) {
+                    if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes &&
+                        (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) {
+                        best = (int)di;
+                    }
+                }
+                if (best >= 0) {
+                    decision.device_idxs.push_back((size_t)best);
+                    continue;
+                }
+                if (comp.splittable && devices.size() > 1) {
+                    int64_t capacity = 0;
+                    for (const Device& d : devices) {
+                        capacity += std::max<int64_t>(d.budget_bytes - comp.reserve_bytes, 0);
+                    }
+                    if (comp.params_bytes <= capacity) {
+                        std::vector<size_t> idxs(devices.size());
+                        for (size_t i = 0; i < idxs.size(); i++) {
+                            idxs[i] = i;
+                        }
+                        std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) {
+                            return devices[a].budget_bytes > devices[b].budget_bytes;
+                        });
+                        decision.device_idxs = std::move(idxs);
+                        continue;
+                    }
+                }
+                decision.on_cpu = true;
+            }
+            plan.valid      = true;
+            plan.time_share = true;
+            return plan;
+        }
+
+        void print_plan(const Plan& plan,
+                        const std::vector<Component>& components,
+                        const std::vector<Device>& devices) {
+            LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : "");
+            LOG_INFO("  devices:");
+            for (const Device& d : devices) {
+                LOG_INFO("    %-12s %-32s free %6lld MiB, budget %6lld MiB",
+                         d.name.c_str(), d.description.c_str(),
+                         (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB));
+            }
+            LOG_INFO("  components:");
+            for (size_t ci = 0; ci < components.size(); ci++) {
+                const Component& comp    = components[ci];
+                const Decision& decision = plan.decisions[ci];
+                std::string target;
+                if (comp.params_bytes == 0) {
+                    target = "(not present)";
+                } else if (decision.on_cpu) {
+                    target = "CPU";
+                } else {
+                    for (size_t k = 0; k < decision.device_idxs.size(); k++) {
+                        if (k > 0) {
+                            target += " & ";
+                        }
+                        target += devices[decision.device_idxs[k]].name;
+                    }
+                    if (decision.device_idxs.size() > 1) {
+                        target += " (split)";
+                    }
+                }
+                LOG_INFO("    %-12s params %6lld MiB, compute reserve %5lld MiB -> %s",
+                         comp.name,
+                         (long long)(comp.params_bytes / MiB),
+                         (long long)(comp.reserve_bytes / MiB),
+                         target.c_str());
+            }
+        }
+
+        void append_assignment(std::string& spec, const char* key, const std::string& value) {
+            if (!spec.empty()) {
+                spec += ",";
+            }
+            spec += key;
+            spec += "=";
+            spec += value;
+        }
+
+        void append_component_decision(const std::vector<Component>& components,
+                                       const std::vector<Device>& devices,
+                                       const Plan& plan,
+                                       ComponentKind kind,
+                                       const char* module_key,
+                                       std::string& runtime_spec,
+                                       std::string& params_spec) {
+            for (size_t ci = 0; ci < components.size(); ci++) {
+                if (components[ci].kind != kind || components[ci].params_bytes == 0) {
+                    continue;
+                }
+                const Decision& decision = plan.decisions[ci];
+                if (decision.on_cpu) {
+                    append_assignment(runtime_spec, module_key, "cpu");
+                    return;
+                }
+                if (decision.device_idxs.empty()) {
+                    return;
+                }
+                std::string device_list;
+                for (size_t k = 0; k < decision.device_idxs.size(); k++) {
+                    if (k > 0) {
+                        device_list += "&";
+                    }
+                    device_list += devices[decision.device_idxs[k]].name;
+                }
+                append_assignment(runtime_spec, module_key, device_list);
+                if (plan.time_share) {
+                    append_assignment(params_spec, module_key, "disk");
+                }
+                return;
+            }
+        }
+
+    }  // namespace
+
+    bool derive_backend_specs(ModelLoader& loader,
+                              ggml_type override_wtype,
+                              sd::ggml_graph_cut::MaxVramAssignment& budgets,
+                              std::string& runtime_spec,
+                              std::string& params_spec) {
+        if (!runtime_spec.empty() || !params_spec.empty()) {
+            LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend");
+        }
+
+        {
+            std::string error;
+            if (!budgets.canonicalize_backend_keys(&error)) {
+                LOG_ERROR("%s", error.c_str());
+                return false;
+            }
+        }
+
+        auto components = estimate_components(loader, override_wtype);
+        auto devices    = enumerate_gpu_devices(budgets);
+        auto plan       = compute_plan(components, devices);
+        if (!plan.valid) {
+            LOG_WARN("auto-fit: no usable GPU devices; using the default backend");
+            runtime_spec.clear();
+            params_spec.clear();
+            return true;
+        }
+
+        print_plan(plan, components, devices);
+
+        std::string derived_runtime_spec;
+        std::string derived_params_spec;
+        append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec);
+        append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec);
+        append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec);
+
+        runtime_spec = std::move(derived_runtime_spec);
+        params_spec  = std::move(derived_params_spec);
+
+        LOG_INFO("auto-fit: --backend \"%s\"%s%s%s",
+                 runtime_spec.empty() ? "(default)" : runtime_spec.c_str(),
+                 params_spec.empty() ? "" : " --params-backend \"",
+                 params_spec.c_str(),
+                 params_spec.empty() ? "" : "\"");
+        return true;
+    }
+
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
+        if (prefer_temporal_tiling) {
+            if (tiling_params.temporal_tiling) {
+                return false;
+            }
+            tiling_params.temporal_tiling = true;
+        } else {
+            if (tiling_params.enabled) {
+                return false;
+            }
+            tiling_params.enabled = true;
+            if (tiling_params.tile_size_x <= 0) {
+                tiling_params.tile_size_x = 256;
+            }
+            if (tiling_params.tile_size_y <= 0) {
+                tiling_params.tile_size_y = 256;
+            }
+        }
+
+        LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
+                 tiling_params.temporal_tiling ? "temporal" : "spatial");
+        return true;
+    }
+
+}  // namespace sd::backend_fit

otherarch/sdcpp/src/core/backend_fit.h

@@ -0,0 +1,23 @@
+#ifndef __SD_BACKEND_FIT_H__
+#define __SD_BACKEND_FIT_H__
+
+#include <string>
+
+#include "core/ggml_graph_cut.h"
+#include "model_loader.h"
+#include "stable-diffusion.h"
+
+namespace sd::backend_fit {
+
+    bool derive_backend_specs(ModelLoader& loader,
+                              ggml_type override_wtype,
+                              sd::ggml_graph_cut::MaxVramAssignment& budgets,
+                              std::string& runtime_spec,
+                              std::string& params_spec);
+
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
+                                         bool prefer_temporal_tiling);
+
+}  // namespace sd::backend_fit
+
+#endif  // __SD_BACKEND_FIT_H__

otherarch/sdcpp/src/core/ggml_extend.hpp

@@ -21,10 +21,12 @@
 #include <sstream>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/layer_split_partition.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
@@ -391,7 +393,7 @@ __STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t*
     int64_t width    = input->ne[0];
     int64_t height   = input->ne[1];
     int64_t channels = input->ne[2];
-    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
     if (image_data == nullptr) {
         image_data = (uint8_t*)malloc(width * height * channels);
     }
@@ -1038,6 +1040,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
 }
 
 __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 int lp0,
                                                 int rp0,
@@ -1063,7 +1066,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
     }
 
     if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
-        x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
     }
     return x;
 }
@@ -1076,7 +1089,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
                                             int p3          = 0,
                                             bool circular_x = false,
                                             bool circular_y = false) {
-    return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
+    return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
 }
 
 // w: [OC，IC, KH, KW]
@@ -1105,7 +1118,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
     }
 
     if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
-        x  = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
         p0 = 0;
         p1 = 0;
     }
@@ -1130,6 +1143,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
 // b: [OC,]
 // result: [N*OC, OD, OH, OW]
 __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 ggml_tensor* w,
                                                 ggml_tensor* b,
@@ -1159,7 +1173,21 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
         x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
     } else {
-        x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
     }
 
     if (b != nullptr) {
@@ -1362,6 +1390,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         }
 
         auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
         ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
         if (kv_scale != 1.0f) {
             out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
@@ -1379,9 +1410,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         if (can_use_flash_attn) {
             kqv = build_kqv(q, k, v, mask);
-            if (!ggml_backend_supports_op(backend, kqv)) {
-                kqv = nullptr;
-            } else {
+            if (kqv != nullptr) {
                 kqv = ggml_view_4d(ctx,
                                    kqv,
                                    d_head,
@@ -1718,6 +1747,13 @@ protected:
     size_t max_graph_vram_bytes           = 0;
     bool stream_layers_enabled            = false;
     size_t observed_max_effective_budget_ = 0;
+    bool graph_cut_layer_split_enabled    = false;
+    std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
+
+    std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
+    ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
+    ggml_backend_t cpu_fallback_backend    = nullptr;    // owned, sched requires a trailing CPU backend
+    bool multi_device_eval_callback_warned = false;
 
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<RunnerWeightManager> weight_manager;
@@ -1744,6 +1780,9 @@ protected:
 
     sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
     std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
 
     template <typename T>
     static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -1842,6 +1881,20 @@ protected:
         params_tensor_set_dirty_ = false;
     }
 
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
+            return tensor;
+        }
+        if (tensor->view_src != nullptr &&
+            params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
+            return tensor->view_src;
+        }
+        return nullptr;
+    }
+
     std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
         std::vector<ggml_tensor*> used_params;
         rebuild_params_tensor_set();
@@ -1854,12 +1907,8 @@ protected:
         seen_params.reserve(static_cast<size_t>(n_leafs));
         for (int i = 0; i < n_leafs; ++i) {
             ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
-            ggml_tensor* param_leaf = leaf;
-            if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
-                param_leaf = param_leaf->view_src;
-            }
+            ggml_tensor* param_leaf = canonical_param_tensor(leaf);
             if (param_leaf != nullptr &&
-                params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
                 seen_params.insert(param_leaf).second) {
                 used_params.push_back(param_leaf);
             }
@@ -1986,7 +2035,127 @@ protected:
         return true;
     }
 
+    // Pass explicit buffer types: synthesized defaults can make CUDA devices
+    // report supporting each other's buffers and skip a required copy.
+    bool ensure_sched(ggml_cgraph* gf) {
+        if (sched != nullptr) {
+            return true;
+        }
+        std::vector<ggml_backend_t> backends;
+        backends.reserve(extra_runtime_backends.size() + 2);
+        backends.push_back(runtime_backend);
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            backends.push_back(backend);
+        }
+        if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) {
+            cpu_fallback_backend = sd_backend_cpu_init();
+        }
+        if (cpu_fallback_backend != nullptr) {
+            backends.push_back(cpu_fallback_backend);
+        }
+
+        std::vector<ggml_backend_buffer_type_t> bufts;
+        bufts.reserve(backends.size());
+        ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend);
+        for (ggml_backend_t backend : backends) {
+            ggml_backend_buffer_type_t buft = nullptr;
+            if (backend == cpu_fallback_backend && main_dev != nullptr) {
+                buft = ggml_backend_dev_host_buffer_type(main_dev);
+            }
+            if (buft == nullptr) {
+                buft = ggml_backend_get_default_buffer_type(backend);
+            }
+            bufts.push_back(buft);
+        }
+
+        size_t graph_size = MAX_GRAPH_SIZE;
+        if (gf != nullptr) {
+            graph_size = std::max<size_t>(graph_size, (size_t)ggml_graph_n_nodes(gf));
+        }
+        sched = ggml_backend_sched_new(backends.data(),
+                                       bufts.data(),
+                                       (int)backends.size(),
+                                       graph_size,
+                                       /*parallel=*/false,
+                                       /*op_offload=*/false);
+        if (sched == nullptr) {
+            LOG_ERROR("%s: failed to create backend sched", get_desc().c_str());
+            return false;
+        }
+        return true;
+    }
+
+    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const {
+        if (tensor == nullptr || tensor->buffer == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+            ggml_backend_buffer_is_host(tensor->buffer)) {
+            return nullptr;
+        }
+        ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+        if (dev == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_get_device(runtime_backend) == dev) {
+            return runtime_backend;
+        }
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            if (ggml_backend_get_device(backend) == dev) {
+                return backend;
+            }
+        }
+        return nullptr;
+    }
+
+    // Weightless ops have no scheduler anchor, so pin them to the most recent
+    // weight device. Views must stay unpinned or cross-device copies can be
+    // skipped for their consumers.
+    void pin_multi_device_nodes(ggml_cgraph* gf) {
+        if (sched == nullptr || gf == nullptr) {
+            return;
+        }
+        ggml_backend_t current = runtime_backend;
+        const int n_nodes      = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; i++) {
+            ggml_tensor* node    = ggml_graph_node(gf, i);
+            auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
+            if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+                current = node_assignment->second;
+            }
+            for (int s = 0; s < GGML_MAX_SRC; s++) {
+                ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+                if (weight_backend != nullptr) {
+                    if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
+                        current = weight_backend;
+                    }
+                }
+            }
+            if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+                node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+                continue;
+            }
+            if (ggml_backend_supports_op(current, node)) {
+                ggml_backend_sched_set_tensor_backend(sched, node, current);
+            }
+        }
+    }
+
+    bool is_multi_device() const {
+        return !extra_runtime_backends.empty();
+    }
+
     bool alloc_compute_buffer(ggml_cgraph* gf) {
+        if (is_multi_device()) {
+            // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
+            // the graph here: reserve runs split_graph, which rewires the
+            // graph's src pointers to sched-internal copy tensors, and the
+            // later ggml_backend_sched_alloc_graph would split the already
+            // rewired graph, silently corrupting every cross-backend input. A
+            // graph must be split at most once; the alloc in execute_graph
+            // performs the real allocation.
+            return ensure_sched(gf);
+        }
         if (compute_allocr != nullptr) {
             return true;
         }
@@ -2202,12 +2371,14 @@ protected:
                plan.valid &&
                max_graph_vram_bytes > 0 &&
                plan.segments.size() > 1 &&
-               !sd_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend) &&
+               !is_multi_device();
     }
 
     bool can_attempt_graph_cut_segmented_compute() const {
         return max_graph_vram_bytes > 0 &&
-               !sd_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend) &&
+               !is_multi_device();
     }
 
     bool resolve_graph_cut_plan(ggml_cgraph* gf,
@@ -2287,6 +2458,123 @@ protected:
         return true;
     }
 
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out) {
+        GGML_ASSERT(plan_out != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                     gf,
+                                                     &graph_cut_plan_cache_,
+                                                     0,
+                                                     params_tensor_set_,
+                                                     get_desc().c_str());
+        return true;
+    }
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+        graph_cut_layer_split_node_assignments_.clear();
+        if (!graph_cut_layer_split_enabled) {
+            return true;
+        }
+        if (!is_multi_device()) {
+            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+            return false;
+        }
+
+        GraphCutPlan plan;
+        if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+            return false;
+        }
+        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+            auto manager = weight_manager.lock();
+            if (manager == nullptr) {
+                LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+                return false;
+            }
+            std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+            if (!graph_params.empty() &&
+                !manager->assign_compute_backend(graph_params, runtime_backend)) {
+                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+                return false;
+            }
+            for (ggml_tensor* param : graph_params) {
+                if (param != nullptr) {
+                    graph_cut_layer_split_assignments_[param] = runtime_backend;
+                }
+            }
+            const int n_nodes = ggml_graph_n_nodes(gf);
+            for (int i = 0; i < n_nodes; i++) {
+                ggml_tensor* node = ggml_graph_node(gf, i);
+                if (node != nullptr) {
+                    graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+                }
+            }
+            if (!graph_cut_layer_split_primary_notice_logged_) {
+                LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                         get_desc().c_str(),
+                         sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                         graph_params.size());
+                graph_cut_layer_split_primary_notice_logged_ = true;
+            } else {
+                LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                          graph_params.size());
+            }
+            return true;
+        }
+
+        std::vector<ggml_backend_t> split_backends;
+        split_backends.reserve(extra_runtime_backends.size() + 1);
+        split_backends.push_back(runtime_backend);
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            if (backend != nullptr) {
+                split_backends.push_back(backend);
+            }
+        }
+
+        auto manager = weight_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+
+        sd::GraphCutLayerSplitAssignment assignment;
+        auto canonicalize_param = [this](ggml_tensor* tensor) {
+            return canonical_param_tensor(tensor);
+        };
+        if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                                 gf,
+                                                 plan,
+                                                 split_backends,
+                                                 graph_cut_layer_split_backend_vram_limits_,
+                                                 max_graph_vram_bytes,
+                                                 graph_cut_layer_split_assignments_,
+                                                 canonicalize_param,
+                                                 &assignment)) {
+            return false;
+        }
+
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+                LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+                return false;
+            }
+        }
+
+        graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+        sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+        return true;
+    }
+
     struct PersistentExternalBinding {
         ggml_backend_buffer_t buffer = nullptr;
         void* data                   = nullptr;
@@ -2463,7 +2751,14 @@ protected:
         };
         ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
 
-        if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
+        if (is_multi_device()) {
+            ggml_backend_sched_reset(sched);
+            pin_multi_device_nodes(gf);  // reset clears the pins; re-apply before alloc
+            if (!ggml_backend_sched_alloc_graph(sched, gf)) {
+                LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str());
+                return std::nullopt;
+            }
+        } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
             LOG_ERROR("%s alloc compute graph failed", get_desc().c_str());
             return std::nullopt;
         }
@@ -2472,8 +2767,27 @@ protected:
         if (sd_backend_is_cpu(runtime_backend)) {
             sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
+        if (cpu_fallback_backend != nullptr) {
+            sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads);
+        }
 
-        ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
+        ggml_status status;
+        if (is_multi_device()) {
+            if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
+                LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
+                         get_desc().c_str());
+                multi_device_eval_callback_warned = true;
+            }
+            status = ggml_backend_sched_graph_compute(sched, gf);
+            if (status == GGML_STATUS_SUCCESS) {
+                ggml_backend_sched_synchronize(sched);
+            }
+        } else {
+            status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
+                                                                 gf,
+                                                                 sd_get_backend_eval_callback(),
+                                                                 sd_get_backend_eval_callback_data());
+        }
         if (status != GGML_STATUS_SUCCESS) {
             LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
             return std::nullopt;
@@ -2650,6 +2964,10 @@ public:
         free_params_ctx();
         free_compute_ctx();
         free_cache_ctx_and_buffer();
+        if (cpu_fallback_backend != nullptr) {
+            ggml_backend_free(cpu_fallback_backend);
+            cpu_fallback_backend = nullptr;
+        }
     }
 
     virtual GGMLRunnerContext get_context() {
@@ -2690,10 +3008,20 @@ public:
             ggml_gallocr_free(compute_allocr);
             compute_allocr = nullptr;
         }
+        if (sched != nullptr) {
+            ggml_backend_sched_free(sched);
+            sched = nullptr;
+        }
     }
 
     // do copy after alloc graph
     void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
+        if (is_multi_device()) {
+            // The sched only assigns a backend (and thus a buffer) to tensors
+            // that participate in the graph; flag standalone data tensors as
+            // inputs so they get one.
+            ggml_set_input(tensor);
+        }
         backend_tensor_data_map[tensor] = data;
     }
 
@@ -2784,6 +3112,11 @@ public:
         GGML_ASSERT(gf != nullptr);
         rebuild_params_tensor_set();
 
+        if (!assign_graph_cut_layer_split_backends(gf)) {
+            free_compute_ctx();
+            return std::nullopt;
+        }
+
         if (can_attempt_graph_cut_segmented_compute()) {
             GraphCutPlan plan;
             if (!resolve_graph_cut_plan(gf, &plan)) {
@@ -2829,8 +3162,50 @@ public:
     }
 
     void set_stream_layers_enabled(bool enabled) {
+        if (enabled && is_multi_device()) {
+            LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
+                     get_desc().c_str());
+            return;
+        }
         stream_layers_enabled = enabled;
     }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) {
+        graph_cut_layer_split_enabled = enabled;
+        if (!enabled) {
+            graph_cut_layer_split_assignments_.clear();
+            graph_cut_layer_split_node_assignments_.clear();
+            graph_cut_layer_split_primary_notice_logged_ = false;
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+        graph_cut_layer_split_backend_vram_limits_ = limits;
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+        extra_runtime_backends.clear();
+        for (ggml_backend_t backend : backends) {
+            if (backend == nullptr || backend == runtime_backend) {
+                continue;
+            }
+            if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
+                extra_runtime_backends.end()) {
+                extra_runtime_backends.push_back(backend);
+            }
+        }
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+        if (is_multi_device() && stream_layers_enabled) {
+            LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
+                     get_desc().c_str());
+            stream_layers_enabled = false;
+        }
+    }
 };
 
 class GGMLBlock {
@@ -3365,7 +3740,7 @@ public:
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
         }
-        return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, 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),

otherarch/sdcpp/src/core/ggml_extend_backend.cpp

@@ -9,6 +9,7 @@
 #include <vector>
 
 #include "core/util.h"
+#include "ggml/src/ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -110,7 +111,67 @@ static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type)
     if (dev == nullptr) {
         return "";
     }
-    return ggml_backend_dev_name(dev);
+    const char* dev_name = ggml_backend_dev_name(dev);
+    if (dev_name != nullptr && dev_name[0] != '\0') {
+        return dev_name;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    const char* reg_name   = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
+    return reg_name != nullptr ? reg_name : "";
+}
+
+static ggml_backend_dev_t resolve_first_device_by_registry_name(const std::string& name) {
+    std::string lower = lower_copy(trim_copy(name));
+    if (lower == "metal") {
+        lower = "mtl";
+    }
+    if (lower.empty()) {
+        return nullptr;
+    }
+
+    const size_t device_count = ggml_backend_dev_count();
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+        if (reg == nullptr) {
+            continue;
+        }
+        const char* reg_name = ggml_backend_reg_name(reg);
+        if (reg_name != nullptr && lower_copy(reg_name) == lower) {
+            return dev;
+        }
+    }
+    return nullptr;
+}
+
+static ggml_backend_dev_t resolve_device_by_name(const std::string& name) {
+    const std::string lower = lower_copy(trim_copy(name));
+    if (lower.empty()) {
+        return nullptr;
+    }
+
+    const size_t device_count = ggml_backend_dev_count();
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        const char* dev_name   = ggml_backend_dev_name(dev);
+        if (dev_name != nullptr && lower_copy(dev_name) == lower) {
+            return dev;
+        }
+    }
+    return nullptr;
+}
+
+static std::string backend_device_name(ggml_backend_dev_t dev) {
+    if (dev == nullptr) {
+        return "";
+    }
+    const char* name = ggml_backend_dev_name(dev);
+    if (name != nullptr && name[0] != '\0') {
+        return name;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    const char* reg_name   = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
+    return reg_name != nullptr ? reg_name : "";
 }
 
 static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
@@ -296,6 +357,10 @@ std::string sd_backend_resolve_name(const std::string& name) {
         return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
     }
 
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(requested)) {
+        return backend_device_name(dev);
+    }
+
     const size_t device_count = ggml_backend_dev_count();
     for (size_t i = 0; i < device_count; ++i) {
         ggml_backend_dev_t dev = ggml_backend_dev_get(i);
@@ -328,7 +393,20 @@ static ggml_backend_t init_named_backend(const std::string& name) {
         return ggml_backend_init_best();
     }
 
+    if (ggml_backend_dev_t dev = resolve_device_by_name(name)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(name)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+
     std::string resolved = sd_backend_resolve_name(name);
+    if (ggml_backend_dev_t dev = resolve_device_by_name(resolved)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
+    if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(resolved)) {
+        return ggml_backend_dev_init(dev, nullptr);
+    }
     if (resolved.empty()) {
         return nullptr;
     }
@@ -364,6 +442,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
     return false;
 }
 
+static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
+    ggml_cgraph cgraph = {
+        /*.size             =*/0,
+        /*.n_nodes          =*/i1 - i0,
+        /*.n_leafs          =*/0,
+        /*.nodes            =*/cgraph0->nodes + i0,
+        /*.grads            =*/nullptr,
+        /*.grad_accs        =*/nullptr,
+        /*.leafs            =*/nullptr,
+        /*.use_counts       =*/cgraph0->use_counts,
+        /*.visited_hash_set =*/cgraph0->visited_hash_set,
+        /*.order            =*/cgraph0->order,
+        /*.uid              =*/0,
+    };
+    return cgraph;
+}
+
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data) {
+    if (callback_eval == nullptr) {
+        return ggml_backend_graph_compute(backend, gf);
+    }
+
+    ggml_status status = GGML_STATUS_SUCCESS;
+    const int n_nodes  = ggml_graph_n_nodes(gf);
+    bool stopped       = false;
+
+    for (int j0 = 0; j0 < n_nodes; ++j0) {
+        ggml_tensor* t = ggml_graph_node(gf, j0);
+        bool need      = callback_eval(t, true, callback_eval_user_data);
+        int j1         = j0;
+
+        while (!need && j1 < n_nodes - 1) {
+            t    = ggml_graph_node(gf, ++j1);
+            need = callback_eval(t, true, callback_eval_user_data);
+        }
+
+        ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
+        status         = ggml_backend_graph_compute_async(backend, &gv);
+        if (status != GGML_STATUS_SUCCESS) {
+            break;
+        }
+
+        ggml_backend_synchronize(backend);
+
+        if (need && !callback_eval(t, false, callback_eval_user_data)) {
+            stopped = true;
+            break;
+        }
+
+        j0 = j1;
+    }
+
+    ggml_backend_synchronize(backend);
+    if (stopped && status == GGML_STATUS_SUCCESS) {
+        status = GGML_STATUS_ABORTED;
+    }
+    return status;
+}
+
 const char* sd_get_system_info() {
     static std::string cache_info = []() -> std::string {
         ggml_backend_load_all_once();
@@ -525,12 +665,52 @@ SDBackendManager::~SDBackendManager() {
 
 void SDBackendManager::reset() {
     backends_.clear();
-    runtime_assignment_ = {};
-    params_assignment_  = {};
+    runtime_assignment_    = {};
+    params_assignment_     = {};
+    split_mode_assignment_ = {};
+}
+
+static std::vector<std::string> split_device_list(const std::string& value) {
+    std::vector<std::string> names;
+    for (const std::string& raw : split_copy(value, '&')) {
+        const std::string name = trim_copy(raw);
+        if (!name.empty()) {
+            names.push_back(name);
+        }
+    }
+    return names;
+}
+
+static std::string primary_device_name(const std::string& value) {
+    std::vector<std::string> names = split_device_list(value);
+    return names.empty() ? std::string() : names.front();
 }
 
 ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) {
-    return init_cached_backend(runtime_assignment_.get(module));
+    return init_cached_backend(primary_device_name(runtime_assignment_.get(module)));
+}
+
+std::vector<ggml_backend_t> SDBackendManager::runtime_backends(SDBackendModule module) {
+    std::vector<ggml_backend_t> backends;
+    for (const std::string& name : split_device_list(runtime_assignment_.get(module))) {
+        ggml_backend_t backend = init_cached_backend(name);
+        if (backend == nullptr) {
+            LOG_ERROR("failed to initialize backend '%s' for module %s",
+                      name.c_str(),
+                      sd_backend_module_name(module));
+            continue;
+        }
+        if (std::find(backends.begin(), backends.end(), backend) == backends.end()) {
+            backends.push_back(backend);
+        }
+    }
+    if (backends.empty()) {
+        ggml_backend_t backend = runtime_backend(module);
+        if (backend != nullptr) {
+            backends.push_back(backend);
+        }
+    }
+    return backends;
 }
 
 ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
@@ -556,6 +736,10 @@ bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const {
     return is_disk_backend_token(params_assignment_.get(module));
 }
 
+bool SDBackendManager::params_backend_follows_runtime(SDBackendModule module) const {
+    return params_assignment_.get(module).empty();
+}
+
 bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
     ggml_backend_t backend = runtime_backend(module);
     if (backend == nullptr) {
@@ -575,6 +759,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
 
 bool SDBackendManager::init(const char* backend_spec,
                             const char* params_backend_spec,
+                            const char* split_mode_spec,
                             std::string* error) {
     reset();
 
@@ -584,12 +769,53 @@ bool SDBackendManager::init(const char* backend_spec,
     if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), &params_assignment_, error)) {
         return false;
     }
+    if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) {
+        return false;
+    }
 
     return validate(error);
 }
 
+SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const {
+    return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW
+                                                                              : SDSplitMode::LAYER;
+}
+
+ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend,
+                                                               const std::vector<float>& tensor_split) {
+    if (backend == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    if (reg == nullptr) {
+        return nullptr;
+    }
+    auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type");
+    if (fn == nullptr) {
+        return nullptr;
+    }
+    int main_device        = -1;
+    const size_t dev_count = ggml_backend_reg_dev_count(reg);
+    for (size_t i = 0; i < dev_count; ++i) {
+        if (ggml_backend_reg_dev_get(reg, i) == dev) {
+            main_device = (int)i;
+            break;
+        }
+    }
+    if (main_device < 0) {
+        return nullptr;
+    }
+    std::vector<float> padded_split(std::max<size_t>(tensor_split.size(), 64), 0.0f);
+    std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin());
+    return fn(main_device, padded_split.data());
+}
+
 bool SDBackendManager::validate(std::string* error) const {
-    auto validate_runtime_name = [&](const std::string& name) -> bool {
+    auto validate_single_runtime_name = [&](const std::string& name) -> bool {
         if (is_default_backend_token(name)) {
             return true;
         }
@@ -599,7 +825,7 @@ bool SDBackendManager::validate(std::string* error) const {
             }
             return false;
         }
-        if (!sd_backend_resolve_name(name).empty()) {
+        if (!sd_backend_resolve_name(name).empty() || resolve_first_device_by_registry_name(name) != nullptr) {
             return true;
         }
         if (error != nullptr) {
@@ -607,15 +833,56 @@ bool SDBackendManager::validate(std::string* error) const {
         }
         return false;
     };
+    auto validate_runtime_name = [&](const std::string& name) -> bool {
+        if (name.find('&') == std::string::npos) {
+            return validate_single_runtime_name(name);
+        }
+        std::vector<std::string> names = split_device_list(name);
+        if (names.empty()) {
+            if (error != nullptr) {
+                *error = "invalid backend device list '" + name + "'";
+            }
+            return false;
+        }
+        for (const std::string& entry : names) {
+            if (is_default_backend_token(entry)) {
+                if (error != nullptr) {
+                    *error = "default backend token is not allowed in a device list '" + name + "'";
+                }
+                return false;
+            }
+            if (!validate_single_runtime_name(entry)) {
+                return false;
+            }
+        }
+        return true;
+    };
     auto validate_params_name = [&](const std::string& name) -> bool {
         if (is_disk_backend_token(name)) {
             return true;
         }
-        return validate_runtime_name(name);
+        if (name.find('&') != std::string::npos) {
+            if (error != nullptr) {
+                *error = "params_backend does not accept device lists ('" + name + "')";
+            }
+            return false;
+        }
+        return validate_single_runtime_name(name);
+    };
+    auto validate_split_mode_name = [&](const std::string& name) -> bool {
+        const std::string lower = lower_copy(trim_copy(name));
+        if (lower.empty() || lower == "layer" || lower == "row") {
+            return true;
+        }
+        if (error != nullptr) {
+            *error = "invalid split mode '" + name + "' (expected layer or row)";
+        }
+        return false;
     };
 
     if (!validate_runtime_name(runtime_assignment_.default_name) ||
-        !validate_params_name(params_assignment_.default_name)) {
+        !validate_params_name(params_assignment_.default_name) ||
+        !validate_split_mode_name(split_mode_assignment_.default_name)) {
         return false;
     }
     for (const auto& kv : runtime_assignment_.module_names) {
@@ -628,6 +895,11 @@ bool SDBackendManager::validate(std::string* error) const {
             return false;
         }
     }
+    for (const auto& kv : split_mode_assignment_.module_names) {
+        if (!validate_split_mode_name(kv.second)) {
+            return false;
+        }
+    }
     return true;
 }
 

otherarch/sdcpp/src/core/ggml_extend_backend.h

@@ -6,9 +6,11 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "stable-diffusion.h"
 
 enum class SDBackendModule {
     DIFFUSION,
@@ -36,10 +38,16 @@ struct SDBackendHandleDeleter {
 
 using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
 
+enum class SDSplitMode {
+    LAYER,
+    ROW,
+};
+
 class SDBackendManager {
 private:
     SDBackendAssignment runtime_assignment_;
     SDBackendAssignment params_assignment_;
+    SDBackendAssignment split_mode_assignment_;
     std::unordered_map<std::string, SDBackendHandle> backends_;
 
 public:
@@ -51,15 +59,23 @@ public:
 
     bool init(const char* backend_spec,
               const char* params_backend_spec,
+              const char* split_mode_spec,
               std::string* error);
     void reset();
 
     ggml_backend_t runtime_backend(SDBackendModule module);
     ggml_backend_t params_backend(SDBackendModule module);
 
+    std::vector<ggml_backend_t> runtime_backends(SDBackendModule module);
+
+    SDSplitMode split_mode(SDBackendModule module) const;
+    ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend,
+                                                 const std::vector<float>& tensor_split);
+
     bool runtime_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_disk(SDBackendModule module) const;
+    bool params_backend_follows_runtime(SDBackendModule module) const;
     bool runtime_backend_supports_host_buffer(SDBackendModule module);
 
 private:
@@ -71,6 +87,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);

otherarch/sdcpp/src/core/layer_split_partition.cpp

@@ -0,0 +1,257 @@
+#include "core/layer_split_partition.h"
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <limits>
+#include <unordered_set>
+#include <utility>
+
+#include "core/util.h"
+
+namespace sd {
+
+    static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
+        return pos == 0 || name[pos - 1] == '.';
+    }
+
+    static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
+        size_t pos = name.find(segment);
+        while (pos != std::string::npos) {
+            if (layer_split_path_segment_starts_at(name, pos)) {
+                return true;
+            }
+            pos = name.find(segment, pos + 1);
+        }
+        return false;
+    }
+
+    int layer_split_tensor_block_index(const std::string& name) {
+        static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
+                                                    "down_blocks.", "up_blocks.", "mid_block."};
+        for (const char* segment : unet_block_segments) {
+            if (layer_split_has_path_segment(name, segment)) {
+                return -1;
+            }
+        }
+
+        static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
+                                               "single_blocks.", "blocks.", "block.", "layers."};
+        for (const char* keyword : block_keywords) {
+            size_t pos = name.find(keyword);
+            while (pos != std::string::npos) {
+                if (!layer_split_path_segment_starts_at(name, pos)) {
+                    pos = name.find(keyword, pos + 1);
+                    continue;
+                }
+                pos += std::strlen(keyword);
+                size_t end = pos;
+                while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
+                    end++;
+                }
+                if (end > pos && (end == name.size() || name[end] == '.')) {
+                    return std::atoi(name.substr(pos, end - pos).c_str());
+                }
+                break;
+            }
+        }
+        return -1;
+    }
+
+    std::string layer_split_backend_device_display_name(ggml_backend_t backend) {
+        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+        const char* name       = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend);
+        return name != nullptr ? name : "unknown";
+    }
+
+    static size_t graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
+                                                           size_t backend_index,
+                                                           size_t primary_backend_vram_limit) {
+        if (backend_index < backend_vram_limits.size()) {
+            return backend_vram_limits[backend_index];
+        }
+        return backend_index == 0 ? primary_backend_vram_limit : 0;
+    }
+
+    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);
+        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]);
+            size_t free_bytes = 0, total_bytes = 0;
+            if (dev != nullptr) {
+                ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            }
+            if (free_bytes > 0) {
+                capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
+            }
+            size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
+                                                                          i,
+                                                                          primary_backend_vram_limit);
+            if (limit_bytes > 0) {
+                capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
+            }
+        }
+        return capacities;
+    }
+
+    bool partition_graph_cut_layer_split(const char* desc,
+                                         ggml_cgraph* gf,
+                                         const sd::ggml_graph_cut::Plan& plan,
+                                         const std::vector<ggml_backend_t>& split_backends,
+                                         const std::vector<size_t>& backend_vram_limits,
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out) {
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(assignment_out != nullptr);
+        GGML_ASSERT(canonical_param_tensor != nullptr);
+        GGML_ASSERT(!split_backends.empty());
+
+        GraphCutLayerSplitAssignment assignment;
+        assignment.segment_count = plan.segments.size();
+        assignment.tensors_by_backend.resize(split_backends.size());
+        assignment.bytes_by_backend.resize(split_backends.size(), 0);
+        assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
+        assignment.last_segment_by_backend.resize(split_backends.size(), 0);
+
+        std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
+        std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
+        std::unordered_set<ggml_tensor*> seen_params;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
+            for (ggml_tensor* raw_param : params) {
+                ggml_tensor* param = canonical_param_tensor(raw_param);
+                if (param == nullptr || !seen_params.insert(param).second) {
+                    continue;
+                }
+                segment_params[seg_idx].push_back(param);
+                segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
+            }
+        }
+
+        int64_t total_param_bytes = 0;
+        for (int64_t bytes : segment_param_bytes) {
+            total_param_bytes += bytes;
+        }
+        if (total_param_bytes <= 0) {
+            LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
+            return false;
+        }
+
+        std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
+                                                                                           backend_vram_limits,
+                                                                                           primary_backend_vram_limit);
+
+        std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
+        size_t current_backend = 0;
+        int64_t current_used   = 0;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            int64_t bytes = segment_param_bytes[seg_idx];
+            while (current_backend + 1 < split_backends.size() &&
+                   bytes > 0 &&
+                   current_used + bytes > backend_capacities[current_backend]) {
+                current_backend++;
+                current_used = 0;
+            }
+            if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
+                LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
+                          desc,
+                          seg_idx,
+                          (current_used + bytes) / (1024.0 * 1024.0),
+                          layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
+                          backend_capacities[current_backend] / (1024.0 * 1024.0));
+                return false;
+            }
+            current_used += bytes;
+            backend_by_segment[seg_idx] = split_backends[current_backend];
+
+            for (ggml_tensor* param : segment_params[seg_idx]) {
+                ggml_backend_t target_backend = split_backends[current_backend];
+                auto assigned_it              = param_assignments.find(param);
+                if (assigned_it == param_assignments.end()) {
+                    param_assignments[param]            = target_backend;
+                    assignment.has_new_param_assignment = true;
+                } else {
+                    target_backend = assigned_it->second;
+                }
+
+                auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
+                if (backend_it == split_backends.end()) {
+                    LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
+                              desc,
+                              ggml_get_name(param));
+                    return false;
+                }
+                size_t backend_idx                               = (size_t)std::distance(split_backends.begin(), backend_it);
+                assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
+                assignment.last_segment_by_backend[backend_idx]  = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
+                assignment.tensors_by_backend[backend_idx].push_back(param);
+                assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
+            }
+        }
+
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        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];
+            for (int node_index : segment.internal_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+            for (int node_index : segment.output_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+        }
+
+        *assignment_out = std::move(assignment);
+        return true;
+    }
+
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment) {
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (i >= assignment.tensors_by_backend.size() ||
+                assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
+                                       ? 0
+                                       : assignment.first_segment_by_backend[i];
+            size_t last_segment  = assignment.last_segment_by_backend[i];
+            if (assignment.has_new_param_assignment) {
+                LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                         desc,
+                         layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                         first_segment,
+                         last_segment,
+                         assignment.tensors_by_backend[i].size(),
+                         assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
+            } else {
+                LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                          desc,
+                          layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                          first_segment,
+                          last_segment,
+                          assignment.tensors_by_backend[i].size(),
+                          assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
+            }
+        }
+    }
+
+}  // namespace sd

otherarch/sdcpp/src/core/layer_split_partition.h

@@ -0,0 +1,44 @@
+#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
+#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
+
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#include "core/ggml_graph_cut.h"
+
+namespace sd {
+
+    struct GraphCutLayerSplitAssignment {
+        std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
+        std::vector<int64_t> bytes_by_backend;
+        std::vector<size_t> first_segment_by_backend;
+        std::vector<size_t> last_segment_by_backend;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
+        size_t segment_count          = 0;
+        bool has_new_param_assignment = false;
+    };
+
+    std::string layer_split_backend_device_display_name(ggml_backend_t backend);
+    int layer_split_tensor_block_index(const std::string& name);
+    bool partition_graph_cut_layer_split(const char* desc,
+                                         ggml_cgraph* gf,
+                                         const sd::ggml_graph_cut::Plan& plan,
+                                         const std::vector<ggml_backend_t>& split_backends,
+                                         const std::vector<size_t>& backend_vram_limits,
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out);
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment);
+
+}  // namespace sd
+
+#endif  // __SD_CORE_LAYER_SPLIT_PARTITION_H__

otherarch/sdcpp/src/core/util.cpp

@@ -4,6 +4,8 @@
 #include <cmath>
 #include <codecvt>
 #include <cstdarg>
+#include <cstdlib>
+#include <cstring>
 #include <exception>
 #include <filesystem>
 #include <fstream>
@@ -29,6 +31,7 @@
 #include <unistd.h>
 #endif
 
+#include "ggml-backend.h"
 #include "ggml.h"
 #include "stable-diffusion.h"
 
@@ -360,6 +363,9 @@ int sd_preview_interval              = 1;
 bool sd_preview_denoised             = true;
 bool sd_preview_noisy                = false;
 
+static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
+static void* sd_backend_eval_cb_data               = nullptr;
+
 std::u32string utf8_to_utf32(const std::string& utf8_str) {
     std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
     return converter.from_bytes(utf8_str);
@@ -662,6 +668,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b
     sd_preview_noisy    = noisy;
 }
 
+void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
+    sd_backend_eval_cb      = cb;
+    sd_backend_eval_cb_data = data;
+}
+
 sd_preview_cb_t sd_get_preview_callback() {
     return sd_preview_cb;
 }
@@ -682,6 +693,14 @@ bool sd_should_preview_noisy() {
     return sd_preview_noisy;
 }
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback() {
+    return sd_backend_eval_cb;
+}
+
+void* sd_get_backend_eval_callback_data() {
+    return sd_backend_eval_cb_data;
+}
+
 sd_progress_cb_t sd_get_progress_callback() {
     return sd_progress_cb;
 }
@@ -1026,3 +1045,26 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
     }
     return result;
 }
+
+size_t sd_list_devices(char* buffer, size_t buffer_size) {
+    if (ggml_backend_dev_count() == 0) {
+        // dynamic-backend builds discover their backend modules at runtime
+        ggml_backend_load_all();
+    }
+
+    std::ostringstream oss;
+    for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        const char* name       = ggml_backend_dev_name(dev);
+        const char* desc       = ggml_backend_dev_description(dev);
+        oss << (name ? name : "") << '\t' << (desc ? desc : "") << '\n';
+    }
+
+    std::string devices = oss.str();
+    if (buffer != nullptr && buffer_size > 0) {
+        size_t copy_size = std::min(devices.size(), buffer_size - 1);
+        memcpy(buffer, devices.data(), copy_size);
+        buffer[copy_size] = '\0';
+    }
+    return devices.size();
+}

otherarch/sdcpp/src/core/util.h

@@ -101,6 +101,9 @@ int sd_get_preview_interval();
 bool sd_should_preview_denoised();
 bool sd_should_preview_noisy();
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback();
+void* sd_get_backend_eval_callback_data();
+
 // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 

otherarch/sdcpp/src/kcpp_sd_extensions.h

@@ -27,8 +27,6 @@ namespace kcpp_sd {
 
     model_info get_model_info(sd_ctx_t* ctx);
 
-    void SetCircularAxesAll(sd_ctx_t* ctx, bool circular_x, bool circular_y);
-
     void set_lora_cache(sd_ctx_t *ctx, bool enable);
 
     void apply_loras(sd_ctx_t *ctx, const std::vector<sd_lora_t>& lora_specs);

otherarch/sdcpp/src/model.h

@@ -36,6 +36,7 @@ enum SDVersion {
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
     VERSION_QWEN_IMAGE,
+    VERSION_QWEN_IMAGE_LAYERED,
     VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
@@ -46,9 +47,11 @@ enum SDVersion {
     VERSION_OVIS_IMAGE,
     VERSION_ERNIE_IMAGE,
     VERSION_LENS,
+    VERSION_MINIT2I,
     VERSION_LONGCAT,
     VERSION_PID,
     VERSION_IDEOGRAM4,
+    VERSION_SEFI_IMAGE,
     VERSION_KREA2,
     VERSION_ESRGAN,
     VERSION_COUNT,
@@ -125,7 +128,7 @@ static inline bool sd_version_is_wan(SDVersion version) {
 }
 
 static inline bool sd_version_is_qwen_image(SDVersion version) {
-    if (version == VERSION_QWEN_IMAGE) {
+    if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
         return true;
     }
     return false;
@@ -173,6 +176,13 @@ static inline bool sd_version_is_lens(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_minit2i(SDVersion version) {
+    if (version == VERSION_MINIT2I) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_pid(SDVersion version) {
     if (version == VERSION_PID) {
         return true;
@@ -187,6 +197,13 @@ static inline bool sd_version_is_ideogram4(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_sefi_image(SDVersion version) {
+    if (version == VERSION_SEFI_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_krea2(SDVersion version) {
     if (version == VERSION_KREA2) {
         return true;
@@ -202,7 +219,14 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_flux2_vae(SDVersion version) {
-    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_uses_wan_vae(SDVersion version) {
+    if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
         return true;
     }
     return false;
@@ -232,9 +256,11 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_boogu_image(version) ||
         sd_version_is_ernie_image(version) ||
         sd_version_is_lens(version) ||
+        sd_version_is_minit2i(version) ||
         sd_version_is_longcat(version) ||
         sd_version_is_pid(version) ||
         sd_version_is_ideogram4(version) ||
+        sd_version_is_sefi_image(version) ||
         sd_version_is_krea2(version)) {
         return true;
     }

otherarch/sdcpp/src/model/common/rope.hpp

@@ -12,6 +12,16 @@ namespace Rope {
         ErnieImage,
     };
 
+    enum class RefIndexMode {
+        FIXED,
+        INCREASE,
+        DECREASE,
+    };
+
+    __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
+        return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
+    }
+
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -346,7 +356,7 @@ namespace Rope {
                                                                    int axes_dim_num,
                                                                    int start_index,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool scale_rope,
                                                                    int base_offset = 0) {
@@ -357,13 +367,15 @@ namespace Rope {
         for (ggml_tensor* ref : ref_latents) {
             int h_offset = 0;
             int w_offset = 0;
-            if (!increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::FIXED) {
                 if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
                     w_offset = curr_w_offset;
                 } else {
                     h_offset = curr_h_offset;
                 }
                 scale_rope = false;
+            } else if (ref_index_mode == RefIndexMode::DECREASE) {
+                index--;
             }
 
             auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@@ -377,7 +389,7 @@ namespace Rope {
                                             scale_rope);
             ids          = concat_ids(ids, ref_ids, bs);
 
-            if (increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::INCREASE) {
                 index++;
             }
 
@@ -395,7 +407,7 @@ namespace Rope {
                                                                    int context_len,
                                                                    std::set<int> txt_arange_dims,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool is_longcat) {
         int x_index = is_longcat ? 1 : 0;
@@ -406,7 +418,7 @@ namespace Rope {
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -420,7 +432,7 @@ namespace Rope {
                                                      int context_len,
                                                      std::set<int> txt_arange_dims,
                                                      const std::vector<ggml_tensor*>& ref_latents,
-                                                     bool increase_ref_index,
+                                                     RefIndexMode ref_index_mode,
                                                      float ref_index_scale,
                                                      int theta,
                                                      bool circular_h,
@@ -435,7 +447,7 @@ namespace Rope {
                                                            context_len,
                                                            txt_arange_dims,
                                                            ref_latents,
-                                                           increase_ref_index,
+                                                           ref_index_mode,
                                                            ref_index_scale,
                                                            is_longcat);
         std::vector<std::vector<int>> wrap_dims;
@@ -481,17 +493,64 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
+                                                                  int h,
+                                                                  int w,
+                                                                  int pt,
+                                                                  int ph,
+                                                                  int pw,
+                                                                  int bs,
+                                                                  int t_offset    = 0,
+                                                                  int h_offset    = 0,
+                                                                  int w_offset    = 0,
+                                                                  bool scale_rope = false) {
+        int t_len = (t + (pt / 2)) / pt;
+        int h_len = (h + (ph / 2)) / ph;
+        int w_len = (w + (pw / 2)) / pw;
+
+        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
+
+        if (scale_rope) {
+            h_offset -= h_len / 2;
+            w_offset -= w_len / 2;
+        }
+
+        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
+        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
+        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
+
+        for (int i = 0; i < t_len; ++i) {
+            for (int j = 0; j < h_len; ++j) {
+                for (int k = 0; k < w_len; ++k) {
+                    int idx         = i * h_len * w_len + j * w_len + k;
+                    vid_ids[idx][0] = t_ids[i];
+                    vid_ids[idx][1] = h_ids[j];
+                    vid_ids[idx][2] = w_ids[k];
+                }
+            }
+        }
+
+        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < vid_ids.size(); ++j) {
+                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
+            }
+        }
+        return vid_ids_repeated;
+    }
+
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
+                                                                         int h,
                                                                          int w,
                                                                          int patch_size,
                                                                          int bs,
                                                                          int context_len,
                                                                          const std::vector<ggml_tensor*>& ref_latents,
-                                                                         bool increase_ref_index) {
+                                                                         RefIndexMode ref_index_mode) {
         int h_len        = (h + (patch_size / 2)) / patch_size;
         int w_len        = (w + (patch_size / 2)) / patch_size;
-        int txt_id_start = std::max(h_len, w_len);
-        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
+        int txt_id_start = std::max(h_len, w_len) / 2;
+        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len);
         std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
         for (int i = 0; i < bs; ++i) {
             for (int j = 0; j < txt_ids.size(); ++j) {
@@ -499,28 +558,30 @@ namespace Rope {
             }
         }
         int axes_dim_num = 3;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
+        auto img_ids     = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
-            ids           = concat_ids(ids, refs_ids, bs);
+            int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
+            auto refs_ids       = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
+            ids                 = concat_ids(ids, refs_ids, bs);
         }
         return ids;
     }
 
     // Generate qwen_image positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
+    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
+                                                           int h,
                                                            int w,
                                                            int patch_size,
                                                            int bs,
                                                            int context_len,
                                                            const std::vector<ggml_tensor*>& ref_latents,
-                                                           bool increase_ref_index,
+                                                           RefIndexMode ref_index_mode,
                                                            int theta,
                                                            bool circular_h,
                                                            bool circular_w,
                                                            const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
@@ -533,7 +594,7 @@ namespace Rope {
                 // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
                 wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
                 size_t cursor           = context_len;  // ignore text tokens
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
                 for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
                     if (circular_h) {
                         wrap_dims[1][cursor + token_i] = h_len;
@@ -684,46 +745,6 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
-                                                                  int h,
-                                                                  int w,
-                                                                  int pt,
-                                                                  int ph,
-                                                                  int pw,
-                                                                  int bs,
-                                                                  int t_offset = 0,
-                                                                  int h_offset = 0,
-                                                                  int w_offset = 0) {
-        int t_len = (t + (pt / 2)) / pt;
-        int h_len = (h + (ph / 2)) / ph;
-        int w_len = (w + (pw / 2)) / pw;
-
-        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
-
-        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
-        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
-        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
-
-        for (int i = 0; i < t_len; ++i) {
-            for (int j = 0; j < h_len; ++j) {
-                for (int k = 0; k < w_len; ++k) {
-                    int idx         = i * h_len * w_len + j * w_len + k;
-                    vid_ids[idx][0] = t_ids[i];
-                    vid_ids[idx][1] = h_ids[j];
-                    vid_ids[idx][2] = w_ids[k];
-                }
-            }
-        }
-
-        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
-        for (int i = 0; i < bs; ++i) {
-            for (int j = 0; j < vid_ids.size(); ++j) {
-                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
-            }
-        }
-        return vid_ids_repeated;
-    }
-
     // Generate wan positional embeddings
     __STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
                                                     int h,
@@ -785,7 +806,8 @@ namespace Rope {
                                                                       int context_len,
                                                                       int seq_multi_of,
                                                                       const std::vector<ggml_tensor*>& ref_latents,
-                                                                      bool increase_ref_index) {
+                                                                      RefIndexMode ref_index_mode) {
+        SD_UNUSED(ref_index_mode);
         int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
         auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
         for (int i = 0; i < bs * padded_context_len; i++) {
@@ -816,12 +838,12 @@ namespace Rope {
                                                         int context_len,
                                                         int seq_multi_of,
                                                         const std::vector<ggml_tensor*>& ref_latents,
-                                                        bool increase_ref_index,
+                                                        RefIndexMode ref_index_mode,
                                                         int theta,
                                                         bool circular_h,
                                                         bool circular_w,
                                                         const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
             int pad_h = (patch_size - (h % patch_size)) % patch_size;

otherarch/sdcpp/src/model/diffusion/anima.hpp

@@ -227,7 +227,6 @@ namespace Anima {
             k4 = k_norm->forward(ctx, k4);
 
             ggml_tensor* attn_out = nullptr;
-            float scale           = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
             if (pe_q != nullptr || pe_k != nullptr) {
                 if (pe_q == nullptr) {
                     pe_q = pe_k;
@@ -245,8 +244,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      true,
-                                                     ctx->flash_attn_enabled,
-                                                     scale);
+                                                     ctx->flash_attn_enabled);
             } else {
                 auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
                 auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
@@ -258,8 +256,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      false,
-                                                     ctx->flash_attn_enabled,
-                                                     scale);
+                                                     ctx->flash_attn_enabled);
             }
 
             return out_proj->forward(ctx, attn_out);
@@ -615,7 +612,7 @@ namespace Anima {
                                           0,
                                           {},
                                           empty_ref_latents,
-                                          false,
+                                          Rope::RefIndexMode::FIXED,
                                           1.0f,
                                           false);
 

otherarch/sdcpp/src/model/diffusion/control.hpp

@@ -1,11 +1,12 @@
-#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
 
-#define CONTROL_NET_GRAPH_SIZE 1536
+// 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 ===================================

otherarch/sdcpp/src/model/diffusion/dit.hpp

@@ -135,23 +135,23 @@ namespace DiT {
         return x;
     }
 
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
-                                   ggml_tensor* x,
-                                   int64_t t_len,
-                                   int64_t h_len,
-                                   int64_t w_len,
-                                   int pt,
-                                   int ph,
-                                   int pw) {
-        // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
+    inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
+                                      ggml_tensor* x,
+                                      int64_t t_len,
+                                      int64_t h_len,
+                                      int64_t w_len,
+                                      int pt,
+                                      int ph,
+                                      int pw) {
+        // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
-        int64_t N = x->ne[3];
+        int64_t N = x->ne[2];
         int64_t C = x->ne[0] / pt / ph / pw;
 
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-        x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, 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, 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]

otherarch/sdcpp/src/model/diffusion/ernie_image.hpp

@@ -162,8 +162,6 @@ namespace ErnieImage {
             int64_t S = x->ne[1];
             int64_t N = x->ne[2];
 
-            float scale = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
-
             auto q = to_q->forward(ctx, x);
             auto k = to_k->forward(ctx, x);
             auto v = to_v->forward(ctx, x);
@@ -184,7 +182,7 @@ namespace ErnieImage {
             k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, heads, S, head_dim]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled, scale);  // [N, S, hidden_size]
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, S, hidden_size]
             x = to_out_0->forward(ctx, x);
             return x;
         }

otherarch/sdcpp/src/model/diffusion/flux.hpp

@@ -4,10 +4,12 @@
 #include <memory>
 #include <vector>
 
+#include "core/util.h"
 #include "model/adapter/pulid.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
+#include "model/diffusion/sefi_image.hpp"
 #include "model_loader.h"
 
 #define FLUX_GRAPH_SIZE 10240
@@ -26,6 +28,9 @@ namespace Flux {
     struct FluxConfig {
         SDVersion version         = VERSION_FLUX;
         bool is_chroma            = false;
+        bool is_sefi              = false;
+        int64_t semantic_channels = 0;
+        float sefi_delta_t        = 0.1f;
         int patch_size            = 2;
         int64_t in_channels       = 64;
         int64_t out_channels      = 64;
@@ -88,6 +93,21 @@ namespace Flux {
                 config.share_modulation = true;
                 config.ref_index_scale  = 10.f;
                 config.use_mlp_silu_act = true;
+            } else if (sd_version_is_sefi_image(version)) {
+                config.is_sefi           = true;
+                config.semantic_channels = 16;
+                config.in_channels       = 128 + config.semantic_channels;
+                config.patch_size        = 1;
+                config.out_channels      = 128 + config.semantic_channels;
+                config.mlp_ratio         = 3.f;
+                config.theta             = 2000;
+                config.axes_dim          = {32, 32, 32, 32};
+                config.vec_in_dim        = 0;
+                config.qkv_bias          = false;
+                config.disable_bias      = true;
+                config.share_modulation  = true;
+                config.ref_index_scale   = 10.f;
+                config.use_mlp_silu_act  = true;
             } else if (sd_version_is_longcat(version)) {
                 config.context_in_dim = 3584;
                 config.vec_in_dim     = 0;
@@ -723,8 +743,8 @@ namespace Flux {
 
                 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);
-                shift      = m_vec[0];  // [N, hidden_size]
-                scale      = m_vec[1];  // [N, hidden_size]
+                shift      = m_vec[0];
+                scale      = m_vec[1];
             }
 
             x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
@@ -902,6 +922,8 @@ namespace Flux {
             }
             if (config.is_chroma) {
                 blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(config.in_dim, config.hidden_size);
+            } else if (config.is_sefi) {
+                blocks["dual_time_embed"] = std::make_shared<SefiImage::SefiDualTimestepEmbeddings>(256, config.hidden_size);
             } else {
                 blocks["time_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
                 if (config.vec_in_dim > 0) {
@@ -1027,6 +1049,11 @@ namespace Flux {
                 if (y != nullptr) {
                     txt_img_mask = ggml_pad(ctx->ggml_ctx, y, static_cast<int>(img->ne[1]), 0, 0, 0);
                 }
+            } else if (config.is_sefi) {
+                auto dual_time_embed = std::dynamic_pointer_cast<SefiImage::SefiDualTimestepEmbeddings>(blocks["dual_time_embed"]);
+                auto timestep_sem    = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, 0);
+                auto timestep_tex    = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, ggml_element_size(timesteps));
+                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));
@@ -1374,18 +1401,28 @@ namespace Flux {
         std::vector<float> dct_vec;
         sd::Tensor<float> guidance_tensor;
         SDVersion version;
-        bool use_mask = false;
+        bool use_mask = true;
 
         FluxRunner(ggml_backend_t backend,
                    const String2TensorStorage& tensor_storage_map      = {},
                    const std::string prefix                            = "",
                    SDVersion version                                   = VERSION_FLUX,
-                   bool use_mask                                       = false,
-                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                   const char* model_args                              = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
               config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
-              version(version),
-              use_mask(use_mask) {
+              version(version) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "chroma_use_dit_mask") {
+                    bool parsed = true;
+                    if (parse_strict_bool(value, parsed)) {
+                        use_mask = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid Chroma DiT model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+
             if (config.is_chroma) {
                 LOG_INFO("Using pruned modulation (Chroma)");
             }
@@ -1459,7 +1496,7 @@ namespace Flux {
                                  const sd::Tensor<float>& y_tensor                        = {},
                                  const sd::Tensor<float>& guidance_tensor                 = {},
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false,
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED,
                                  std::vector<int> skip_layers                             = {},
                                  const sd::Tensor<float>& pulid_id_tensor                 = {},
                                  float pulid_id_weight                                    = 1.0f) {
@@ -1500,9 +1537,9 @@ namespace Flux {
                 set_backend_tensor_data(mod_index_arange, mod_index_arange_vec.data());
             }
             std::set<int> txt_arange_dims;
-            if (sd_version_is_flux2(version)) {
-                txt_arange_dims    = {3};
-                increase_ref_index = true;
+            if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
+                txt_arange_dims = {3};
+                ref_index_mode  = Rope::RefIndexMode::INCREASE;
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
@@ -1513,7 +1550,7 @@ namespace Flux {
                                             static_cast<int>(context->ne[1]),
                                             txt_arange_dims,
                                             ref_latents,
-                                            increase_ref_index,
+                                            ref_index_mode,
                                             config.ref_index_scale,
                                             config.theta,
                                             circular_y_enabled,
@@ -1573,7 +1610,7 @@ namespace Flux {
                                   const sd::Tensor<float>& y                        = {},
                                   const sd::Tensor<float>& guidance                 = {},
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false,
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED,
                                   std::vector<int> skip_layers                      = std::vector<int>(),
                                   const sd::Tensor<float>& pulid_id                 = {},
                                   float pulid_id_weight                             = 1.0f) {
@@ -1584,7 +1621,7 @@ namespace Flux {
             // guidance: [N, ]
             // pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
+                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
             auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -1606,7 +1643,7 @@ namespace Flux {
                            tensor_or_empty(diffusion_params.y),
                            tensor_or_empty(extra->guidance),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index,
+                           diffusion_params.ref_index_mode,
                            extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
                            tensor_or_empty(extra->pulid_id),
                            extra->pulid_id_weight);
@@ -1657,7 +1694,7 @@ namespace Flux {
                                        {},
                                        guidance,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());
@@ -1692,7 +1729,6 @@ namespace Flux {
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
                                                                             VERSION_FLUX2,
-                                                                            false,
                                                                             model_manager);
 
             if (!model_manager->register_runner_params("Flux test",

otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
+#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 #define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 
 #include <algorithm>

otherarch/sdcpp/src/model/diffusion/minit2i.hpp

@@ -0,0 +1,611 @@
+#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__
+#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+namespace MiniT2I {
+    constexpr int MINIT2I_GRAPH_SIZE = 196608;
+
+    struct MiniT2IConfig {
+        int64_t image_size         = 512;
+        int64_t patch_size         = 16;
+        int64_t in_channels        = 3;
+        int64_t txt_input_size     = 1024;
+        int64_t hidden_size        = 768;
+        int64_t txt_hidden_size    = 768;
+        int64_t cond_vec_size      = 768;
+        int64_t depth_double       = 17;
+        int64_t txt_preamble_depth = 2;
+        int64_t num_heads          = 12;
+        int64_t head_dim           = 64;
+        float mlp_ratio            = 2.6667f;
+        int64_t pca_channels       = 128;
+        int64_t prompt_length      = 256;
+        int64_t n_T                = 100;
+        float cfg_interval_start   = 0.0f;
+        float cfg_interval_end     = 1.0f;
+
+        static MiniT2IConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            MiniT2IConfig config;
+            config.depth_double       = 0;
+            config.txt_preamble_depth = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "img_embedder.proj1.weight") && tensor_storage.n_dims == 4) {
+                    config.patch_size   = tensor_storage.ne[0];
+                    config.in_channels  = tensor_storage.ne[2];
+                    config.pca_channels = tensor_storage.ne[3];
+                } else if (ends_with(name, "img_embedder.proj2.weight") && tensor_storage.n_dims == 4) {
+                    config.pca_channels = tensor_storage.ne[2];
+                    config.hidden_size  = tensor_storage.ne[3];
+                } else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
+                    config.txt_input_size  = tensor_storage.ne[0];
+                    config.txt_hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
+                    config.cond_vec_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
+                    int64_t inner3     = tensor_storage.ne[1];
+                    int64_t inner      = inner3 / 3;
+                    config.hidden_size = tensor_storage.ne[0];
+                    if (config.hidden_size == 768) {
+                        config.num_heads = 12;
+                        config.head_dim  = 64;
+                    } else if (config.hidden_size == 1248) {
+                        config.num_heads = 24;
+                        config.head_dim  = 52;
+                    } else if (inner > 0) {
+                        config.head_dim  = 64;
+                        config.num_heads = std::max<int64_t>(1, inner / config.head_dim);
+                    }
+                } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area = config.patch_size * config.patch_size;
+                    config.hidden_size = tensor_storage.ne[0];
+                    config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
+                } else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
+                    config.prompt_length = tensor_storage.ne[1];
+                }
+
+                size_t pos = name.find("double_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int64_t idx         = atoi(items[1].c_str());
+                        config.depth_double = std::max<int64_t>(config.depth_double, idx + 1);
+                    }
+                }
+                pos = name.find("txt_preamble_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int64_t idx               = atoi(items[1].c_str());
+                        config.txt_preamble_depth = std::max<int64_t>(config.txt_preamble_depth, idx + 1);
+                    }
+                }
+            }
+
+            if (config.depth_double <= 0) {
+                config.depth_double = config.hidden_size == 1248 ? 23 : 17;
+            }
+            if (config.txt_preamble_depth <= 0) {
+                config.txt_preamble_depth = 2;
+            }
+            if (config.head_dim <= 0 || config.num_heads <= 0) {
+                config.head_dim  = config.hidden_size == 1248 ? 52 : 64;
+                config.num_heads = config.hidden_size / config.head_dim;
+            }
+            LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
+                      config.hidden_size,
+                      config.txt_hidden_size,
+                      config.num_heads,
+                      config.head_dim,
+                      config.depth_double,
+                      config.txt_preamble_depth,
+                      config.patch_size,
+                      config.in_channels);
+            return config;
+        }
+    };
+
+    inline std::vector<float> make_2d_sincos_pos_embed(int grid_size, int dim) {
+        GGML_ASSERT(dim % 4 == 0);
+        int half_dim = dim / 2;
+        int quarter  = half_dim / 2;
+        std::vector<float> out(static_cast<size_t>(grid_size) * grid_size * dim);
+        std::vector<float> omega(quarter);
+        for (int i = 0; i < quarter; ++i) {
+            omega[i] = 1.0f / std::pow(10000.0f, static_cast<float>(i) / static_cast<float>(quarter));
+        }
+        for (int y = 0; y < grid_size; ++y) {
+            for (int x = 0; x < grid_size; ++x) {
+                size_t base = static_cast<size_t>(y * grid_size + x) * dim;
+                for (int i = 0; i < quarter; ++i) {
+                    float ay                           = y * omega[i];
+                    float ax                           = x * omega[i];
+                    out[base + i]                      = std::sin(ax);
+                    out[base + quarter + i]            = std::cos(ax);
+                    out[base + half_dim + i]           = std::sin(ay);
+                    out[base + half_dim + quarter + i] = std::cos(ay);
+                }
+            }
+        }
+        return out;
+    }
+
+    inline std::vector<float> make_text_rope(int length, int head_dim) {
+        return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
+    }
+
+    inline std::vector<float> make_vision_rope(int side, int head_dim) {
+        GGML_ASSERT(head_dim % 4 == 0);
+        int dim     = head_dim / 2;
+        int quarter = dim / 2;
+        int length  = side * side;
+        std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
+        std::vector<float> freqs(quarter);
+        for (int i = 0; i < quarter; ++i) {
+            freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
+        }
+        for (int y = 0; y < side; ++y) {
+            for (int x = 0; x < side; ++x) {
+                int pos     = y * side + x;
+                size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
+                for (int i = 0; i < quarter; ++i) {
+                    float ay        = y * freqs[i];
+                    float ax        = x * freqs[i];
+                    float angles[2] = {ay, ax};
+                    for (int axis = 0; axis < 2; ++axis) {
+                        int j                 = axis * quarter + i;
+                        out[base + 4 * j]     = std::cos(angles[axis]);
+                        out[base + 4 * j + 1] = -std::sin(angles[axis]);
+                        out[base + 4 * j + 2] = std::sin(angles[axis]);
+                        out[base + 4 * j + 3] = std::cos(angles[axis]);
+                    }
+                }
+            }
+        }
+        return out;
+    }
+
+    struct SwiGLUMlp : public GGMLBlock {
+        SwiGLUMlp(int64_t in_features, int64_t hidden_features) {
+            int64_t hidden_dim = ((hidden_features + 7) / 8) * 8;
+            blocks["w1"]       = std::make_shared<Linear>(in_features, hidden_dim, false);
+            blocks["w3"]       = std::make_shared<Linear>(in_features, hidden_dim, false);
+            blocks["w2"]       = std::make_shared<Linear>(hidden_dim, in_features, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto w1   = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
+            auto w3   = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
+            auto w2   = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
+            auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto up   = w3->forward(ctx, x);
+            return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up));
+        }
+    };
+
+    struct BottleneckPatchEmbed : public GGMLBlock {
+        int64_t patch_size;
+
+        BottleneckPatchEmbed(int64_t patch_size, int64_t in_channels, int64_t pca_channels, int64_t hidden_size)
+            : patch_size(patch_size) {
+            blocks["proj1"] = std::make_shared<Conv2d>(in_channels,
+                                                       pca_channels,
+                                                       std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
+                                                       std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
+                                                       std::pair<int, int>{0, 0},
+                                                       std::pair<int, int>{1, 1},
+                                                       false);
+            blocks["proj2"] = std::make_shared<Conv2d>(pca_channels,
+                                                       hidden_size,
+                                                       std::pair<int, int>{1, 1},
+                                                       std::pair<int, int>{1, 1},
+                                                       std::pair<int, int>{0, 0},
+                                                       std::pair<int, int>{1, 1},
+                                                       true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto proj1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj1"]);
+            auto proj2 = std::dynamic_pointer_cast<Conv2d>(blocks["proj2"]);
+            x          = proj1->forward(ctx, x);
+            x          = proj2->forward(ctx, x);
+            x          = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+            x          = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            return x;
+        }
+    };
+
+    struct TimestepEmbedder : public GGMLBlock {
+        int frequency_embedding_size;
+
+        TimestepEmbedder(int64_t hidden_size, int frequency_embedding_size = 256)
+            : frequency_embedding_size(frequency_embedding_size) {
+            blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
+            auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
+            t_emb      = mlp_0->forward(ctx, t_emb);
+            t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
+            return mlp_2->forward(ctx, t_emb);
+        }
+    };
+
+    inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
+        int64_t N = qkv->ne[2];
+        int64_t L = qkv->ne[1];
+        auto q    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
+        auto k    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads);
+        auto v    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
+                                 qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads * 2);
+        return {q, k, v};
+    }
+
+    struct PlainTextTransformerBlock : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        PlainTextTransformerBlock(int64_t hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
+            : num_heads(num_heads), head_dim(head_dim) {
+            int64_t inner_dim   = num_heads * head_dim;
+            blocks["norm1"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["norm2"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["qkv"]       = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
+            blocks["attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
+            blocks["mlp"]       = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
+            blocks["q_norm"]    = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm"]    = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) {
+            auto norm1     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+            auto norm2     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+            auto qkv_proj  = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+            auto attn_proj = std::dynamic_pointer_cast<Linear>(blocks["attn_proj"]);
+            auto mlp       = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["mlp"]);
+            auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+
+            auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto q   = q_norm->forward(ctx, qkv[0]);
+            auto k   = k_norm->forward(ctx, qkv[1]);
+            auto v   = qkv[2];
+            auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
+            txt      = ggml_add(ctx->ggml_ctx, txt, attn_proj->forward(ctx, out));
+            txt      = ggml_add(ctx->ggml_ctx, txt, mlp->forward(ctx, norm2->forward(ctx, txt)));
+            return txt;
+        }
+    };
+
+    struct DoubleStreamDiTBlock : public GGMLBlock {
+        int64_t num_heads;
+        int64_t head_dim;
+
+        DoubleStreamDiTBlock(int64_t hidden_size, int64_t txt_hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
+            : num_heads(num_heads), head_dim(head_dim) {
+            int64_t inner_dim       = num_heads * head_dim;
+            blocks["img_norm1"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["img_norm2"]     = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["txt_norm1"]     = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
+            blocks["txt_norm2"]     = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
+            blocks["img_qkv"]       = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
+            blocks["txt_qkv"]       = std::make_shared<Linear>(txt_hidden_size, inner_dim * 3, true);
+            blocks["q_norm"]        = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm"]        = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["img_attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
+            blocks["txt_attn_proj"] = std::make_shared<Linear>(inner_dim, txt_hidden_size, true);
+            blocks["img_mlp"]       = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
+            blocks["txt_mlp"]       = std::make_shared<SwiGLUMlp>(txt_hidden_size, static_cast<int64_t>(txt_hidden_size * mlp_ratio));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* pe) {
+            auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
+            auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
+            auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
+            auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
+            auto img_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
+            auto txt_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
+            auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+            auto img_proj  = std::dynamic_pointer_cast<Linear>(blocks["img_attn_proj"]);
+            auto txt_proj  = std::dynamic_pointer_cast<Linear>(blocks["txt_attn_proj"]);
+            auto img_mlp   = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["img_mlp"]);
+            auto txt_mlp   = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["txt_mlp"]);
+
+            int64_t li = img->ne[1];
+            int64_t lt = txt->ne[1];
+
+            auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
+            auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
+
+            auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
+            auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
+            auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
+
+            auto out     = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
+            auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
+            auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
+
+            img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
+            txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
+            img = ggml_add(ctx->ggml_ctx, img, img_mlp->forward(ctx, img_norm2->forward(ctx, img)));
+            txt = ggml_add(ctx->ggml_ctx, txt, txt_mlp->forward(ctx, txt_norm2->forward(ctx, txt)));
+            return {img, txt};
+        }
+    };
+
+    struct FinalLayer : public GGMLBlock {
+        FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels) {
+            blocks["norm_final"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["linear"]     = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm_final = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_final"]);
+            auto linear     = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            return linear->forward(ctx, norm_final->forward(ctx, x));
+        }
+    };
+
+    struct MMJiT : public GGMLBlock {
+        MiniT2IConfig config;
+
+        MMJiT(const MiniT2IConfig& config)
+            : config(config) {
+            blocks["img_embedder"]    = std::make_shared<BottleneckPatchEmbed>(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size);
+            blocks["txt_embedder"]    = std::make_shared<Linear>(config.txt_input_size, config.txt_hidden_size, false);
+            blocks["t_embedder"]      = std::make_shared<TimestepEmbedder>(config.cond_vec_size);
+            blocks["pooled_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.cond_vec_size, false);
+            for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
+                blocks["txt_preamble_blocks." + std::to_string(i)] = std::make_shared<PlainTextTransformerBlock>(config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
+            }
+            for (int64_t i = 0; i < config.depth_double; ++i) {
+                blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamDiTBlock>(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
+            }
+            blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.in_channels);
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+            enum ggml_type wtype = get_type(prefix + "mask_token", tensor_storage_map, GGML_TYPE_F32);
+            params["mask_token"] = ggml_new_tensor_3d(ctx, wtype, config.txt_input_size, 1, 1);
+        }
+
+        ggml_tensor* apply_text_mask(GGMLRunnerContext* ctx, ggml_tensor* context, ggml_tensor* mask) {
+            if (mask == nullptr) {
+                return context;
+            }
+            mask            = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
+            mask            = ggml_repeat(ctx->ggml_ctx, mask, context);
+            auto keep       = ggml_mul(ctx->ggml_ctx, context, mask);
+            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
+            auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
+            return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
+        }
+
+        ggml_tensor* pool_context(GGMLRunnerContext* ctx, ggml_tensor* context) {
+            int64_t dim = context->ne[0];
+            int64_t len = context->ne[1];
+            int64_t N   = context->ne[2];
+            auto x      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3));
+            x           = ggml_reshape_3d(ctx->ggml_ctx, x, len, dim, N);
+            x           = ggml_mean(ctx->ggml_ctx, x);
+            x           = ggml_reshape_2d(ctx->ggml_ctx, x, dim, N);
+            return x;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* img,
+                             ggml_tensor* context,
+                             ggml_tensor* mask,
+                             ggml_tensor* pos_embed,
+                             ggml_tensor* txt_pe,
+                             ggml_tensor* joint_pe) {
+            auto img_embedder = std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["img_embedder"]);
+            auto txt_embedder = std::dynamic_pointer_cast<Linear>(blocks["txt_embedder"]);
+            auto final_layer  = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
+
+            int64_t W  = img->ne[0];
+            int64_t H  = img->ne[1];
+            int64_t hp = H / config.patch_size;
+            int64_t wp = W / config.patch_size;
+
+            context = apply_text_mask(ctx, context, mask);
+            auto x  = img_embedder->forward(ctx, img);
+            x       = ggml_add(ctx->ggml_ctx, x, pos_embed);
+
+            auto txt = txt_embedder->forward(ctx, context);
+            for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
+                auto block = std::dynamic_pointer_cast<PlainTextTransformerBlock>(blocks["txt_preamble_blocks." + std::to_string(i)]);
+                txt        = block->forward(ctx, txt, txt_pe);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.txt_preamble_blocks." + std::to_string(i), "txt");
+            }
+            for (int64_t i = 0; i < config.depth_double; ++i) {
+                auto block = std::dynamic_pointer_cast<DoubleStreamDiTBlock>(blocks["double_blocks." + std::to_string(i)]);
+                auto out   = block->forward(ctx, x, txt, joint_pe);
+                x          = out.first;
+                txt        = out.second;
+                sd::ggml_graph_cut::mark_graph_cut(x, "minit2i.double_blocks." + std::to_string(i), "x");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.double_blocks." + std::to_string(i), "txt");
+            }
+            auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
+            auto out      = final_layer->forward(ctx, combined);
+            auto img_out  = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
+            return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
+        }
+    };
+
+    struct MiniT2IRunner : public DiffusionModelRunner {
+        MiniT2IConfig config;
+        MMJiT model;
+        ggml_context* position_cache_ctx            = nullptr;
+        ggml_backend_buffer_t position_cache_buffer = nullptr;
+        ggml_tensor* cached_pos_embed               = nullptr;
+        ggml_tensor* cached_txt_pe                  = nullptr;
+        ggml_tensor* cached_joint_pe                = nullptr;
+        int64_t cached_img_side                     = -1;
+        int64_t cached_txt_len                      = -1;
+        int64_t cached_hidden_size                  = -1;
+        int64_t cached_head_dim                     = -1;
+
+        MiniT2IRunner(ggml_backend_t backend,
+                      const String2TensorStorage& tensor_storage_map      = {},
+                      const std::string prefix                            = "",
+                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(MiniT2IConfig::detect_from_weights(tensor_storage_map, this->prefix)),
+              model(config) {
+            model.init(params_ctx, tensor_storage_map, this->prefix);
+        }
+
+        ~MiniT2IRunner() override {
+            free_position_cache();
+        }
+
+        std::string get_desc() override {
+            return "MiniT2I";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        void free_position_cache() {
+            if (position_cache_buffer != nullptr) {
+                ggml_backend_buffer_free(position_cache_buffer);
+                position_cache_buffer = nullptr;
+            }
+            if (position_cache_ctx != nullptr) {
+                ggml_free(position_cache_ctx);
+                position_cache_ctx = nullptr;
+            }
+            cached_pos_embed   = nullptr;
+            cached_txt_pe      = nullptr;
+            cached_joint_pe    = nullptr;
+            cached_img_side    = -1;
+            cached_txt_len     = -1;
+            cached_hidden_size = -1;
+            cached_head_dim    = -1;
+        }
+
+        void ensure_position_cache(int64_t img_side, int64_t txt_len) {
+            if (cached_img_side == img_side &&
+                cached_txt_len == txt_len &&
+                cached_hidden_size == config.hidden_size &&
+                cached_head_dim == config.head_dim &&
+                cached_pos_embed != nullptr &&
+                cached_txt_pe != nullptr &&
+                cached_joint_pe != nullptr) {
+                return;
+            }
+
+            free_position_cache();
+
+            auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
+            auto txt_pe_vec    = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
+            auto img_pe_vec    = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
+            auto joint_pe_vec  = txt_pe_vec;
+            joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
+
+            ggml_init_params params;
+            params.mem_size    = static_cast<size_t>(3 * ggml_tensor_overhead());
+            params.mem_buffer  = nullptr;
+            params.no_alloc    = true;
+            position_cache_ctx = ggml_init(params);
+            GGML_ASSERT(position_cache_ctx != nullptr);
+
+            cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
+            ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
+            cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
+            ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
+            cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
+            ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
+
+            position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
+            GGML_ASSERT(position_cache_buffer != nullptr);
+            ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+            ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
+            ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
+            ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
+            ggml_backend_synchronize(runtime_backend);
+
+            cached_img_side    = img_side;
+            cached_txt_len     = txt_len;
+            cached_hidden_size = config.hidden_size;
+            cached_head_dim    = config.head_dim;
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& mask_tensor) {
+            ggml_cgraph* gf      = new_graph_custom(MINIT2I_GRAPH_SIZE);
+            ggml_tensor* x       = make_input(x_tensor);
+            ggml_tensor* context = make_input(context_tensor);
+            ggml_tensor* mask    = make_input(mask_tensor);
+            SD_UNUSED(timesteps_tensor);
+
+            int64_t W        = x->ne[0];
+            int64_t H        = x->ne[1];
+            int64_t img_side = H / config.patch_size;
+            int64_t txt_len  = context->ne[1];
+            ensure_position_cache(img_side, txt_len);
+
+            auto runner_ctx = get_context();
+            auto out        = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const sd::Tensor<float>& mask) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, mask);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            GGML_ASSERT(diffusion_params.context != nullptr);
+            const auto* extra = diffusion_extra_as<MiniT2IDiffusionExtra>(diffusion_params);
+            GGML_ASSERT(extra->mask != nullptr);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *diffusion_params.context,
+                           *extra->mask);
+        }
+    };
+}  // namespace MiniT2I
+
+#endif  // __SD_MODEL_DIFFUSION_MINIT2I_HPP__

otherarch/sdcpp/src/model/diffusion/model.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
 #define __SD_MODEL_DIFFUSION_MODEL_HPP__
 
 #include <string>
@@ -7,6 +7,7 @@
 
 #include "core/ggml_extend.hpp"
 #include "core/tensor_ggml.hpp"
+#include "model/common/rope.hpp"
 #include "model_manager.h"
 
 struct UNetDiffusionExtra {
@@ -52,6 +53,10 @@ struct LTXAVDiffusionExtra {
     const sd::Tensor<float>* video_positions = nullptr;
 };
 
+struct MiniT2IDiffusionExtra {
+    const sd::Tensor<float>* mask = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
@@ -59,7 +64,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           AnimaDiffusionExtra,
                                           WanDiffusionExtra,
                                           HiDreamO1DiffusionExtra,
-                                          LTXAVDiffusionExtra>;
+                                          LTXAVDiffusionExtra,
+                                          MiniT2IDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;
@@ -68,7 +74,7 @@ struct DiffusionParams {
     const sd::Tensor<float>* c_concat                 = nullptr;
     const sd::Tensor<float>* y                        = nullptr;
     const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    bool increase_ref_index                           = false;
+    Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED;
     DiffusionExtraParams extra                        = std::monostate{};
 };
 

otherarch/sdcpp/src/model/diffusion/qwen_image.hpp

@@ -3,7 +3,9 @@
 
 #include <memory>
 
+#include "core/util.h"
 #include "model/common/block.hpp"
+#include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
@@ -23,6 +25,7 @@ namespace Qwen {
         std::vector<int> axes_dim   = {16, 56, 56};
         int axes_dim_sum            = 128;
         bool zero_cond_t            = false;
+        bool use_additional_t_cond  = false;
 
         static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
             QwenImageConfig config;
@@ -88,19 +91,33 @@ namespace Qwen {
     };
 
     struct QwenTimestepProjEmbeddings : public GGMLBlock {
+    protected:
+        bool use_additional_t_cond = false;
+
     public:
-        QwenTimestepProjEmbeddings(int64_t embedding_dim) {
+        QwenTimestepProjEmbeddings(int64_t embedding_dim, bool use_additional_t_cond = false)
+            : use_additional_t_cond(use_additional_t_cond) {
             blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
+            if (use_additional_t_cond) {
+                blocks["addition_t_embedding"] = std::shared_ptr<GGMLBlock>(new Embedding(2, embedding_dim));
+            }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps) {
+                             ggml_tensor* timesteps,
+                             ggml_tensor* addition_t_cond = nullptr) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
 
             auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
             auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
+            if (use_additional_t_cond) {
+                GGML_ASSERT(addition_t_cond != nullptr);
+                auto addition_t_embedding = std::dynamic_pointer_cast<Embedding>(blocks["addition_t_embedding"]);
+                auto addition_t_emb       = addition_t_embedding->forward(ctx, addition_t_cond);
+                timesteps_emb             = ggml_add(ctx->ggml_ctx, timesteps_emb, addition_t_emb);
+            }
             return timesteps_emb;
         }
     };
@@ -402,7 +419,7 @@ namespace Qwen {
         QwenImageModel(QwenImageConfig config)
             : config(config) {
             int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
-            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
+            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
             blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
             blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
             blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
@@ -424,6 +441,7 @@ namespace Qwen {
         ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
                                   ggml_tensor* x,
                                   ggml_tensor* timestep,
+                                  ggml_tensor* addition_t_cond,
                                   ggml_tensor* context,
                                   ggml_tensor* pe,
                                   ggml_tensor* modulate_index = nullptr) {
@@ -434,9 +452,9 @@ namespace Qwen {
             auto norm_out        = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
             auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
 
-            auto t_emb = time_text_embed->forward(ctx, timestep);
+            auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
             if (config.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
+                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
             auto img = img_in->forward(ctx, x);
@@ -469,33 +487,50 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
                              ggml_tensor* timestep,
+                             ggml_tensor* addition_t_cond,
                              ggml_tensor* context,
                              ggml_tensor* pe,
                              std::vector<ggml_tensor*> ref_latents = {},
                              ggml_tensor* modulate_index           = nullptr) {
             // Forward pass of DiT.
-            // x: [N, C, H, W]
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timestep: [N,]
             // context: [N, L, D]
             // pe: [L, d_head/2, 2, 2]
-            // return: [N, C, H, W]
+            // return: [N, C, H, W] or [N*C, T, H, W]
 
-            int64_t W = x->ne[0];
-            int64_t H = x->ne[1];
-            int64_t C = x->ne[2];
-            int64_t N = x->ne[3];
+            int64_t W          = x->ne[0];
+            int64_t H          = x->ne[1];
+            int64_t T          = 1;
+            int64_t N          = addition_t_cond != nullptr ? addition_t_cond->ne[0] : x->ne[3];
+            bool has_time_axis = false;
+            if (x->ne[3] != 1) {
+                T             = x->ne[2];
+                has_time_axis = true;
+            }
 
-            auto img           = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
+            auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
+                input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
+                if (!has_time_axis) {
+                    return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
+                }
+                if (input->ne[3] == 1) {
+                    input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
+                }
+                return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+            };
+
+            auto img           = patchify_input(x);
             int64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
+                    ref = patchify_input(ref);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
+            auto out = forward_orig(ctx, img, timestep, addition_t_cond, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
 
             if (out->ne[1] > img_tokens) {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
@@ -503,7 +538,17 @@ namespace Qwen {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            if (has_time_axis) {
+                int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size;
+                int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
+                int h_len = static_cast<int>((H + pad_h) / config.patch_size);
+                int w_len = static_cast<int>((W + pad_w) / config.patch_size);
+                out       = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            } else {
+                out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            }
 
             return out;
         }
@@ -515,18 +560,32 @@ namespace Qwen {
         QwenImageModel qwen_image;
         std::vector<float> pe_vec;
         std::vector<float> modulate_index_vec;
+        std::vector<int32_t> additional_t_cond_vec;
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
                         const String2TensorStorage& tensor_storage_map      = {},
                         const std::string prefix                            = "",
                         SDVersion version                                   = VERSION_QWEN_IMAGE,
-                        bool zero_cond_t                                    = false,
-                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                        const char* model_args                              = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
-              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
-            config.zero_cond_t = config.zero_cond_t || zero_cond_t;
-            qwen_image         = QwenImageModel(config);
+              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
+              version(version) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "qwen_image_zero_cond_t") {
+                    bool parsed = false;
+                    if (parse_strict_bool(value, parsed)) {
+                        config.zero_cond_t = config.zero_cond_t || parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid Qwen Image model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                config.use_additional_t_cond = true;
+            }
+            qwen_image = QwenImageModel(config);
             qwen_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -542,11 +601,11 @@ namespace Qwen {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::INCREASE) {
             ggml_cgraph* gf        = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
-            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(x->ne[3] == 1 || x_tensor.dim() == 5);
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
             std::vector<ggml_tensor*> ref_latents;
@@ -555,13 +614,29 @@ namespace Qwen {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
+            int batch_size = static_cast<int>(x->ne[3]);
+            int time_len   = 1;
+            if (x_tensor.dim() == 5) {
+                time_len   = static_cast<int>(x_tensor.shape()[2]);
+                batch_size = static_cast<int>(x_tensor.shape()[4]);
+            }
+
+            ggml_tensor* addition_t_cond = nullptr;
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                additional_t_cond_vec.assign(static_cast<size_t>(batch_size), 0);
+                addition_t_cond = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, batch_size);
+                set_backend_tensor_data(addition_t_cond, additional_t_cond_vec.data());
+                ref_index_mode = Rope::RefIndexMode::DECREASE;
+            }
+
+            pe_vec      = Rope::gen_qwen_image_pe(time_len,
+                                                  static_cast<int>(x->ne[1]),
                                                   static_cast<int>(x->ne[0]),
                                                   config.patch_size,
-                                                  static_cast<int>(x->ne[3]),
+                                                  batch_size,
                                                   static_cast<int>(context->ne[1]),
                                                   ref_latents,
-                                                  increase_ref_index,
+                                                  ref_index_mode,
                                                   config.theta,
                                                   circular_y_enabled,
                                                   circular_x_enabled,
@@ -604,6 +679,7 @@ namespace Qwen {
             ggml_tensor* out = qwen_image.forward(&runner_ctx,
                                                   x,
                                                   timesteps,
+                                                  addition_t_cond,
                                                   context,
                                                   pe,
                                                   ref_latents,
@@ -619,12 +695,12 @@ namespace Qwen {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false) {
-            // x: [N, in_channels, h, w]
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::INCREASE) {
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -640,7 +716,7 @@ namespace Qwen {
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index);
+                           diffusion_params.ref_index_mode);
         }
 
         void test() {
@@ -674,7 +750,7 @@ namespace Qwen {
                                        timesteps,
                                        context,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());
@@ -709,7 +785,6 @@ namespace Qwen {
                                                                                             tensor_storage_map,
                                                                                             "model.diffusion_model",
                                                                                             VERSION_QWEN_IMAGE,
-                                                                                            false,
                                                                                             model_manager);
 
             if (!model_manager->register_runner_params("Qwen image test",

otherarch/sdcpp/src/model/diffusion/sefi_image.hpp

@@ -0,0 +1,91 @@
+#ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
+
+#include <memory>
+
+#include "model/common/block.hpp"
+
+namespace SefiImage {
+    struct SefiImageConfig {
+        int64_t semantic_channels        = 16;
+        int64_t texture_latent_channels  = 32;
+        int64_t timestep_guidance_in_dim = 256;
+        int64_t hidden_size              = 3072;
+        float timestep_shift_alpha       = 0.3f;
+        float delta_t                    = 0.1f;
+
+        int64_t packed_texture_channels(int patch_size) const {
+            return texture_latent_channels * patch_size * patch_size;
+        }
+
+        int64_t packed_input_channels(int patch_size) const {
+            return semantic_channels + packed_texture_channels(patch_size);
+        }
+
+        static SefiImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                   const std::string& prefix) {
+            SefiImageConfig config;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
+                    config.timestep_guidance_in_dim = tensor_storage.ne[0];
+                    config.hidden_size              = tensor_storage.ne[1] * 2;
+                }
+            }
+            LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64,
+                      config.semantic_channels,
+                      config.texture_latent_channels,
+                      config.hidden_size);
+            return config;
+        }
+    };
+
+    struct SefiTimestepEmbedding : public GGMLBlock {
+    public:
+        SefiTimestepEmbedding(int64_t in_channels, int64_t time_embed_dim) {
+            blocks["linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim, false));
+            blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim, false));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* sample) {
+            auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
+
+            sample = linear_1->forward(ctx, sample);
+            sample = ggml_silu_inplace(ctx->ggml_ctx, sample);
+            sample = linear_2->forward(ctx, sample);
+            return sample;
+        }
+    };
+
+    struct SefiDualTimestepEmbeddings : public GGMLBlock {
+    public:
+        SefiDualTimestepEmbeddings(int64_t in_channels, int64_t embedding_dim) {
+            GGML_ASSERT(embedding_dim % 2 == 0);
+            int64_t half_dim            = embedding_dim / 2;
+            blocks["semantic_embedder"] = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
+            blocks["texture_embedder"]  = std::make_shared<SefiTimestepEmbedding>(in_channels, half_dim);
+            timestep_guidance_in_dim    = in_channels;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* timestep_sem,
+                             ggml_tensor* timestep_tex) {
+            auto semantic_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["semantic_embedder"]);
+            auto texture_embedder  = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["texture_embedder"]);
+
+            auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto sem_emb  = semantic_embedder->forward(ctx, sem_proj);
+            auto tex_emb  = texture_embedder->forward(ctx, tex_proj);
+            return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0);
+        }
+
+    private:
+        int64_t timestep_guidance_in_dim = 256;
+    };
+}  // namespace SefiImage
+
+#endif  // __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__