Draft: portable GGML GPU offload for Vulkan, Metal, and CUDA

CMakeLists.txt

@@ -17,6 +17,7 @@ add_subdirectory(ggml)
 
 # ── Core library (diffusion + autoregressive) ─────────────────
 add_library(diffuse
+    src/diffuse-backend.cpp
     src/diffuse-model.cpp
     src/diffuse-graph.cpp
     src/diffuse-sampler.cpp
@@ -53,6 +54,30 @@ endif()
 # ── Tests ───────────────────────────────────────────────────────
 if(DIFFUSE_BUILD_TESTS)
     enable_testing()
+    find_package(Python3 COMPONENTS Interpreter QUIET)
+    if(Python3_Interpreter_FOUND)
+        add_test(
+            NAME test-validate-logits
+            COMMAND ${Python3_EXECUTABLE}
+                    ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-validate-logits.py)
+        add_test(
+            NAME test-generate
+            COMMAND ${Python3_EXECUTABLE}
+                    ${CMAKE_CURRENT_SOURCE_DIR}/tests/test-generate.py)
+    endif()
+
+    add_executable(test-cache tests/test-cache.cpp)
+    target_link_libraries(test-cache PRIVATE diffuse)
+    add_test(NAME test-cache COMMAND test-cache)
+
+    add_executable(test-ar tests/test-ar.cpp)
+    target_link_libraries(test-ar PRIVATE diffuse)
+    add_test(NAME test-ar COMMAND test-ar)
+
+    add_executable(test-backend tests/test-backend.cpp)
+    target_link_libraries(test-backend PRIVATE diffuse)
+    add_test(NAME test-backend COMMAND test-backend)
+
     add_executable(test-forward tests/test-forward.cpp)
     target_link_libraries(test-forward PRIVATE diffuse)
     add_test(NAME test-forward COMMAND test-forward)

README.md

@@ -119,6 +119,40 @@ cmake --build build -j$(nproc)
     --remasking entropy_exit
 ```
 
+### GPU offload (experimental)
+
+Build GGML with the backend for your platform. Metal is enabled by default on
+macOS; Vulkan and CUDA are opt-in:
+
+```bash
+# Linux / Windows with Vulkan
+cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_VULKAN=ON
+
+# NVIDIA CUDA
+cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_CUDA=ON
+
+cmake --build build -j
+```
+
+List the exact GGML device names available in that build, then select one:
+
+```bash
+./build/diffuse-cli --list-devices
+./build/diffuse-cli -m model.gguf --device Vulkan0 --tokens "..." -n 64 -s 16
+```
+
+The selected device holds the model weights and receives the supported graph
+operations. The CLI reports model allocation and per-backend graph node counts
+so an offload cannot silently become CPU-only execution. Integrated GPUs are
+valid devices and are reported separately as `IGPU`.
+
+GPU offload supports full and inter-step-cached masked-diffusion graphs, plus
+autoregressive prefill, decoding, layer profiling, and speculative decoding.
+Use `--device` with `diffuse-ar` as well; speculative decoding also accepts a
+separate `--draft-device`. The full model must fit the selected device's buffer
+type. The GGML scheduler keeps a CPU backend available for unsupported graph
+operations and reports any such placement.
+
 **Note**: diffuse-cpp operates on token IDs, not raw text. Use the HuggingFace transformers library to tokenize your prompts:
 
 ```python
@@ -301,7 +335,8 @@ Current limitations:
 - No integrated tokenizer (use transformers)
 - Default 256 generated tokens per call (configurable via -n flag)
 - Single-model inference only (no batching)
-- CPU-only (GPU support via GGML is possible but not prioritized)
+- GPU offload is experimental and currently requires the full model to fit the
+  selected device's buffer type
 
 ## Contributing
 

include/diffuse.h

@@ -2,6 +2,7 @@
 
 // diffuse-cpp public API
 
+#include <cstddef>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -11,6 +12,26 @@
 struct diffuse_model;
 struct diffuse_context;
 
+enum class diffuse_device_type {
+    CPU,
+    GPU,
+    IGPU,
+    ACCEL,
+    UNKNOWN,
+};
+
+struct diffuse_device_info {
+    std::string         name;
+    std::string         description;
+    diffuse_device_type type         = diffuse_device_type::UNKNOWN;
+    size_t              memory_free  = 0;
+    size_t              memory_total = 0;
+    bool                async        = false;
+};
+
+std::vector<diffuse_device_info> diffuse_available_devices();
+const char * diffuse_device_type_name(diffuse_device_type type);
+
 // ── Hyperparameters ────────────────────────────────────────────
 struct diffuse_hparams {
     uint32_t n_vocab      = 0;
@@ -39,7 +60,7 @@ enum class diffuse_remasking {
     RANDOM,
     ENTROPY_EXIT,   // Unmask all low-entropy tokens early (semantic scheduling)
     MASKGIT_PLUS,   // Dream: unmask highest top-1 confidence (similar to LOW_CONFIDENCE)
-    TOPK_MARGIN,    // Dream: unmask by margin between top-1 and top-2 logits
+    TOPK_MARGIN,    // Dream: unmask by probability margin between top-1 and top-2
 };
 
 struct diffuse_sampler_params {
@@ -69,6 +90,8 @@ using diffuse_step_callback = std::function<void(
 
 // ── Model API ──────────────────────────────────────────────────
 diffuse_model * diffuse_model_load(const std::string & path, int n_threads);
+diffuse_model * diffuse_model_load(const std::string & path, int n_threads,
+                                   const std::string & device_name);
 void            diffuse_model_free(diffuse_model * model);
 const diffuse_hparams & diffuse_model_hparams(const diffuse_model * model);