ggml-hrx: materialize weights on demand and keep them resident

ggml/include/ggml-hrx.h

@@ -11,6 +11,13 @@ struct ggml_backend_hrx_cache_stats {
     uint64_t graph_program_hits;
     uint64_t prepared_program_builds;
     uint64_t prepared_program_hits;
+    uint64_t weight_materializations;
+    uint64_t weight_reuses;
+    uint64_t weight_uploaded_bytes;
+    // Current weight-cache residency.
+    uint64_t weight_live_allocations;
+    uint64_t weight_live_bytes;
+    uint64_t weight_duplicate_warnings;
 };
 
 GGML_BACKEND_API ggml_backend_t             ggml_backend_hrx_init(size_t device);

ggml/src/ggml-hrx/README.md

@@ -1,5 +1,11 @@
 # HRX Backend
 
+## Weight residency diagnostics
+
+HRX materializes imported host weights in the layout requested by the selected kernel and caches each result for the backend's lifetime. If one weight is materialized in two layouts, HRX warns about the duplicate allocations and their memory cost.
+
+To log residency totals, set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v`. HRX prints a `ggml_hrx: weight_stats` line when the backend is released, with materialization and reuse counts, cached allocations and bytes, and upload totals.
+
 ## Compilation diagnostics
 
 Set `GGML_HRX_TRACE_COMPILE=1` before starting an application using the ggml HRX backend to log compilation activity. Each record is a JSON object on stderr prefixed with `hrx compile: `. This output bypasses the application's log-level filter, so it also works with quiet benchmark runs. The setting is read once per process; unset or any value other than `1` disables it.

ggml/src/ggml-hrx/backend-buffer-binding.cpp

@@ -27,7 +27,13 @@ bool ggml_backend_hrx_tensor_binding(const ggml_tensor *                tensor,
     }
     auto *       context = ggml_backend_hrx_buffer_context_from_buffer(buffer);
     const size_t offset  = ggml_backend_hrx_tensor_offset(context, tensor);
-    if (context->buffer == nullptr || offset > buffer->size || ggml_nbytes(tensor) > buffer->size - offset) {
+    const bool has_device_buffer = context->buffer != nullptr;
+    const bool is_imported = context->imported;
+    const bool has_storage = has_device_buffer || is_imported;
+    const bool offset_fits = offset <= buffer->size;
+    const bool tensor_size_fits = offset_fits ? ggml_nbytes(tensor) <= buffer->size - offset : false;
+    const bool valid_binding = has_storage && offset_fits && tensor_size_fits;
+    if (!valid_binding) {
         return false;
     }
     *out_context = context;
@@ -68,19 +74,86 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx
         binding.generation = 1;
         binding.capacity   = capacity;
         binding.weight     = ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+        // Standard GGUF weights are immutable, so their native copy can use the shared weight cache.
+        // Layout eligibility remains a separate import-only decision.
+        binding.source_host_data = base;
         return true;
     }
     ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
-    const bool directly_bindable = !ggml_backend_buffer_is_host(buffer) || context->direct_host_binding;
+    const bool is_imported = context->imported;
+    const bool is_host_buffer = ggml_backend_buffer_is_host(buffer);
+    const bool ordinary_device_buffer = !is_imported && !is_host_buffer;
+    const bool directly_bindable = ordinary_device_buffer || context->direct_host_binding;
     // Coherent HRX host allocations are directly device-addressable. Represent them with an HRX buffer handle so
     // command-program preparation bypasses host materialization. Noncoherent host allocations remain host data.
     binding.buffer               = directly_bindable ? context->buffer : nullptr;
-    binding.host_data            = directly_bindable ? nullptr : context->base;
-    binding.offset               = offset;
+    const bool shared_host_region = context->retained_base != nullptr;
+    binding.host_data            = directly_bindable ? nullptr :
+                                   (shared_host_region ? context->retained_base : context->base);
+    binding.offset               = shared_host_region ? context->physical_offset + offset : offset;
     binding.length               = ggml_nbytes(tensor);
     binding.identity             = context->identity;
     binding.generation           = context->generation;
-    binding.capacity             = buffer != nullptr ? buffer->size : 0;
+    binding.capacity             = shared_host_region ? context->physical_capacity : buffer->size;
     binding.weight = buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+    const bool retained_host_source = is_imported || is_host_buffer;
+    binding.source_host_data = retained_host_source ?
+        (shared_host_region ? context->retained_base : context->base) : nullptr;
     return true;
 }
+
+bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_retained_region & region) {
+    ggml_backend_hrx_buffer_context * context = nullptr;
+    size_t                            offset  = 0;
+    const bool hrx_bound = ggml_backend_hrx_tensor_binding(tensor, &context, &offset);
+    if (hrx_bound) {
+        ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+        const bool retained = context->imported || ggml_backend_buffer_is_host(buffer);
+        if (!retained) {
+            return false;
+        }
+        region.identity   = context->identity;
+        region.offset     = context->retained_base != nullptr ? context->physical_offset + offset : offset;
+    } else {
+        ggml::hrx::ValueBufferBinding binding;
+        const bool resolved = ggml_backend_hrx_resolve_value_buffer(tensor, binding);
+        const bool has_host_source = resolved && binding.host_data != nullptr;
+        if (!has_host_source) {
+            return false;
+        }
+        region.identity = binding.identity;
+        region.offset = binding.offset;
+    }
+    region.length     = ggml_nbytes(tensor);
+    return true;
+}
+
+bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor) {
+    if (tensor == nullptr) {
+        return false;
+    }
+    const ggml_tensor * root = tensor->view_src != nullptr ? tensor->view_src : tensor;
+    ggml_backend_buffer_t buffer = root->buffer;
+    if (buffer == nullptr) {
+        return false;
+    }
+    if (buffer->iface.get_base == ggml_backend_hrx_buffer_base) {
+        const auto * context = ggml_backend_hrx_buffer_context_from_buffer(buffer);
+        const bool borrowed_import = context->imported;
+        const bool known_host_alias = context->retained_base != nullptr;
+        const bool repackable_source = borrowed_import && !known_host_alias;
+        return repackable_source;
+    }
+    return false;
+}
+
+bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor) {
+    if (tensor == nullptr) {
+        return true;
+    }
+    const ggml_tensor * root = tensor->view_src != nullptr ? tensor->view_src : tensor;
+    ggml_backend_buffer_t buffer = root->buffer;
+    const bool is_weight = buffer != nullptr &&
+                           ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+    return !is_weight || ggml_backend_hrx_weight_has_repackable_source(tensor);
+}

ggml/src/ggml-hrx/backend-buffer-binding.h

@@ -8,6 +8,14 @@
 
 struct ggml_tensor;
 
+// The buffer identity and byte range name the source independently of a resident HRX allocation.
+// GGML's tensor metadata supplies the logical shape, including views into a shared import.
+struct ggml_backend_hrx_retained_region {
+    uint64_t                          identity;
+    size_t                            offset;
+    size_t                            length;
+};
+
 ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer);
 size_t ggml_backend_hrx_tensor_offset(const ggml_backend_hrx_buffer_context * context, const ggml_tensor * tensor);
 void * ggml_backend_hrx_buffer_base(ggml_backend_buffer_t buffer);
@@ -16,3 +24,7 @@ bool ggml_backend_hrx_tensor_binding(const ggml_tensor *                tensor,
                                      ggml_backend_hrx_buffer_context ** out_context,
                                      size_t *                           out_offset);
 bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx::ValueBufferBinding & binding);
+// Only borrowed imports outside HRX-owned host storage may use alternative weight layouts.
+bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor);
+bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor);
+bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_retained_region & region);

ggml/src/ggml-hrx/backend-context.h

@@ -34,6 +34,10 @@ struct ggml_backend_hrx_buffer_context {
     uint64_t                          identity;
     uint64_t                          generation;
     bool                              direct_host_binding;
+    bool                              imported;
+    size_t                            physical_offset = 0;
+    size_t                            physical_capacity = 0;
+    uint8_t *                         retained_base = nullptr;
 };
 
 struct ggml_backend_hrx_device_context {

ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp

@@ -1,6 +1,5 @@
 #include "command-program-bindings.h"
 
-#include <sstream>
 #include <utility>
 
 namespace ggml::hrx {
@@ -27,9 +26,7 @@ CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & v
             result.status.log("external value %d is not bound", id.value);
             continue;
         }
-        bindings.push_back({ value->id, buffer->buffer, buffer->offset, buffer->length, buffer->identity,
-                             buffer->generation, buffer->capacity, buffer->host_data, buffer->weight,
-                             value->byte_count == 0 });
+        bindings.push_back(make_command_program_binding(value->id, *buffer, value->byte_count == 0));
     }
     return from_bindings(std::move(bindings), result.status);
 }
@@ -50,7 +47,7 @@ CommandProgramBindings CommandProgramBindings::from_bindings(std::vector<Command
     return result;
 }
 
-const CommandProgramBinding * CommandProgramBindings::find(ValueId value) const {
+const CommandProgramBinding * CommandProgramBindings::find_by_value(ValueId value) const {
     for (const CommandProgramBinding & binding : bindings_) {
         if (binding.value == value) {
             return &binding;
@@ -59,34 +56,42 @@ const CommandProgramBinding * CommandProgramBindings::find(ValueId value) const
     return nullptr;
 }
 
+const CommandProgramBinding * CommandProgramBindings::find(const CommandBinding & request) const {
+    for (const CommandProgramBinding & binding : bindings_) {
+        const bool same_value = binding.value == request.value;
+        const bool same_layout = binding.layout == request.layout;
+        const bool native_layout = request.layout == kNativeWeightLayout;
+        const bool same_transformation = native_layout ||
+            (binding.layout_source_type == request.source_type && binding.input_size == request.input_size &&
+             binding.output_size == request.output_size && binding.source_length == request.source_length &&
+             binding.length == request.length);
+        const bool matches_request = same_value && same_layout && same_transformation;
+        if (matches_request) {
+            return &binding;
+        }
+    }
+    return nullptr;
+}
+
 CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBindings & bindings) {
     uint64_t hash = UINT64_C(1469598103934665603);
     mix_hash(hash, UINT64_C(0x6872782d62696e64));
     for (const CommandProgramBinding & binding : bindings.bindings()) {
         mix_hash(hash, static_cast<uint64_t>(static_cast<int64_t>(binding.value.value)));
         mix_hash(hash, binding.host_data != nullptr ? 1 : 0);
         mix_hash(hash, binding.identity);
-        mix_hash(hash, binding.generation);
+        const bool retained_weight = binding.weight && binding.source_host_data != nullptr;
+        mix_hash(hash, retained_weight ? binding.generation : 0);
         mix_hash(hash, static_cast<uint64_t>(binding.capacity));
         mix_hash(hash, static_cast<uint64_t>(binding.offset));
         mix_hash(hash, static_cast<uint64_t>(binding.length));
         mix_hash(hash, binding.weight ? 1 : 0);
         mix_hash(hash, binding.empty_value ? 1 : 0);
+        mix_hash(hash, binding.source_host_data != nullptr ? 1 : 0);
     }
     mix_hash(hash, static_cast<uint64_t>(bindings.bindings().size()));
     return { hash };
 }
 
-CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings) {
-    std::ostringstream out;
-    out << "hrx-bindings-v1";
-    for (const CommandProgramBinding & binding : bindings.bindings()) {
-        out << "|value=" << binding.value.value << "|kind=" << (binding.host_data != nullptr ? "host" : "device")
-            << "|identity=" << binding.identity << "|generation=" << binding.generation
-            << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length
-            << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0);
-    }
-    return { out.str() };
-}
 
 }  // namespace ggml::hrx

ggml/src/ggml-hrx/dispatch/command-program-bindings.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "graph/value-map.h"
+#include "command-program.h"
 #include "status.h"
 
 #include <cstddef>
@@ -23,13 +24,28 @@ struct CommandProgramBinding {
     void *       host_data  = nullptr;
     bool         weight     = false;
     bool         empty_value = false;
+    // Retained host bytes are separate from the address used to bind a command.
+    const void * source_host_data = nullptr;
+    // Materialized bindings are selected by the requesting command, not just by ValueId.
+    std::string  layout = kNativeWeightLayout;
+    ggml_type    layout_source_type = GGML_TYPE_COUNT;
+    int64_t      input_size = 0;
+    int64_t      output_size = 0;
+    size_t       source_length = 0;
 
     bool requires_materialization() const { return host_data != nullptr; }
 };
 
-struct CommandProgramBindingsFingerprint {
-    std::string value;
-};
+inline CommandProgramBinding make_command_program_binding(
+    ValueId value, const ValueBufferBinding & buffer, bool empty_value) {
+    CommandProgramBinding binding = { value, buffer.buffer, buffer.offset, buffer.length, buffer.identity,
+                                      buffer.generation, buffer.capacity, buffer.host_data, buffer.weight,
+                                      empty_value };
+    binding.source_host_data = buffer.source_host_data;
+    return binding;
+}
+
+
 
 struct CommandProgramBindingsHash {
     uint64_t value = 0;
@@ -41,7 +57,8 @@ class CommandProgramBindings {
     static CommandProgramBindings from_bindings(std::vector<CommandProgramBinding> bindings,
                                                 const Status &                     errors = {});
 
-    const CommandProgramBinding * find(ValueId value) const;
+    const CommandProgramBinding * find_by_value(ValueId value) const;
+    const CommandProgramBinding * find(const CommandBinding & request) const;
 
     const std::vector<CommandProgramBinding> & bindings() const { return bindings_; }
 
@@ -54,6 +71,5 @@ class CommandProgramBindings {
 };
 
 CommandProgramBindingsHash        command_program_bindings_hash(const CommandProgramBindings & bindings);
-CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings);
 
 }  // namespace ggml::hrx

ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp

@@ -23,7 +23,7 @@ static Status resolve_command_binding(const Command &                     comman
     switch (binding.origin) {
         case CommandBindingOrigin::GraphValue:
             {
-                const CommandProgramBinding * concrete = bindings.find(binding.value);
+                const CommandProgramBinding * concrete = bindings.find(binding);
                 if (concrete == nullptr) {
                     status.log("%s %s is not bound", command_context.c_str(), binding_context.c_str());
                     return status;

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp

@@ -723,10 +723,16 @@ static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLU
         return false;
     }
 
-    const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 &&
-                         match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 &&
-                         match.gate_format == CommonMulMatWeightFormat::Q4K &&
-                         match.up_format == CommonMulMatWeightFormat::Q4K;
+    const bool has_few_tokens = match.token_count <= 5;
+    const bool has_aligned_output = match.output_size % 64 == 0;
+    const bool has_direct_gate = match.gate_weight->alias_source.value < 0;
+    const bool has_direct_up = match.up_weight->alias_source.value < 0;
+    const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K;
+    const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K;
+    const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight);
+    const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight);
+    const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up &&
+                         has_q4_gate && has_q4_up && can_transform_gate && can_transform_up;
     const bool use_q8 = pack_q4;
     const bool use_direct_dot = match.token_count <= 5 && match.input_size % 256 == 0;
     return !use_q8 && !use_direct_dot && match.token_count >= 2;
@@ -958,10 +964,16 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context,
         }
     }
 
-    const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 &&
-                         match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 &&
-                         match.gate_format == CommonMulMatWeightFormat::Q4K &&
-                         match.up_format == CommonMulMatWeightFormat::Q4K;
+    const bool has_few_tokens = match.token_count <= 5;
+    const bool has_aligned_output = match.output_size % 64 == 0;
+    const bool has_direct_gate = match.gate_weight->alias_source.value < 0;
+    const bool has_direct_up = match.up_weight->alias_source.value < 0;
+    const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K;
+    const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K;
+    const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight);
+    const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight);
+    const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up &&
+                         has_q4_gate && has_q4_up && can_transform_gate && can_transform_up;
     const bool use_q8 = pack_q4;
     const CommonActivationPublicationDemand q8_demand =
         use_q8 ? select_swiglu_q8_publication(context, match) : CommonActivationPublicationDemand{};