Save inference parity implementation and evaluation harness

native/metal/ds4_gpu.h

@@ -47,10 +47,36 @@ int ds4_gpu_metal4_tensor_api_enabled(void);
 
 ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes);
 ds4_gpu_tensor *ds4_gpu_tensor_alloc_untracked(uint64_t bytes);
+void *ds4_gpu_mtplx_heap_create(uint64_t bytes);
+void ds4_gpu_mtplx_heap_free(void *heap);
+ds4_gpu_tensor *ds4_gpu_mtplx_tensor_alloc(uint64_t bytes, void *heap);
+ds4_gpu_tensor *ds4_gpu_mtplx_tensor_empty(void);
+int ds4_gpu_tensor_has_heap(const ds4_gpu_tensor *tensor);
+int ds4_gpu_mtplx_residency_supported(void);
+void *ds4_gpu_mtplx_residency_create(void);
+void ds4_gpu_mtplx_residency_free(void *set);
+void ds4_gpu_mtplx_residency_allocation(void *set, void *allocation, int add);
+void ds4_gpu_mtplx_residency_commit(void *set);
+void ds4_gpu_mtplx_residency_attach(void *queue, void *const *sets, uint64_t count);
+void *ds4_gpu_mtplx_command_queue(void);
+void *ds4_gpu_mtplx_retain_command_queue(void);
+void ds4_gpu_mtplx_queue_free(void *queue);
+int ds4_gpu_mtplx_submission_hook(void *context, void (*before_commit)(void *, void *));
+void *ds4_gpu_mtplx_tensor_allocation(const ds4_gpu_tensor *tensor);
+uint64_t ds4_gpu_mtplx_allocated_size(void *allocation);
+uint64_t ds4_gpu_mtplx_residency_count(void *set);
+int ds4_gpu_mtplx_residency_contains(void *set, void *allocation);
 ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes);
 ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes);
 void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor);
+/* Transfer context to a completion callback without committing or waiting.
+ * Releases immediately when there is no open or pending work. An empty open
+ * batch returns 0 without taking ownership; callers must retain the context. */
+int ds4_gpu_defer_until_completed(void *context, void (*release)(void *));
+uint64_t ds4_gpu_max_buffer_length(void);
 uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor);
+/* Borrowed CPU pointer, valid while the caller owns the tensor. This does not
+ * retain storage beyond that ownership or synchronize pending GPU work. */
 void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor);
 int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count);
 int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes);
@@ -80,6 +106,18 @@ int ds4_gpu_begin_commands(void);
 int ds4_gpu_flush_encoder(void);
 int ds4_gpu_flush_commands(void);
 int ds4_gpu_commands_active(void);
+typedef struct {
+    double busy_seconds;
+    double max_busy_seconds;
+    uint64_t command_buffers;
+} ds4_gpu_busy_stats;
+void ds4_gpu_busy_stats_reset(void);
+void ds4_gpu_busy_stats_get(ds4_gpu_busy_stats *stats);
+typedef struct {
+    double scheduled_seconds;
+    double completed_seconds;
+} ds4_gpu_canary_sample;
+int ds4_gpu_canary_probe(ds4_gpu_canary_sample *sample);
 #ifdef __APPLE__
 int ds4_gpu_parallel_ffn_finish(void);
 void ds4_gpu_parallel_ffn_abort(void);
@@ -156,6 +194,66 @@ int ds4_gpu_qwen_dispatch(
         const ds4_gpu_qwen_kernel_args *args,
         uint32_t grid_x,
         uint32_t grid_y);
+/* Ordered contiguous buffers: reference inputs first, then outputs. No dtype
+ * conversion, argument remapping, extra flush, or inferred thread geometry. */
+int ds4_gpu_qwen_dispatch_buffers(
+        const char *kernel,
+        const ds4_gpu_tensor *const *buffers,
+        uint32_t count,
+        const uint32_t grid[3],
+        const uint32_t threadgroup[3]);
+/* Native Metal bindings only; Rust selects kernels, parameters and geometry. */
+typedef struct {
+    const ds4_gpu_tensor *tensor;
+    const void *bytes;
+    uint32_t size;
+    uint32_t index;
+    uint64_t offset;
+} ds4_gpu_metal_binding;
+typedef struct {
+    uint32_t index;
+    uint32_t value;
+    uint32_t kind; /* 0: Bool, 1: I32 bit pattern. */
+} ds4_gpu_metal_constant;
+/* Optional Rust allocation owner: clear *resource_owner only when transferred
+ * to the command completion handler. On earlier failure the caller retains it. */
+int ds4_gpu_mtplx_dispatch(
+        const char *kernel,
+        const ds4_gpu_metal_binding *bindings, uint32_t count,
+        const ds4_gpu_metal_constant *constants, uint32_t constant_count,
+        const uint32_t groups[3], const uint32_t threads[3], uint32_t individual_threads,
+        void **resource_owner, void (*release_owner)(void *));
+/* Same dispatch, on a Rust-owned command buffer/concurrent encoder. Null
+ * handles select the existing dispatch path. No implicit end/commit/wait. */
+int ds4_gpu_mtplx_encode(
+        const char *kernel,
+        const ds4_gpu_metal_binding *bindings, uint32_t count,
+        const ds4_gpu_metal_constant *constants, uint32_t constant_count,
+        const uint32_t groups[3], const uint32_t threads[3], uint32_t individual_threads,
+        void **resource_owner, void (*release_owner)(void *),
+        void *command_buffer, void *encoder, uint32_t barrier);
+void *ds4_gpu_mtplx_queue_create(void);
+void *ds4_gpu_mtplx_command_create(void *queue);
+void *ds4_gpu_mtplx_encoder_create(void *command);
+void *ds4_gpu_mtplx_fence_create(void);
+void ds4_gpu_mtplx_object_free(void *object);
+void ds4_gpu_mtplx_encoder_barrier(void *encoder);
+void ds4_gpu_mtplx_encoder_fence(void *encoder, void *fence, int update);
+void ds4_gpu_mtplx_encoder_end(void *encoder);
+void ds4_gpu_mtplx_command_completed(void *command, void *context,
+        void (*completed)(void *, const char *));
+void ds4_gpu_mtplx_command_commit(void *command);
+void ds4_gpu_mtplx_command_wait(void *command);
+uint32_t ds4_gpu_mtplx_arch_suffix(void);
+void ds4_gpu_mtplx_nax_platform(uint32_t out[4]);
+void *ds4_gpu_mtplx_event_create(void);
+void ds4_gpu_mtplx_command_event(void *command, void *event, uint64_t value, int signal);
+void ds4_gpu_mtplx_event_signal(void *event, uint64_t value);
+void ds4_gpu_mtplx_event_wait(void *event, uint64_t value);
+uint64_t ds4_gpu_mtplx_event_value(void *event);
+/* With source == NULL, only query the existing pipeline cache. */
+int ds4_gpu_mtplx_compile(const char *kernel, const char *source);
+uint32_t ds4_gpu_mtplx_kernel_max_threads(const char *kernel);
 int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label);
 int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label);
 int ds4_gpu_q8_cache_suppressed(void);
@@ -229,6 +327,7 @@ static inline int ds4_gpu_device_is_m5_apple_silicon(void) { return 0; }
 void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts);
 void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes);
 uint64_t ds4_gpu_recommended_working_set_size(void);
+uint64_t ds4_gpu_current_allocated_size(void);
 uint32_t ds4_gpu_stream_expert_cache_configured_count(void);
 uint32_t ds4_gpu_stream_expert_cache_current_count(void);
 typedef struct ds4_gpu_stream_expert_cache_stats {
@@ -1858,6 +1957,19 @@ int ds4_gpu_glm_attention_flash_tensor(
         uint32_t              value_dim,
         bool                  cache_f16);
 
+/* Qwen prefill attention over already staged head-major F16 K/V tensors. */
+int ds4_gpu_qwen_attention_flash_staged_tensor(
+        ds4_gpu_tensor       *heads,
+        const ds4_gpu_tensor *q,
+        const ds4_gpu_tensor *key_cache,
+        const ds4_gpu_tensor *value_cache,
+        uint32_t              pos0,
+        uint32_t              n_tokens,
+        uint32_t              cache_len,
+        uint32_t              n_head,
+        uint32_t              n_kv_head,
+        uint32_t              head_dim);
+
 /* Release decode fused KV finalizer: after the standalone RoPE kernel, this
  * performs DS4's FP8 non-RoPE KV round trip and writes the F16-rounded raw
  * attention cache row in one dispatch. */