ggml-hrx: graph inputs are never written back to host; non-replay programs wait for their input upload

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

@@ -23,6 +23,10 @@ struct CommandProgramBinding {
     void *       host_data  = nullptr;
     bool         weight     = false;
     bool         empty_value = false;
+    // ggml graph input (GGML_TENSOR_FLAG_INPUT). A kernel may rewrite its device copy in place (the Qwen
+    // attention metadata kernel regenerates positions, cache indices and the mask), but that copy is never
+    // written back to the caller's host tensor: ggml reuses that memory for the next graph's inputs.
+    bool         graph_input = false;
 
     bool requires_materialization() const { return host_data != nullptr; }
 };

ggml/src/ggml-hrx/runtime/command-program-executor.cpp

@@ -643,7 +643,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec
         staging.value                        = binding.value.value;
         staging.host_data                    = static_cast<uint8_t *>(binding.host_data) + binding.offset;
         staging.upload                       = access.read;
-        staging.download                     = access.write;
+        staging.download                     = access.write && !binding.graph_input;
         if (staging.upload && context.host_buffers != nullptr) {
             staging.source_host_buffer = context.host_buffers->find(staging.host_data, staging.length);
         }
@@ -1966,6 +1966,24 @@ bool execute_prepared_command_program(const CommandProgramExecutionContext & con
         return false;
     }
     HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-download-post");
+    // Uploads from registered host buffers are device-timeline copies that read
+    // the host memory when they execute. ggml writes the next graph's inputs
+    // into that same memory as soon as graph_compute returns, so wait for this
+    // program here, as the graph replay path does after its launch.
+    bool reads_registered_host_memory = false;
+    for (const HostStagingBuffer & staging : commands.host_staging) {
+        reads_registered_host_memory =
+            reads_registered_host_memory || (staging.upload && staging.source_host_buffer.valid());
+    }
+    if (reads_registered_host_memory) {
+        static thread_local WaitHistory registered_upload_wait_history;
+        if (ErrorResult error =
+                take_status(stream_synchronize_sleeping(context.stream, registered_upload_wait_history))) {
+            GGML_LOG_ERROR("%s: wait for HRX command program reading host memory: %s\n", __func__, error->c_str());
+            return false;
+        }
+        mark_graph_replay_synchronized(context);
+    }
     return true;
 }
 

ggml/src/ggml-hrx/runtime/graph-executor.cpp

@@ -64,6 +64,9 @@ CommandProgramBindings GraphExecutor::bind_external_value_buffers(const GraphPro
             binding.capacity   = value_binding.capacity;
             binding.weight     = value_binding.weight;
             binding.empty_value = ggml_nbytes(external.tensor) == 0;
+            binding.graph_input = (external.tensor->flags & GGML_TENSOR_FLAG_INPUT) != 0 ||
+                                  (external.tensor->view_src != nullptr &&
+                                   (external.tensor->view_src->flags & GGML_TENSOR_FLAG_INPUT) != 0);
         } else {
             status.log("external value %d is not bound", external.value.value);
         }