| /* ds4_gpu_mgpu.h — multi-GPU plumbing types and APIs (v0). | |
| * | |
| * This header carries the new multi-GPU additions for the multi-GPU plumbing PP work | |
| * (device-aware CUDA). It is included from ds4_cuda.cu and from | |
| * downstream tasks that need access to g_gpu[], g_n_gpus, g_gpu_peer_ok[], | |
| * the ds4_gpu_config struct, and the new tensor APIs. | |
| * | |
| * Why not in ds4_gpu.h? The legacy ds4_gpu.h is included from C-only | |
| * callers (ds4.c, ds4_cli.c, etc.) and from the Metal build, but is NOT | |
| * included from ds4_cuda.cu historically. That asymmetry hid pre-existing | |
| * signature mismatches between the legacy header and ds4_cuda.cu. We keep | |
| * the legacy header opaque and put the new shared types here, so this | |
| * file is the single source of truth for both ds4_cuda.cu and downstream | |
| * multi-GPU tasks without disturbing the legacy contract. | |
| * | |
| * The struct definitions reference CUDA-specific handle types via void * | |
| * placeholders so the header is safe to include from C builds, Metal | |
| * builds, and the CUDA build (where ds4_cuda.cu casts the void * back | |
| * to cudaStream_t / cublasHandle_t / cudaEvent_t internally). | |
| */ | |
| extern "C" { | |
| /* Complete definition of the previously-opaque ds4_gpu_tensor, plus a | |
| * typedef so the new API prototypes below can use the bare name | |
| * `ds4_gpu_tensor *` in both C and C++ without forcing callers to | |
| * include ds4_gpu.h first. Callers that include this header can | |
| * stack-allocate or struct-embed tensors and pass to | |
| * ds4_gpu_tensor_alloc_on. */ | |
| struct ds4_gpu_tensor { | |
| void *ptr; | |
| uint64_t bytes; | |
| int owner; | |
| int device_id; /* -1 means legacy/untagged → treat as device 0 */ | |
| }; | |
| typedef struct ds4_gpu_tensor ds4_gpu_tensor; | |
| typedef struct { | |
| uint64_t raw_kv; | |
| uint64_t comp_kv; | |
| uint64_t topk; | |
| uint32_t pos; | |
| uint32_t n_raw; | |
| uint32_t raw_cap; | |
| uint32_t raw_start; | |
| uint32_t n_comp; | |
| uint32_t top_k; | |
| uint32_t window; | |
| uint32_t ratio; | |
| uint32_t indexed; | |
| } ds4_gpu_attention_decode_row; | |
| /* Tagged so headers (notably ds4.h) can forward-declare `struct | |
| * ds4_gpu_config` without dragging in this entire header. */ | |
| typedef struct ds4_gpu_config { | |
| int device_indices[DS4_MAX_GPUS]; /* CUDA device IDs to use */ | |
| /* Explicit per-device budget in bytes. The engine does NOT auto-fill | |
| * missing budgets - a value of 0 means "zero bytes of budget for | |
| * this slot" and (combined with reserves) will push placement to | |
| * CPU spill. Auto-detection (e.g. mapping --gpu-vram auto to | |
| * cudaMemGetInfo) is the caller's job; see CLI flag wiring for the | |
| * canonical CLI path. The engine emits a clear stderr and refuses | |
| * if n_gpus > 0 and every vram_bytes[] is 0 (almost certainly a | |
| * caller bug from zero-initializing the struct). */ | |
| size_t vram_bytes[DS4_MAX_GPUS]; | |
| int n_gpus; | |
| size_t safety_margin_bytes; /* per-device reserve */ | |
| } ds4_gpu_config; | |
| typedef struct { | |
| int device_id; | |
| void *stream; /* cudaStream_t under CUDA */ | |
| void *cublas; /* cublasHandle_t under CUDA */ | |
| int cublas_ready; | |
| void *scratch; | |
| size_t scratch_bytes; | |
| size_t budget_bytes; | |
| size_t used_bytes; | |
| void *boundary_event; /* cudaEvent_t under CUDA */ | |
| } ds4_gpu_ctx; | |
| extern ds4_gpu_ctx g_gpu[DS4_MAX_GPUS]; | |
| extern int g_n_gpus; | |
| extern int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS]; | |
| /* Primary multi-device init. The existing ds4_gpu_init (declared in | |
| * ds4_gpu.h) is a thin shim that builds a single-device config for | |
| * device 0 and calls this. */ | |
| int ds4_gpu_init_multi(const ds4_gpu_config *cfg); | |
| /* Caller-supplied struct alloc on a specific device. Returns 0 on | |
| * success, nonzero on error. Pair with ds4_gpu_tensor_free_in_place. */ | |
| int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id, uint64_t bytes); | |
| void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t); | |
| /* Heap-allocated tensor on a specific logical tier; mirrors the legacy | |
| * ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *) but with a tier | |
| * parameter. Returns NULL on failure. Used by the multi-tier graph | |
| * allocations in ds4.c. Single-tier callers can continue using the | |
| * legacy ds4_gpu_tensor_alloc(bytes) which is equivalent to | |
| * ds4_gpu_tensor_alloc_ptr_on(0, bytes). */ | |
| ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes); | |
| int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst, const ds4_gpu_tensor *src, uint64_t bytes); | |
| void ds4_gpu_enable_q8_dequant_gemm(void); | |
| /* Heap-allocated MANAGED-memory tensor on a specific logical tier. Used | |
| * for the KV-cache pool. Stamps tier on the tensor so subsequent | |
| * ds4_gpu_tensor_free runs under the correct device. Note managed memory | |
| * pages on first-touch and is not strictly device-bound; the stamped | |
| * tier records the home tier for accounting + free. */ | |
| ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes); | |
| /* Cross-device tensor copy. Same-device → cudaMemcpyAsync; peer-capable | |
| * cross-device → cudaMemcpyPeerAsync with event sync; non-peer → pinned | |
| * host bounce (per src→dst pair). Honors DS4_FORCE_HOST_BOUNCE=1. */ | |
| int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst, | |
| const ds4_gpu_tensor *src, | |
| uint64_t bytes); | |
| int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst, | |
| const ds4_gpu_tensor *src, | |
| uint64_t bytes); | |
| /* Grouped default-stream handoff whose copies execute on the destination | |
| * device. The source default stream records readiness; the destination waits, | |
| * performs all three copies, then naturally orders its following kernels. | |
| * This allows source-device work submitted afterward to overlap the handoff. */ | |
| int ds4_gpu_tensor_copy_xdev3_default_dst( | |
| ds4_gpu_tensor *dst0, | |
| const ds4_gpu_tensor *src0, | |
| uint64_t bytes0, | |
| ds4_gpu_tensor *dst1, | |
| const ds4_gpu_tensor *src1, | |
| uint64_t bytes1, | |
| ds4_gpu_tensor *dst2, | |
| const ds4_gpu_tensor *src2, | |
| uint64_t bytes2); | |
| /* Cross-device copy of three buffers between the same source/destination | |
| * tiers, recording a single readiness event for the destination stream. This | |
| * is meant for tiny grouped activation handoffs where three independent | |
| * ds4_gpu_tensor_copy_xdev calls would spend more time in event plumbing than | |
| * in the copies themselves. Falls back internally when a grouped peer copy is | |
| * not applicable. */ | |
| int ds4_gpu_tensor_copy_xdev3(ds4_gpu_tensor *dst0, | |
| const ds4_gpu_tensor *src0, | |
| uint64_t bytes0, | |
| ds4_gpu_tensor *dst1, | |
| const ds4_gpu_tensor *src1, | |
| uint64_t bytes1, | |
| ds4_gpu_tensor *dst2, | |
| const ds4_gpu_tensor *src2, | |
| uint64_t bytes2); | |
| /* Cross-device copy that also orders against prior work on the destination | |
| * stream before overwriting dst. This is needed for pipelined prefill slots | |
| * where the next producer reuses a destination buffer that the destination | |
| * stage just consumed. */ | |
| int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst, | |
| const ds4_gpu_tensor *src, | |
| uint64_t bytes); | |
| /* Order a peer read without copying: records readiness on src's stream and | |
| * makes dst_tier's stream wait for it. */ | |
| int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier); | |
| int ds4_gpu_tensor_wait_xdev_default(const ds4_gpu_tensor *src, int dst_tier); | |
| /* Cross-device float add used by CUDA tensor-parallel reductions. | |
| * See ds4_gpu.h for the full contract. */ | |
| int ds4_gpu_add_xdev_tensor(ds4_gpu_tensor *out, | |
| const ds4_gpu_tensor *local, | |
| const ds4_gpu_tensor *remote, | |
| ds4_gpu_tensor *remote_tmp, | |
| uint32_t n); | |
| /* Returns the device_id recorded on the tensor; -1 if untagged. */ | |
| int ds4_gpu_tensor_device(const ds4_gpu_tensor *t); | |
| /* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1). | |
| * | |
| * This is the canonical shim for per-layer device routing in the | |
| * multi-tier execution path. The caller passes a logical tier index; | |
| * the shim internally indexes g_gpu[tier].device_id and calls | |
| * cudaSetDevice. Returns 0 on success, nonzero on error or if the | |
| * tier index is out of range. | |
| * | |
| * Wave-2 multi-GPU placement scaffolding adds this shim but does not | |
| * exercise the multi-tier execution path; multi-GPU execution | |
| * (follow-up) is its first caller. */ | |
| int ds4_gpu_set_current_device(int logical_tier); | |
| int ds4_gpu_set_current_device_fenced(int logical_tier); | |
| /* Register the mmap'd host model pointer for selective-cache lookups | |
| * WITHOUT triggering any device-side copy. This bypasses the | |
| * DS4_CUDA_COPY_MODEL environment-variable branch that | |
| * ds4_gpu_set_model_map normally honors, which is essential for | |
| * multi-tier startup: we want only per-device selective tensor caches, | |
| * never the whole-model copy. Returns 1 on success, 0 on error. */ | |
| int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size); | |
| /* Strict per-device selective-cache lookup (no fallback). | |
| * | |
| * Returns 1 only if a covering entry exists whose device_id matches the | |
| * caller-supplied expected_device (a PHYSICAL CUDA device id). Otherwise | |
| * returns 0 — no host-pointer fallback, no different-device match. The | |
| * caller is expected to have cudaSetDevice'd to expected_device before | |
| * invoking; the returned pointer is valid to consume from that device's | |
| * kernel. | |
| * | |
| * Used by multi-tier kernel-dispatch resolvers in | |
| * multi-GPU execution (multi-GPU execution). Single-tier callers should | |
| * keep using ds4_gpu_lookup_cache for back-compat behavior. */ | |
| int ds4_gpu_lookup_cache_strict(uint64_t source_offset, | |
| uint64_t bytes, | |
| int expected_device, | |
| void **out_device_ptr); | |
| } /* extern "C" */ | |