Skip to content

metal: accelerate M5 Max indexed prefill #758

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
rinaldofesta wants to merge 1 commit into
base: main
Choose a base branch
from
+400 −113
Open

metal: accelerate M5 Max indexed prefill #758

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
46 changes: 34 additions & 12 deletions ds4_metal.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@
static id<MTLComputePipelineState> g_soft_max_f32_4_pipeline;
static id<MTLComputePipelineState> g_argsort_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_argsort_merge_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_topk_stream512_pipeline;
static id<MTLComputePipelineState> g_topk_fused512_pipeline;
static id<MTLComputePipelineState> g_sum_rows_f32_f32_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_scatter_pipeline;
Expand All @@ -170,6 +170,7 @@
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_rb16_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads16_dual_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads16_dual_rb16_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_split_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_split_reduce_pipeline;
static id<MTLComputePipelineState> g_dsv4_softplus_sqrt_pipeline;
Expand Down Expand Up @@ -7295,16 +7296,16 @@ int ds4_gpu_init(void) {
return 0;
}

fn = [library newFunctionWithName:@"kernel_topk_stream512"];
fn = [library newFunctionWithName:@"kernel_topk_fused512"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_topk_stream512 function not found\n");
fprintf(stderr, "ds4: Metal kernel_topk_fused512 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_topk_stream512_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_topk_stream512_pipeline) {
fprintf(stderr, "ds4: Metal kernel_topk_stream512 pipeline failed: %s\n",
g_topk_fused512_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_topk_fused512_pipeline) {
fprintf(stderr, "ds4: Metal kernel_topk_fused512 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
Expand Down Expand Up @@ -7707,6 +7708,8 @@ int ds4_gpu_init(void) {
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16");
g_dsv4_indexed_attention_heads16_dual_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual");
g_dsv4_indexed_attention_heads16_dual_rb16_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual_rb16");
g_dsv4_indexed_attention_heads8_split_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split");
g_dsv4_indexed_attention_heads8_split_reduce_pipeline =
Expand Down Expand Up @@ -7849,6 +7852,7 @@ int ds4_gpu_init(void) {
!g_dsv4_indexed_attention_heads8_pipeline ||
!g_dsv4_indexed_attention_heads8_rb16_pipeline ||
!g_dsv4_indexed_attention_heads16_dual_pipeline ||
!g_dsv4_indexed_attention_heads16_dual_rb16_pipeline ||
!g_dsv4_indexed_attention_heads8_split_pipeline ||
!g_dsv4_indexed_attention_heads8_split_reduce_pipeline ||
!g_dsv4_softplus_sqrt_pipeline ||
Expand Down Expand Up @@ -9191,7 +9195,7 @@ void ds4_gpu_cleanup(void) {
g_soft_max_f32_4_pipeline = nil;
g_argsort_f32_i32_desc_pipeline = nil;
g_argsort_merge_f32_i32_desc_pipeline = nil;
g_topk_stream512_pipeline = nil;
g_topk_fused512_pipeline = nil;
g_sum_rows_f32_f32_pipeline = nil;
g_dsv4_topk_mask_pipeline = nil;
g_dsv4_topk_mask_scatter_pipeline = nil;
Expand All @@ -9202,6 +9206,7 @@ void ds4_gpu_cleanup(void) {
g_dsv4_indexed_attention_heads8_pipeline = nil;
g_dsv4_indexed_attention_heads8_rb16_pipeline = nil;
g_dsv4_indexed_attention_heads16_dual_pipeline = nil;
g_dsv4_indexed_attention_heads16_dual_rb16_pipeline = nil;
g_dsv4_indexed_attention_heads8_split_pipeline = nil;
g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil;
g_dsv4_softplus_sqrt_pipeline = nil;
Expand Down Expand Up @@ -16650,7 +16655,15 @@ int ds4_gpu_indexer_topk_tensor(
fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n");
return 0;
}
if (top_k == 512u && n_comp > 1024u && n_tokens >= 32u) {
/* Fuse the two-run M5 case. Wider rows benefit from the canonical
* path's parallel block sorts, so leave them on argsort/merge. */
if (top_k == 512u && n_comp > 1024u && n_comp <= 2048u &&
n_tokens >= 32u &&
(ds4_gpu_device_name_contains("M5") ||
getenv("DS4_METAL_ENABLE_TOPK_FUSED512") != NULL) &&
g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup >= 1024u &&
g_topk_fused512_pipeline.maxTotalThreadsPerThreadgroup >= 1024u &&
getenv("DS4_METAL_DISABLE_TOPK_FUSED512") == NULL) {
ds4_gpu_kargs_argsort args = {
.ne00 = (int32_t)n_comp,
.ne01 = (int32_t)n_tokens,
Expand All @@ -16670,14 +16683,14 @@ int ds4_gpu_indexer_topk_tensor(
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_topk_stream512_pipeline];
[enc setComputePipelineState:g_topk_fused512_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1];
[enc setBuffer:selbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "indexer streaming top-k");
return ds4_gpu_finish_command_buffer(cb, owned, "indexer fused top-k");
}
NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
if (max_threads == 0) max_threads = 256;
Expand Down Expand Up @@ -26479,6 +26492,11 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
!decode_one_token && !g_quality_mode && ds4_gpu_mpp_available() &&
n_head == 64u &&
top_k == 512u && window == 128u && head_dim == 512u;
const bool prefill_dual_heads_rb16 =
prefill_dual_heads && n_tokens >= 32u &&
(ds4_gpu_device_name_contains("M5 Max") ||
getenv("DS4_METAL_ENABLE_INDEXED_ATTN_DUAL_RB16") != NULL) &&
getenv("DS4_METAL_DISABLE_INDEXED_ATTN_DUAL_RB16") == NULL;
const uint32_t decode_splits =
decode_one_token && !g_quality_mode ? 12u : 1u;
const bool split_decode = decode_splits > 1u;
Expand All @@ -26490,6 +26508,9 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
decode_one_token ?
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads8_rb16") :
prefill_dual_heads_rb16 ?
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_rb16_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads16_dual_rb16") :
prefill_dual_heads ?
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads16_dual") :
Expand Down Expand Up @@ -26627,7 +26648,8 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
atIndex:4];
[enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6];
[enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) *
[enc setThreadgroupMemoryLength:
(decode_one_token || prefill_dual_heads_rb16 ? 16u : 1u) *
128u * 4u * sizeof(uint16_t)
atIndex:0];
[enc dispatchThreadgroups:
Expand Down
153 changes: 76 additions & 77 deletions metal/argsort.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,105 +274,104 @@ kernel void kernel_argsort_merge_f32_i32(
// Host-visible merge variant used by DS4 top-k selection.
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<DS4_SORT_ORDER_DESC>;

// Exact top-512 selection for wide prefill rows. Scores and row indices are
// packed into one key, giving a deterministic descending total order. The
// threshold is the 512th-best key seen so far, so discarded entries cannot
// enter the final result.
kernel void kernel_topk_stream512(
// Exact fused top-512 selection for wide M5 prefill rows. This reproduces the
// canonical 1024-wide block sort and left-biased merge order, including its
// tie behavior, while keeping intermediate indices in threadgroup memory.
kernel void kernel_topk_fused512(
constant ds4_metal_args_argsort & args [[buffer(0)]],
device const float *scores [[buffer(1)]],
device uint32_t *selected [[buffer(2)]],
uint token [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
uint ntg [[threads_per_threadgroup]]) {
constexpr uint stream_cap = 1024u;
uint tid [[thread_index_in_threadgroup]]) {
constexpr uint block_width = 1024u;
constexpr uint keep_count = 512u;
threadgroup uint64_t keys[stream_cap];
threadgroup atomic_uint count;
threadgroup uint64_t threshold;
threadgroup int32_t block_idx[block_width];
threadgroup float block_scores[block_width];
threadgroup int32_t retained[keep_count];
threadgroup int32_t merged[keep_count];

if (token >= (uint)args.ne01) return;
if (tid == 0u) {
atomic_store_explicit(&count, 0u, memory_order_relaxed);
threshold = 0u;
}
threadgroup_barrier(mem_flags::mem_threadgroup);

device const float *row = scores + (uint64_t)token * (uint)args.ne00;
const uint n_comp = (uint)args.ne00;
const uint start = (uint)(((uint64_t)(token + 1u) * 0x9E3779B9u) % n_comp);

for (uint base = 0u; base < n_comp; base += ntg) {
const uint i = base + tid;
if (i < n_comp) {
uint col = start + i;
if (col >= n_comp) col -= n_comp;
const uint bits = as_type<uint>(row[col]);
const uint ordered = (bits & 0x80000000u) ? ~bits :
(bits ^ 0x80000000u);
const uint64_t key = ((uint64_t)ordered << 32u) |
(uint64_t)(0xffffffffu - col);
if (key > threshold) {
const uint pos = atomic_fetch_add_explicit(
&count, 1u, memory_order_relaxed);
keys[pos] = key;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);

const uint used = atomic_load_explicit(&count, memory_order_relaxed);
if (used > stream_cap - ntg) {
for (uint j = used + tid; j < stream_cap; j += ntg) keys[j] = 0u;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint base = 0u; base < n_comp; base += block_width) {
const uint col = base + tid;
block_idx[tid] = (int32_t)col;
if (col < n_comp) block_scores[tid] = row[col];
threadgroup_barrier(mem_flags::mem_threadgroup);

for (uint k = 2u; k <= stream_cap; k <<= 1u) {
for (uint j = k >> 1u; j > 0u; j >>= 1u) {
for (uint p = tid; p < stream_cap; p += ntg) {
const uint other = p ^ j;
if (other > p) {
const uint64_t a = keys[p];
const uint64_t b = keys[other];
const bool descending = (p & k) == 0u;
if ((descending && a < b) || (!descending && a > b)) {
keys[p] = b;
keys[other] = a;
}
// Keep this comparison network identical to
// kernel_argsort_f32_i32_desc above. In particular, strict float
// comparisons preserve its deterministic, non-stable tie order.
for (uint k = 2u; k <= block_width; k <<= 1u) {
for (uint j = k >> 1u; j > 0u; j >>= 1u) {
const uint other = tid ^ j;
if (other > tid) {
const int32_t lhs = block_idx[tid];
const int32_t rhs = block_idx[other];
if ((tid & k) == 0u) {
if (lhs >= (int32_t)n_comp ||
(rhs < (int32_t)n_comp &&
block_scores[lhs - (int32_t)base] <
block_scores[rhs - (int32_t)base])) {
block_idx[tid] = rhs;
block_idx[other] = lhs;
}
} else {
if (rhs >= (int32_t)n_comp ||
(lhs < (int32_t)n_comp &&
block_scores[lhs - (int32_t)base] >
block_scores[rhs - (int32_t)base])) {
block_idx[tid] = rhs;
block_idx[other] = lhs;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) {
threshold = keys[keep_count - 1u];
atomic_store_explicit(&count, keep_count, memory_order_relaxed);
}
}

const uint block_len = min(block_width, n_comp - base);
const uint run_len = min(keep_count, block_len);
if (base == 0u) {
if (tid < keep_count) retained[tid] = block_idx[tid];
threadgroup_barrier(mem_flags::mem_threadgroup);
continue;
}
}

const uint used = atomic_load_explicit(&count, memory_order_relaxed);
for (uint j = used + tid; j < stream_cap; j += ntg) keys[j] = 0u;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k = 2u; k <= stream_cap; k <<= 1u) {
for (uint j = k >> 1u; j > 0u; j >>= 1u) {
for (uint p = tid; p < stream_cap; p += ntg) {
const uint other = p ^ j;
if (other > p) {
const uint64_t a = keys[p];
const uint64_t b = keys[other];
const bool descending = (p & k) == 0u;
if ((descending && a < b) || (!descending && a > b)) {
keys[p] = b;
keys[other] = a;
}
if (tid < keep_count) {
const int k0 = (int)tid;
int low = k0 > (int)run_len ? k0 - (int)run_len : 0;
int high = min(k0, (int)keep_count);
while (low < high) {
const int mid = (low + high) >> 1;
const int32_t idx0 = retained[mid];
const int32_t idx1 = block_idx[k0 - mid - 1];
if (row[idx0] >= row[idx1]) {
low = mid + 1;
} else {
high = mid;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);

const int i = low;
const int j = k0 - i;
if (i >= (int)keep_count) {
merged[tid] = block_idx[j];
} else if (j >= (int)run_len) {
merged[tid] = retained[i];
} else if (row[retained[i]] >= row[block_idx[j]]) {
merged[tid] = retained[i];
} else {
merged[tid] = block_idx[j];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid < keep_count) retained[tid] = merged[tid];
threadgroup_barrier(mem_flags::mem_threadgroup);
}

device uint32_t *out = selected + (uint64_t)token * keep_count;
for (uint i = tid; i < keep_count; i += ntg) {
out[i] = 0xffffffffu - (uint)keys[i];
if (tid < keep_count) {
selected[(uint64_t)token * keep_count + tid] = (uint32_t)retained[tid];
}
}
Loading