Skip to content

rocm: fix distributed inference on unified-memory APUs (strix halo / gfx1151) #407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kyuz0 wants to merge 5 commits into
base: main
Choose a base branch
from
+98 −39
Open

rocm: fix distributed inference on unified-memory APUs (strix halo / gfx1151) #407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
96abadf
rocm: fix split-model memory allocation for distributed inference
Jun 10, 2026
9fb79b4
rocm: support non-contiguous span groups for distributed inference
Jun 10, 2026
00e64ea
rocm: add comments for non-obvious distributed inference decisions
Jun 13, 2026
ef7eb54
Use safer ROCm prefill default
antirez Jun 16, 2026
89c1ea9
Fix ROCm MTP startup residency
antirez Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 1 addition & 6 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8216,11 +8216,7 @@ static uint32_t ds4_prefill_cap_for_prompt(int prompt_len,
cap = (uint32_t)v;
}
} else if (prompt_len > 4096) {
#ifdef DS4_ROCM_BUILD
cap = 8192u;
#else
cap = DS4_MODEL_VARIANT == DS4_VARIANT_PRO ? 8192u : 4096u;
#endif
}
}

Expand Down Expand Up @@ -20369,8 +20365,7 @@ static uint32_t metal_graph_raw_cap_for_context(int ctx_size, uint32_t prefill_c
}

/* Choose the prefill ubatch size. Whole-batch is fastest for normal prompts.
* Long prompts default to 4096-token chunks, except PRO and Strix Halo ROCm
* where larger raw/cache chunks are measurably better on long contexts. */
* Long Flash prompts default to 4096-token chunks; PRO defaults to 8192. */
static uint32_t metal_graph_prefill_cap_for_prompt(int prompt_len,
uint32_t prefill_chunk) {
return ds4_prefill_cap_for_prompt(prompt_len, prefill_chunk);
Expand Down
130 changes: 97 additions & 33 deletions rocm/ds4_rocm_runtime.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ struct cuda_model_image {
const void *host_base;
uint64_t size;
char *device_ptr;
uint64_t device_offset;
};

struct cuda_q8_f16_range {
Expand Down Expand Up @@ -78,6 +79,7 @@ static std::unordered_map<uint64_t, size_t> g_q8_f16_transpose_by_offset;
static uint64_t g_model_range_bytes;
static uint64_t g_q8_f16_bytes;
static int g_q8_f16_disabled_after_oom;
static int g_q8_f16_disabled_for_multi_model;
static int g_q8_f16_budget_notice_printed;
static uint64_t g_model_load_progress_next;
static double g_model_load_progress_last;
Expand Down Expand Up @@ -206,11 +208,14 @@ static int cuda_model_image_find(const void *model_map) {
}

static const char *cuda_model_image_ptr(const void *model_map, uint64_t offset) {
const int idx = cuda_model_image_find(model_map);
if (idx < 0) return NULL;
const cuda_model_image &img = g_model_images[(size_t)idx];
if (offset > img.size) return NULL;
return img.device_ptr + offset;
for (size_t i = 0; i < g_model_images.size(); i++) {
const cuda_model_image &img = g_model_images[i];
if (img.host_base != model_map) continue;
if (offset < img.device_offset) continue;
if (offset - img.device_offset >= img.size) continue;
return img.device_ptr + (offset - img.device_offset);
}
return NULL;
}

static int cuda_model_image_owned(const void *model_map) {
Expand Down Expand Up @@ -268,10 +273,6 @@ static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset,
if (bytes == 0) return cuda_model_ptr(model_map, offset);
if (cuda_model_image_owned(model_map)) return cuda_model_ptr(model_map, offset);

if (model_map != g_model_host_base) {
return cuda_model_range_copy_uncached(model_map, offset, bytes, what);
}

const uint64_t end = offset + bytes;
auto exact = g_model_range_by_offset.find(offset);
if (exact != g_model_range_by_offset.end()) {
Expand All @@ -294,6 +295,10 @@ static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset,
const char *fd_ptr = cuda_model_range_ptr_from_fd(model_map, offset, bytes, what);
if (fd_ptr) return fd_ptr;

if (model_map != g_model_host_base) {
return cuda_model_range_copy_uncached(model_map, offset, bytes, what);
}

cudaError_t err = cudaSuccess;
if (g_model_range_mapping_supported && model_map == g_model_host_base) {
const long page_sz_l = sysconf(_SC_PAGESIZE);
Expand Down Expand Up @@ -550,6 +555,8 @@ static void cuda_q8_f16_cache_disable_after_failure(const char *what, uint64_t r
static int cuda_q8_f16_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) {
if (g_quality_mode) return 0;
if (g_q8_f16_disabled_after_oom) return 0;
if (g_q8_f16_disabled_for_multi_model) return 0;
if (getenv("DS4_CUDA_NO_Q8_F16_CACHE") != NULL) return 0;
if (!label) return 0;
if (strstr(label, "attn_output_a") != NULL ||
strstr(label, "attn_output_b") != NULL ||
Expand Down Expand Up @@ -1046,13 +1053,29 @@ static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) {
void *dev = NULL;
cudaError_t err = cudaMalloc(&dev, (size_t)chunk);
if (err != cudaSuccess) {
fprintf(stderr, DS4_GPU_LOG_PREFIX "model arena alloc failed for %s (%.2f MiB chunk): %s\n",
what ? what : "weights",
(double)chunk / 1048576.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
g_model_cache_full = 1;
return NULL;
uint64_t fallback = chunk / 2u;
while (fallback >= aligned) {
err = cudaMalloc(&dev, (size_t)fallback);
if (err == cudaSuccess) break;
(void)cudaGetLastError();
fallback /= 2u;
}
if (err != cudaSuccess) {
err = cudaMalloc(&dev, (size_t)aligned);
if (err != cudaSuccess) {
fprintf(stderr, DS4_GPU_LOG_PREFIX "model arena alloc failed for %s (%.2f MiB): %s\n",
what ? what : "weights",
(double)aligned / 1048576.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
g_model_cache_full = 1;
return NULL;
}
fallback = aligned;
}
g_model_arenas.push_back({(char *)dev, fallback, aligned});
return (char *)dev;
}
g_model_arenas.push_back({(char *)dev, chunk, aligned});
return (char *)dev;
Expand Down Expand Up @@ -1143,25 +1166,19 @@ static const char *cuda_model_range_ptr_from_fd(

static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) {
if (!model_map || model_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0;
if (cuda_model_image_owned(model_map)) {
g_model_host_base = model_map;
g_model_device_base = cuda_model_image_ptr(model_map, 0);
g_model_registered_size = model_size;
g_model_device_owned = 1;
return 1;
}
if (map_size == 0) return 0;

void *dev = NULL;
const double t0 = cuda_wall_sec();
cudaError_t err = cudaMalloc(&dev, (size_t)model_size);
cudaError_t err = cudaMalloc(&dev, (size_t)map_size);
if (err != cudaSuccess) {
fprintf(stderr, DS4_GPU_LOG_PREFIX "model allocation skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}

fprintf(stderr, DS4_GPU_LOG_PREFIX "chunk-copying %.2f GiB model image\n",
(double)model_size / 1073741824.0);
(double)map_size / 1073741824.0);

const uint64_t chunk = cuda_model_copy_chunk_bytes();
const uint64_t stage_bytes = chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1);
Expand All @@ -1172,8 +1189,8 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u

uint64_t copied = 0;
uint64_t chunk_idx = 0;
while (copied < model_size) {
const uint64_t n = (model_size - copied < chunk) ? (model_size - copied) : chunk;
while (copied < map_size) {
const uint64_t n = (map_size - copied < chunk) ? (map_size - copied) : chunk;
const uint64_t bi = chunk_idx % 4u;
if (chunk_idx >= 4u) {
err = cudaEventSynchronize(g_model_stage_event[bi]);
Expand All @@ -1186,7 +1203,7 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u
}
const char *payload = NULL;
if (!cuda_model_stage_read(g_model_stage[bi], g_model_stage_bytes,
copied, n, &payload)) {
map_offset + copied, n, &payload)) {
fprintf(stderr, DS4_GPU_LOG_PREFIX "model staged read failed at %.2f GiB: %s\n",
(double)copied / 1073741824.0, strerror(errno));
(void)cudaFree(dev);
Expand All @@ -1208,11 +1225,11 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u
(void)cudaGetLastError();
return 0;
}
cuda_model_drop_file_pages(copied, n);
cuda_model_discard_source_pages(model_map, model_size, copied, n);
cuda_model_drop_file_pages(map_offset + copied, n);
cuda_model_discard_source_pages(model_map, model_size, map_offset + copied, n);
copied += n;
chunk_idx++;
cuda_model_load_progress_note(copied > map_offset ? copied - map_offset : 0);
cuda_model_load_progress_note(copied);
}
err = cudaStreamSynchronize(g_model_upload_stream);
if (err != cudaSuccess) {
Expand All @@ -1221,9 +1238,11 @@ static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, u
(void)cudaGetLastError();
return 0;
}
g_model_images.push_back({model_map, model_size, (char *)dev});
g_model_images.push_back({model_map, map_size, (char *)dev, map_offset});
g_model_host_base = model_map;
g_model_device_base = (const char *)dev;
/* With multiple disjoint images there is no single base pointer;
* all tensor lookups go through cuda_model_image_ptr() instead. */
g_model_device_base = NULL;
g_model_registered_size = model_size;
g_model_device_owned = 1;
const double t1 = cuda_wall_sec();
Expand Down Expand Up @@ -1304,6 +1323,7 @@ extern "C" void ds4_gpu_cleanup(void) {
cuda_model_range_release_all();
cuda_q8_f16_cache_release_all();
g_q8_f16_disabled_after_oom = 0;
g_q8_f16_disabled_for_multi_model = 0;
g_q8_f16_budget_notice_printed = 0;
if (g_cuda_tmp) {
(void)cudaFree(g_cuda_tmp);
Expand Down Expand Up @@ -1474,10 +1494,23 @@ extern "C" int ds4_gpu_synchronize(void) { return cuda_ok(cudaDeviceSynchronize(
extern "C" int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) {
if (!model_map || model_size == 0) return 0;
if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1;
const int multi_model =
g_model_host_base != NULL &&
(g_model_host_base != model_map || g_model_registered_size != model_size);
cuda_model_range_release_all();
cuda_q8_f16_cache_release_all();
g_q8_f16_disabled_after_oom = 0;
g_q8_f16_budget_notice_printed = 0;
if (multi_model) {
/*
* MTP loads a second GGUF mapping. Its weights are small, but on UMA
* ROCm systems the optional expanded Q8->F16 cache can consume the
* memory margin needed for session/context tensors once both model
* mappings are resident. The cache is only a speed path; the normal
* Q8 kernels remain available and keep MTP startup reliable.
*/
g_q8_f16_disabled_for_multi_model = 1;
}
g_model_host_base = model_map;
g_model_device_base = cuda_model_image_owned(model_map) ?
cuda_model_image_ptr(model_map, 0) :
Expand Down Expand Up @@ -1511,12 +1544,14 @@ extern "C" int ds4_gpu_set_model_map_spans(
uint64_t max_tensor_bytes) {
(void)max_tensor_bytes;
if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0;
uint64_t span_bytes = 0;
for (uint32_t i = 0; i < count; i++) {
if (offsets[i] > model_size ||
sizes[i] == 0 ||
sizes[i] > model_size - offsets[i]) {
return 0;
}
span_bytes += sizes[i];
}
uint64_t min_offset = offsets[0];
uint64_t max_end = offsets[0] + sizes[0];
Expand All @@ -1526,7 +1561,36 @@ extern "C" int ds4_gpu_set_model_map_spans(
if (end > max_end) max_end = end;
}
if (!ds4_gpu_set_model_map(model_map, model_size)) return 0;
return cuda_model_copy_chunked(model_map, model_size, min_offset, max_end - min_offset);
/* If the bounding box is tight (<=10% waste), copy it as one
* contiguous range. Otherwise sort and merge to avoid copying
* large gaps between non-adjacent spans. */
const uint64_t bbox = max_end - min_offset;
if (bbox <= span_bytes + span_bytes / 10u) {
return cuda_model_copy_chunked(model_map, model_size, min_offset, bbox);
}
/* Bounding box has large gaps. Sort the spans, merge adjacent ones,
* and bulk-copy each contiguous group individually. */
std::vector<std::pair<uint64_t, uint64_t>> sorted(count);
for (uint32_t i = 0; i < count; i++) sorted[i] = {offsets[i], offsets[i] + sizes[i]};
std::sort(sorted.begin(), sorted.end());
uint64_t grp_off = sorted[0].first;
uint64_t grp_end = sorted[0].second;
for (uint32_t i = 1; i <= count; i++) {
/* Merge spans separated by <= 64 KiB to reduce the number of
* device allocations without wasting meaningful memory. */
const uint64_t gap = 64ull * 1024ull;
if (i < count && sorted[i].first <= grp_end + gap) {
if (sorted[i].second > grp_end) grp_end = sorted[i].second;
continue;
}
if (!cuda_model_copy_chunked(model_map, model_size, grp_off, grp_end - grp_off))
return 0;
if (i < count) {
grp_off = sorted[i].first;
grp_end = sorted[i].second;
}
}
return 1;
}

extern "C" int ds4_gpu_set_model_fd(int fd) {
Expand Down