Add Windows ARM64 and Qualcomm Hexagon NPU support

CMakeLists.txt

@@ -89,6 +89,8 @@ endif()
 
 # --- core library ----------------------------------------------------------
 add_library(trellis_core STATIC
+  src/trellis_debug.cpp
+  src/trellis_sched.cpp
   src/trellis_model.cpp
   src/dit.cpp
   src/flow_runner.cpp
@@ -263,6 +265,10 @@ add_executable(trellis-smoke src/smoke_test.cpp)
 target_link_libraries(trellis-smoke PRIVATE trellis_core)
 set_target_properties(trellis-smoke PROPERTIES BUILD_RPATH "${GGML_RPATH}")
 
+add_executable(trellis-devices src/list_devices.cpp)
+target_link_libraries(trellis-devices PRIVATE trellis_core)
+set_target_properties(trellis-devices PROPERTIES BUILD_RPATH "${GGML_RPATH}")
+
 add_executable(trellis-decode-replay src/decode_replay.cpp)
 target_link_libraries(trellis-decode-replay PRIVATE trellis_core)
 set_target_properties(trellis-decode-replay PROPERTIES BUILD_RPATH "${GGML_RPATH}")

README.md

@@ -118,6 +118,13 @@ The most useful ones:
 | `--atlas PX` | UV atlas size (default 2048 @1024 / 1024 @512) |
 | `--box-uv` | voxel-native 6-way box projection instead of the default xatlas unwrap (O(faces), faster, looser packing) |
 | `--seed N` | RNG seed |
+| `--steps N` | flow sampler steps (default 12; lower values are useful for backend bring-up) |
+| `--backend NAME` | select a registered ggml backend explicitly, such as `HTP`, `Vulkan`, or `CPU` |
+| `--threads N` | CPU threads (defaults to the detected hardware thread count) |
+| `--sched on\|off` | enable multi-backend scheduling; automatic for the partial-coverage HTP backend |
+| `--vulkan-fallback` | with HTP, try Vulkan before the CPU for unsupported operations |
+| `--fa-fast` / `--fa-f32` | select fast F16 or BF16/F32 FlashAttention accumulation |
+| `--verbose` | graph timings and progress heartbeats for long-running stages |
 | `--require-gpu` | fail instead of falling back to the (very slow, RAM-hungry) CPU path |
 
 The postprocess matches the reference pipeline op for op (see
@@ -132,8 +139,8 @@ inpaint port, and exported as a GLB with smooth normals and **lossy-WebP texture
 (`EXT_texture_webp`; PNG fallback when built with `-DTRELLIS_WEBP=OFF`). Output
 quality is at parity with the reference CUDA postprocess on identical inputs.
 
-`TRELLIS_DBG_*` environment variables toggle developer debug logging only; no
-behavior-driving environment variables remain — use the flags above.
+Prefer the CLI flags above. The corresponding `TRELLIS_*` environment variables remain
+available for test binaries and backwards-compatible automation.
 
 ### trellis-server
 
@@ -193,7 +200,10 @@ The 1024 cascade runs on a 16 GB card thanks to **FlashAttention with padded K/V
 the sparse-structure stage, and at the HR token count (≈53k) ggml's tiled FA NaN'd on
 the unpadded last key-tile — zero-padding K/V to a 256 multiple + BF16 fixes both.
 f16 compute is the default and matches torch (`--f32` forces f32; `--no-fa` restores
-the plain-softmax path for A/B testing).
+the plain-softmax path for A/B testing). On the Qualcomm HTP backend, FlashAttention
+defaults to F16 K/V and fast accumulation after a 12-step quality gate showed a 3.53x
+end-to-end speedup with identical sparse voxels; `--fa-f32` restores BF16 K/V + F32
+accumulation. Other backends retain BF16/F32 by default (`--fa-fast` forces the HTP mode).
 
 Every neural component is validated against PyTorch (the `trellis-test-*` binaries +
 `tools/ref_*.py`): SS sampler matches torch to rel 4.3e-3 (exact voxel match), DiT
@@ -264,6 +274,35 @@ See `.github/workflows/release.yml` for the exact flags the release binaries use
 6.0/6.1/7.0); the standalone installers select it automatically for devices such
 as the Tesla P100.
 
+### Windows ARM64 and Qualcomm HTP
+
+The native ARM64 helper configures clang for Windows-on-ARM and can build CPU, Vulkan,
+or Hexagon variants. For HTP, install the Hexagon SDK using the official
+[llama.cpp Windows Snapdragon guide](https://github.com/ggml-org/llama.cpp/blob/master/docs/backend/snapdragon/windows.md).
+The trimmed `hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz` package described there is sufficient;
+extract it anywhere and pass the directory containing `hexagon_sdk.json` to
+`-HexagonSdk`:
+
+```powershell
+scripts\build-arm64.ps1 -Backend cpu
+scripts\build-arm64.ps1 -Backend vulkan
+scripts\build-arm64.ps1 -Backend hexagon -HexagonSdk C:\Qualcomm\Hexagon_SDK\6.6.0.0
+```
+
+Windows requires the generated HTP Ops libraries and catalog to be signed with a trusted
+certificate before the NPU driver will load them. The linked guide documents the required
+driver, certificate, test-signing, and `HEXAGON_HTP_CERT` setup.
+
+List registered devices with `trellis-devices --init`. A typical NPU run uses:
+
+```powershell
+build-arm64-hexagon\trellis-cli.exe input.png output.glb `
+  --models models\q4 --backend HTP --sched on --verbose
+```
+
+Add `--vulkan-fallback` when the same build includes Vulkan and unsupported HTP operations
+should prefer the GPU over the CPU.
+
 ## Layout
 
 ```

include/flow_runner.h

@@ -1,5 +1,6 @@
 // Flow-DiT runner (dense grid or sparse voxel) + FlowEuler guidance-interval sampler.
 #pragma once
+#include <memory>
 #include <vector>
 #include <array>
 #include <functional>
@@ -14,6 +15,7 @@ typedef struct ggml_gallocr* ggml_gallocr_t;
 
 namespace trellis {
 struct Model;
+class GraphExec;
 
 struct SamplerParams {
     int   steps             = 12;
@@ -38,11 +40,14 @@ class DitRunner {
     int N() const { return N_; }
 private:
     const Model& m_; DiTParams p_; int N_, Lc_;
-    ggml_context* ctx_ = nullptr; ggml_cgraph* g_ = nullptr; ggml_gallocr_t alloc_ = nullptr;
+    ggml_context* ctx_ = nullptr; ggml_cgraph* g_ = nullptr;
+    // Owns graph allocation and execution for both direct and scheduled backends.
+    std::unique_ptr<GraphExec> exec_;
     ggml_tensor *gh0_, *gtf_, *gcond_, *gcos_, *gsin_, *gout_;
     std::vector<float> rcos_, rsin_;   // re-uploaded each forward (gallocr may reuse input buffers)
     std::map<std::string, ggml_tensor*> inter_;   // [dbg] named intermediates for NaN localization
     bool dbg_nan_ = false, dbg_done_ = false;
+    int  n_fwd_ = 0;
 };
 
 // Dense factory: RoPE from R^3 grid (ij meshgrid, z fastest). N = R^3.

include/trellis_args.h

@@ -9,8 +9,10 @@ namespace trellis {
 // binaries (which don't parse args) keep their historical TRELLIS_* behavior.
 extern bool g_sparse_cast_f32;  // defined in sparse.cpp        (TRELLIS_F32)
 extern bool g_no_fa;            // defined in dit.cpp           (TRELLIS_NOFA)
+extern int  g_fa_fast;          // defined in dit.cpp: -1 auto/env, 0 BF16/F32, 1 F16/fast
 extern bool g_require_gpu;      // defined in trellis_model.cpp (TRELLIS_REQUIRE_GPU)
 extern int  g_cpu_threads;      // defined in trellis_model.cpp (TRELLIS_THREADS)
+extern std::string g_backend;   // defined in trellis_model.cpp (TRELLIS_BACKEND)
 
 // Every knob for one TRELLIS.2 image->3D run. Resolved as default -> environment
 // (the historical TRELLIS_* / GSS / GSH names) -> CLI flag, with the CLI winning.
@@ -24,6 +26,11 @@ struct TrellisParams {
     std::string host   = "127.0.0.1";                           // trellis-server only
     int      port = 8080;                                       // trellis-server only
     int      gpu  = 0;                                          // >=0 GPU index, <0 CPU
+    std::string backend;        // force a ggml backend by name ("Vulkan", "HTP",
+                                // "CPU", ...); empty = auto-select. --gpu N then indexes
+                                // within that backend's devices.
+    int threads = 0;            // CPU backend thread count; 0 = auto (all cores).
+                                // ggml's own default is GGML_DEFAULT_N_THREADS == 4.
     uint32_t seed = 0;
 
     bool cascade    = true;     // 1024 cascade (default); --res 512 selects the light path
@@ -51,15 +58,26 @@ struct TrellisParams {
     int  webp     = -1;         // GLB texture encoding: -1 auto (WebP if built with it), 1 on, 0 off (PNG)
     bool f32      = false;      // f32 sparse-conv compute
     bool no_fa    = false;      // disable FlashAttention (manual softmax)
+    int  fa_fast  = -1;         // FA precision: -1 auto (fast on HTP), 0 BF16/F32, 1 F16/fast
     bool require_gpu = false;   // refuse CPU fallback if no GPU is usable
-    int  threads  = 0;          // CPU backend thread count; 0 = all cores
     float gss = 7.5f;           // sparse-structure guidance strength
     float gsh = 7.5f;           // shape-SLAT guidance strength
     bool voxply = false;        // dump out/myvox.ply              (debug)
     bool dump_slat = false;     // dump /tmp/hr_slat.bin           (debug)
     bool dump_bg = false;       // also write the bg-removal cutout as <out>_cutout.png
     bool bg_only = false;       // background removal only: write the cutout and skip the rest
 
+    int  steps    = 0;          // flow sampler steps; 0 = model default (12). Lowering this
+                                //   trades output quality for turnaround and is meant for
+                                //   backend bring-up: a CPU and an NPU run at the same low
+                                //   step count are still directly comparable to each other.
+    int  sched = -1;            // multi-backend scheduler: -1 auto (on unless the primary
+                                // backend is the CPU), 0 off, 1 on. Required for partial
+                                // op-coverage accelerators like the Hexagon NPU.
+    int  vulkan_fallback = -1;   // add Vulkan between HTP and CPU: -1 environment/default,
+                                // 0 disabled, 1 enabled
+    bool verbose = false;       // --verbose: per-stage timings, graph shapes, and a
+                                // heartbeat while a backend compute is in flight
     bool help = false;          // --help requested
 
     // 512 -> light single-res path; 1024/1536 -> cascade with that HR target.

include/trellis_debug.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#include <chrono>
+#include <string>
+
+namespace trellis {
+
+// --verbose / TRELLIS_VERBOSE=1 enables graph timings and progress heartbeats.
+bool verbose();
+void set_verbose(bool v);
+
+// Scoped wall-clock timer. Logs "<label>: <n> ms" on destruction when verbose().
+class ScopedTimer {
+public:
+    explicit ScopedTimer(std::string label);
+    ~ScopedTimer();
+    double ms() const;
+
+private:
+    std::string                                    label_;
+    std::chrono::steady_clock::time_point          t0_;
+};
+
+// Instrumented graph execution for a single backend or scheduler.
+ggml_status compute_graph(ggml_backend_t backend, ggml_cgraph* graph, const char* label);
+ggml_status compute_graph_sched(ggml_backend_sched_t sched, ggml_cgraph* graph, const char* label);
+
+// Log a graph's shape (node/leaf counts, and the op histogram at verbose level 2).
+void log_graph(const char* label, const ggml_cgraph* graph);
+
+}  // namespace trellis

include/trellis_model.h

@@ -11,6 +11,7 @@ struct ggml_context;
 struct gguf_context;
 struct ggml_backend;
 struct ggml_backend_buffer;
+struct ggml_backend_sched;
 
 namespace trellis {
 
@@ -20,7 +21,15 @@ struct Model {
     gguf_context*         gguf   = nullptr;  // KV metadata + tensor table
     ggml_context*         meta   = nullptr;  // owns the ggml_tensor structs
     ggml_backend*         backend = nullptr; // where weights live
-    ggml_backend_buffer*  buffer = nullptr;  // the weight buffer
+    ggml_backend_buffer*  buffer = nullptr;  // weights in the backend's default buffer
+    // Optional backend-specific buffer for repacked quantized matrix weights.
+    ggml_backend_buffer*  buffer_repack = nullptr;
+    ggml_context*         ctx_host   = nullptr;   // owns the tensors in `buffer`
+    ggml_context*         ctx_repack = nullptr;   // owns the tensors in `buffer_repack`
+    // Multi-backend scheduler over [backend, ...fallbacks]. Non-null only when the
+    // primary backend has partial op coverage (e.g. the Hexagon NPU), in which case every
+    // graph runs through it so unsupported ops land on a backend that can execute them.
+    ggml_backend_sched*   sched  = nullptr;
     std::map<std::string, ggml_tensor*> tensors;
     std::string arch;          // general.architecture
     std::string config_json;   // trellis.config_json (raw model config)

include/trellis_sched.h

@@ -0,0 +1,78 @@
+#pragma once
+
+#include "ggml.h"
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+
+namespace trellis {
+
+struct Model;
+
+// Multi-backend execution.
+//
+// trellis historically ran every graph with ggml_backend_graph_compute() on a single
+// backend, which requires that backend to support *every* op in the graph. That holds for
+// CUDA and Vulkan, whose op coverage is near-complete. It does not hold for the Qualcomm
+// Hexagon NPU: the HTP skel implements a subset of ops and rejects tensors larger than
+// VTCM, and with no scheduler those ops are simply never executed -- the run reports
+// "dsp-rsp: UNKNOWN" and continues with undefined data until the DSP queue desynchronizes.
+//
+// ggml_backend_sched assigns each node to a backend that actually supports it, keeping
+// weight-consuming ops on the backend that owns the weights and inserting copies where a
+// node has to run elsewhere. That is what makes partial-coverage accelerators usable.
+//
+// Mode: -1 auto (on unless the primary backend is the CPU, which needs no fallback),
+//        0 off (legacy single-backend path), 1 on.
+void set_sched_mode(int mode);
+
+// Add Vulkan between an HTP primary and the CPU catch-all. Mode -1 uses the
+// TRELLIS_VULKAN_FALLBACK environment fallback, 0 disables it, and 1 enables it.
+void set_vulkan_fallback_mode(int mode);
+
+bool sched_enabled_for(ggml_backend_t primary);
+
+// Fallback backends shared process-wide. Created on demand, never freed; they outlive the
+// per-stage Models that borrow them.
+ggml_backend_t fallback_cpu_backend();
+
+// Build a scheduler for a freshly loaded model's backend, or nullptr when the single
+// backend path is sufficient. Called by Model::load; the Model owns the result.
+ggml_backend_sched_t make_sched(ggml_backend_t primary);
+
+// Applies --threads / TRELLIS_THREADS to a CPU backend. Defined in trellis_model.cpp;
+// declared here because the scheduler's CPU fallback needs it too -- a fallback stuck on
+// ggml's 4-thread default would quietly become the bottleneck.
+void tune_cpu_threads(ggml_backend_t b);
+
+// Allocation + execution strategy for one graph.
+//
+// Both paths follow the same contract, which mirrors what the call sites already did:
+//   GraphExec ex(m);
+//   ex.alloc(graph);                       // must precede input uploads
+//   ggml_backend_tensor_set(input, ...);   // graph tensors exist only after alloc
+//   ex.compute(graph, "label");
+class GraphExec {
+public:
+    explicit GraphExec(const Model& m);
+    ~GraphExec();
+
+    GraphExec(const GraphExec&)            = delete;
+    GraphExec& operator=(const GraphExec&) = delete;
+
+    // Allocate a graph. Call again only when the graph itself changes.
+    bool alloc(ggml_cgraph* graph);
+
+    ggml_status compute(ggml_cgraph* graph, const char* label);
+
+    // Bytes of compute buffer, for the existing [stage-alloc] diagnostics.
+    size_t buffer_size() const;
+
+    bool uses_sched() const { return sched_ != nullptr; }
+
+private:
+    ggml_backend_t       backend_ = nullptr;
+    ggml_gallocr_t       galloc_  = nullptr;   // single-backend path
+    ggml_backend_sched_t sched_   = nullptr;   // borrowed from the Model
+};
+
+}  // namespace trellis