mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize Metal prefill and decode kernels

Browse Source
This commit is contained in:
Ivan Fioravanti authored and antirez committed 2026-07-20 11:09:49 +02:00
1 parent e0824ddcaa
commit 427e281d24
17 files changed
+8435 -657
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

ds4.c
+200 -73
View File
@@ -19500,6 +19500,11 @@ static bool metal_graph_use_reference_qkv_norm(void) {
return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_NORM_FUSION", &cache);
}
static bool metal_graph_use_reference_qkv_pair_proj(void) {
static int cache = -1;
return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_PAIR_PROJ", &cache);
}
static bool metal_graph_use_reference_compressor_pair_proj(void) {
static int cache = -1;
return metal_graph_env_flag("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ", &cache);
@@ -21555,7 +21560,7 @@ static bool metal_graph_encode_decode_layer_phase(
if (!resume_after_qkv) {
bool qkv_pair_projected = resume_after_qa_kv_raw;
if (!resume_after_qa_kv_raw && ok && qkv_rms_fused &&
g->cuda_qkv_pair) {
g->cuda_qkv_pair && !metal_graph_use_reference_qkv_pair_proj()) {
qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor(
metal_graph_qr(g),
metal_graph_kv_raw(g),
@@ -21768,17 +21773,41 @@ static bool metal_graph_encode_decode_layer_phase(
fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il);
ok = false;
}
bool comp_state_already_stored = false;
if (ok && !metal_graph_use_reference_compressor_pair_proj()) {
ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
model->map,
model->size,
layer->attn_compressor_kv->abs_offset,
layer->attn_compressor_gate->abs_offset,
DS4_N_EMBD,
comp_width,
metal_graph_attn_norm(g),
1) != 0;
const int fused_store =
ds4_gpu_matmul_f16_pair_compressor_store_tensor(
metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
g->layer_attn_state_kv[il],
g->layer_attn_state_score[il],
model->map,
model->size,
layer->attn_compressor_kv->abs_offset,
layer->attn_compressor_gate->abs_offset,
layer->attn_compressor_ape->abs_offset,
layer->attn_compressor_ape->type,
DS4_N_EMBD,
comp_width,
metal_graph_attn_norm(g),
ratio,
pos);
if (fused_store < 0) {
ok = false;
} else if (fused_store > 0) {
comp_state_already_stored = true;
} else {
ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
model->map,
model->size,
layer->attn_compressor_kv->abs_offset,
layer->attn_compressor_gate->abs_offset,
DS4_N_EMBD,
comp_width,
metal_graph_attn_norm(g),
1) != 0;
}
} else {
if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size,
layer->attn_compressor_kv->abs_offset,
@@ -21814,7 +21843,8 @@ static bool metal_graph_encode_decode_layer_phase(
attn_factor,
DS4_ROPE_YARN_BETA_FAST,
DS4_ROPE_YARN_BETA_SLOW,
DS4_RMS_EPS) != 0;
DS4_RMS_EPS,
comp_state_already_stored) != 0;
DS4_METAL_PROFILE_DECODE_STAGE("compressor_update");
if (ok && emit) {
ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row);
@@ -21850,17 +21880,41 @@ static bool metal_graph_encode_decode_layer_phase(
fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il);
ok = false;
}
bool index_state_already_stored = false;
if (ok && !metal_graph_use_reference_compressor_pair_proj()) {
ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
model->map,
model->size,
layer->indexer_compressor_kv->abs_offset,
layer->indexer_compressor_gate->abs_offset,
DS4_N_EMBD,
index_width,
metal_graph_attn_norm(g),
1) != 0;
const int fused_store =
ds4_gpu_matmul_f16_pair_compressor_store_tensor(
metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
g->layer_index_state_kv[il],
g->layer_index_state_score[il],
model->map,
model->size,
layer->indexer_compressor_kv->abs_offset,
layer->indexer_compressor_gate->abs_offset,
layer->indexer_compressor_ape->abs_offset,
layer->indexer_compressor_ape->type,
DS4_N_EMBD,
index_width,
metal_graph_attn_norm(g),
ratio,
pos);
if (fused_store < 0) {
ok = false;
} else if (fused_store > 0) {
index_state_already_stored = true;
} else {
ok = ds4_gpu_matmul_f16_pair_tensor(metal_graph_comp_kv_cur(g),
metal_graph_comp_sc_cur(g),
model->map,
model->size,
layer->indexer_compressor_kv->abs_offset,
layer->indexer_compressor_gate->abs_offset,
DS4_N_EMBD,
index_width,
metal_graph_attn_norm(g),
1) != 0;
}
} else {
if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_kv_cur(g), model->map, model->size,
layer->indexer_compressor_kv->abs_offset,
@@ -21896,7 +21950,8 @@ static bool metal_graph_encode_decode_layer_phase(
attn_factor,
DS4_ROPE_YARN_BETA_FAST,
DS4_ROPE_YARN_BETA_SLOW,
DS4_RMS_EPS) != 0;
DS4_RMS_EPS,
index_state_already_stored) != 0;
DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_update");
if (ok && emit) {
#if defined(__APPLE__)
@@ -23929,22 +23984,47 @@ static bool metal_graph_encode_output_head(
if (ok) {
metal_graph_debug_dump_tensor("result_hc_weights", metal_graph_output_weights(g), DS4_N_HC, DS4_N_LAYER, 0);
}
if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g),
metal_graph_cur_hc(g),
metal_graph_output_weights(g),
DS4_N_EMBD,
DS4_N_HC) != 0;
bool output_sum_norm_fused = false;
#if defined(__APPLE__)
if (ok) {
output_sum_norm_fused =
ds4_gpu_hc_weighted_sum_norm_tensor(
metal_graph_output_embd(g),
metal_graph_output_norm(g),
metal_graph_cur_hc(g),
metal_graph_output_weights(g),
model->map,
model->size,
weights->output_norm->abs_offset,
DS4_N_EMBD,
DS4_N_HC,
DS4_RMS_EPS) != 0;
if (!output_sum_norm_fused &&
getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL) {
ok = false;
}
}
#endif
if (ok && !output_sum_norm_fused) {
ok = ds4_gpu_hc_weighted_sum_tensor(metal_graph_output_embd(g),
metal_graph_cur_hc(g),
metal_graph_output_weights(g),
DS4_N_EMBD,
DS4_N_HC) != 0;
}
DS4_METAL_PROFILE_OUTPUT_STAGE("hc_weighted_sum");
if (ok) {
metal_graph_debug_dump_tensor("result_hc", metal_graph_output_embd(g), DS4_N_EMBD, DS4_N_LAYER, 0);
}
if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g),
metal_graph_output_embd(g),
model->map,
model->size,
weights->output_norm->abs_offset,
DS4_N_EMBD,
DS4_RMS_EPS) != 0;
if (ok && !output_sum_norm_fused) {
ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_output_norm(g),
metal_graph_output_embd(g),
model->map,
model->size,
weights->output_norm->abs_offset,
DS4_N_EMBD,
DS4_RMS_EPS) != 0;
}
DS4_METAL_PROFILE_OUTPUT_STAGE("output_norm");
if (ok) {
metal_graph_debug_dump_tensor("result_norm", metal_graph_output_norm(g), DS4_N_EMBD, DS4_N_LAYER, 0);
@@ -26265,6 +26345,52 @@ static bool metal_graph_upload_prompt_embeddings_hc(
n_tokens);
}
static bool metal_graph_hc_rms_scale_project(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_scratch,
const ds4_model *model,
const ds4_tensor *weight,
const ds4_gpu_tensor *x,
uint64_t in_dim,
uint32_t n_tokens) {
if (!out || !norm_scratch || !model || !weight || !x ||
in_dim > UINT32_MAX) {
return false;
}
#if defined(__APPLE__)
return ds4_gpu_hc_rms_scale_project_f16_tensor(
out,
norm_scratch,
model->map,
model->size,
weight->abs_offset,
(uint32_t)in_dim,
2u * DS4_N_HC + DS4_N_HC * DS4_N_HC,
x,
n_tokens,
DS4_RMS_EPS) != 0;
#else
bool ok = ds4_gpu_rms_norm_plain_rows_tensor(
norm_scratch,
x,
(uint32_t)in_dim,
n_tokens,
DS4_RMS_EPS) != 0;
if (ok) {
ok = ds4_gpu_matmul_f16_tensor(
out,
model->map,
model->size,
weight->abs_offset,
in_dim,
2u * DS4_N_HC + DS4_N_HC * DS4_N_HC,
norm_scratch,
n_tokens) != 0;
}
return ok;
#endif
}
static bool metal_graph_warmup_prefill_kernels(
ds4_gpu_graph *g,
const ds4_model *model,
@@ -26293,18 +26419,17 @@ static bool metal_graph_warmup_prefill_kernels(
if (!metal_graph_set_active_tier_batch(g, g->placement[1], n_tokens)) return false;
}
const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD;
const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC;
bool ok = ds4_gpu_begin_commands() != 0;
if (ok) {
ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_hc_mix(g),
model->map,
model->size,
weights->layer[0].hc_attn_fn->abs_offset,
hc_dim,
mix_hc,
metal_graph_batch_flat_hc(g),
n_tokens) != 0;
ok = metal_graph_hc_rms_scale_project(
metal_graph_batch_hc_mix(g),
metal_graph_batch_flat_hc(g),
model,
weights->layer[0].hc_attn_fn,
metal_graph_batch_cur_hc(g),
hc_dim,
n_tokens);
}
if (ok) ok = ds4_gpu_end_commands() != 0;
if (!ok) {
@@ -26633,22 +26758,17 @@ static bool metal_graph_encode_layer_attention_batch(
ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view(
metal_graph_batch_after_attn_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float));
bool ok = hc_mix_view && hc_split_view && attn_cur_view && after_attn_hc_view;
const bool fuse_hc_norm = DS4_N_HC == 4 &&
const bool fuse_hc_norm = n_tokens > 1 &&
DS4_N_HC == 4 &&
!metal_graph_use_reference_hc_decode() &&
metal_graph_enable_batch_hc_norm_fusion();
if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g),
metal_graph_batch_cur_hc(g),
(uint32_t)hc_dim,
n_tokens,
DS4_RMS_EPS) != 0;
if (ok) ok = ds4_gpu_matmul_f16_tensor(hc_mix_view,
model->map,
model->size,
layer->hc_attn_fn->abs_offset,
hc_dim,
mix_hc,
metal_graph_batch_flat_hc(g),
n_tokens) != 0;
if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view,
metal_graph_batch_flat_hc(g),
model,
layer->hc_attn_fn,
metal_graph_batch_cur_hc(g),
hc_dim,
n_tokens);
if (metal_graph_use_reference_hc_decode()) {
if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view,
hc_mix_view,
@@ -27308,7 +27428,8 @@ static bool metal_graph_encode_layer_attention_batch(
attn_factor,
DS4_ROPE_YARN_BETA_FAST,
DS4_ROPE_YARN_BETA_SLOW,
DS4_RMS_EPS) != 0;
DS4_RMS_EPS,
false) != 0;
if (ok && emit) {
ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row);
ok = comp_row_view &&
@@ -27602,7 +27723,8 @@ static bool metal_graph_encode_layer_attention_batch(
attn_factor,
DS4_ROPE_YARN_BETA_FAST,
DS4_ROPE_YARN_BETA_SLOW,
DS4_RMS_EPS) != 0;
DS4_RMS_EPS,
false) != 0;
if (ok && emit) {
ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view(
g->layer_index_comp_cache[il],
@@ -28332,21 +28454,17 @@ static bool metal_graph_encode_layer_ffn_batch(
ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view(
metal_graph_batch_next_hc(g), 0, (uint64_t)n_tokens * hc_dim * sizeof(float));
bool ok = hc_mix_view && hc_split_view && ffn_cur_view && next_hc_view;
const bool fuse_hc_norm = DS4_N_HC == 4 &&
const bool fuse_hc_norm = n_tokens > 1 &&
DS4_N_HC == 4 &&
!metal_graph_use_reference_hc_decode() &&
metal_graph_enable_batch_hc_norm_fusion();
if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(metal_graph_batch_flat_hc(g),
metal_graph_batch_after_attn_hc(g),
(uint32_t)hc_dim,
n_tokens,
DS4_RMS_EPS) != 0;
if (ok) ok = metal_graph_matmul_plain_tensor(hc_mix_view,
model,
layer->hc_ffn_fn,
hc_dim,
mix_hc,
metal_graph_batch_flat_hc(g),
n_tokens);
if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view,
metal_graph_batch_flat_hc(g),
model,
layer->hc_ffn_fn,
metal_graph_batch_after_attn_hc(g),
hc_dim,
n_tokens);
if (metal_graph_use_reference_hc_decode()) {
if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view,
hc_mix_view,
@@ -32954,6 +33072,9 @@ static bool metal_graph_prefill_layer_major(
}
return false;
}
#ifdef __APPLE__
ds4_gpu_release_zero_prefix_prefill_mask_cache();
#endif
const double t_before_read = profile ? now_sec() : 0.0;
if (logits) {
@@ -33363,6 +33484,12 @@ static bool metal_graph_prefill_layer_major(
}
return false;
}
#ifdef __APPLE__
/* Zero-prefix masks are shared across the 43 per-layer command batches,
* then become dead weight. Release them before the output head and later
* replay/decode chunks so the prefill win does not add residency pressure. */
ds4_gpu_release_zero_prefix_prefill_mask_cache();
#endif
if (show_progress) fputc('\n', stderr);
metal_graph_stream_prefill_selected_profile_summary(g);
#ifdef DS4_ROCM_BUILD
ds4_cuda.cu
+42 -5
View File
@@ -13376,6 +13376,40 @@ extern "C" int ds4_gpu_matmul_f16_pair_tensor(
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
}
extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
(void)out_kv;
(void)out_score;
(void)state_kv;
(void)state_score;
(void)model_map;
(void)model_size;
(void)weight_kv_offset;
(void)weight_score_offset;
(void)ape_offset;
(void)ape_type;
(void)in_dim;
(void)width;
(void)x;
(void)ratio;
(void)pos;
return 0;
}
extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) return 0;
if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0;
@@ -13980,7 +14014,8 @@ extern "C" int ds4_gpu_compressor_update_tensor(
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
float rms_eps,
bool state_already_stored) {
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
!model_map || head_dim == 0 || ratio == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
@@ -14004,10 +14039,12 @@ extern "C" int ds4_gpu_compressor_update_tensor(
(emit && comp_cache->bytes < comp_bytes)) {
return 0;
}
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
if (!state_already_stored) {
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
}
}
if (!emit) return 1;
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
ds4_gpu.h
+65 -1
View File
@@ -140,6 +140,9 @@ void ds4_gpu_set_quality(bool quality);
void ds4_gpu_set_glm_model(bool enabled);
void ds4_gpu_set_ssd_streaming(bool enabled);
void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled);
#ifdef __APPLE__
void ds4_gpu_release_zero_prefix_prefill_mask_cache(void);
#endif
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);
@@ -738,6 +741,26 @@ int ds4_gpu_matmul_f16_pair_tensor(
const ds4_gpu_tensor *x,
uint64_t n_tok);
/* Optional Metal decode fusion. Returns 1 when the paired projection and
* recurrent compressor-state store were encoded, 0 when the optimized path
* is unavailable, and -1 on an attempted-path error. */
int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos);
int ds4_gpu_matmul_f32_tensor(
ds4_gpu_tensor *out,
const void *model_map,
@@ -1569,7 +1592,8 @@ int ds4_gpu_compressor_update_tensor(
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps);
float rms_eps,
bool state_already_stored);
int ds4_gpu_compressor_store_batch_tensor(
const ds4_gpu_tensor *kv,
@@ -1717,6 +1741,19 @@ int ds4_gpu_attention_decode_rows_rope_tensor(
float attn_factor,
float beta_fast,
float beta_slow);
/* Diagnostic/public form of the dk=512 gathered decode-attention KV staging
* step. The compressed source must be F16; dst writes chronological raw-ring
* rows followed by compressed rows and must not overlap either source. */
int ds4_gpu_flash_kv_stage_f16_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
const ds4_gpu_tensor *comp,
uint32_t comp_is_f16,
uint32_t n_comp,
uint32_t head_dim);
int ds4_gpu_attention_prefill_raw_heads_tensor(
ds4_gpu_tensor *heads,
@@ -2356,6 +2393,18 @@ int ds4_gpu_hc_weighted_sum_tensor(
uint32_t n_embd,
uint32_t n_hc);
int ds4_gpu_hc_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
float norm_eps);
int ds4_gpu_hc_weighted_sum_split_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
@@ -2396,6 +2445,21 @@ int ds4_gpu_hc_split_weighted_sum_norm_tensor(
float eps,
float norm_eps);
/* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the
* tuned Metal path, scale_scratch stores one float per row instead of the
* full normalized HC tensor; other shapes retain the established fallback. */
int ds4_gpu_hc_rms_scale_project_f16_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *scale_scratch,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows,
float eps);
int ds4_gpu_output_hc_weights_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *pre,
ds4_metal.m
+2176 -398
View File
File diff suppressed because it is too large. Load diff
metal/argsort.metal
+13 -4
View File
@@ -65,6 +65,15 @@ kernel void kernel_argsort_f32_i32(
// initialize indices
shmem_i32[col] = i00 + col;
// Stage this block's score slice in threadgroup memory (indices stay in
// [i00, i00+ntg.x), so shmem_f32[idx - i00] replaces the device gather).
// The host allocates ntg.x extra floats after the index array. Values and
// the comparison network are unchanged, so the permutation is identical.
threadgroup float * shmem_f32 = (threadgroup float *) (shmem_i32 + ntg.x);
if (i00 + col < args.ne00) {
shmem_f32[col] = src0_row[i00 + col];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int k = 2; k <= ntg.x; k *= 2) {
@@ -74,16 +83,16 @@ kernel void kernel_argsort_f32_i32(
if ((col & k) == 0) {
if (shmem_i32[col] >= args.ne00 ||
(shmem_i32[ixj] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00] :
shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00]))
) {
SWAP(shmem_i32[col], shmem_i32[ixj]);
}
} else {
if (shmem_i32[ixj] >= args.ne00 ||
(shmem_i32[col] < args.ne00 && (order == DS4_SORT_ORDER_ASC ?
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00] :
shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00]))
) {
SWAP(shmem_i32[col], shmem_i32[ixj]);
}
metal/cpy.metal
+189
View File
@@ -56,3 +56,192 @@ template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<
template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, half>;
template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<half, float>;
template [[host_name("kernel_cpy_f16_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<half, half>;
// Contiguous 1D conversions avoid the generic tensor-index reconstruction
// above. Packed vector types retain scalar alignment, so tensor views whose
// offsets are float/half aligned do not need additional 16/8-byte alignment.
// The final vector is converted element-by-element when n is not divisible by
// four, preserving the generic kernel's bounds and conversion semantics.
kernel void kernel_cpy_contig_f32_f16_4(
constant uint & n,
device const packed_float4 * src,
device packed_half4 * dst,
uint gid [[thread_position_in_grid]]) {
const uint i = gid * 4u;
if (i >= n) {
return;
}
const uint remaining = n - i;
if (remaining >= 4u) {
const float4 value = float4(src[gid]);
dst[gid] = packed_half4(half4(value));
return;
}
device const float * src_scalar = (device const float *)src;
device half * dst_scalar = (device half *)dst;
for (uint lane = 0; lane < remaining; ++lane) {
dst_scalar[i + lane] = half(src_scalar[i + lane]);
}
}
kernel void kernel_cpy_contig_f16_f32_4(
constant uint & n,
device const packed_half4 * src,
device packed_float4 * dst,
uint gid [[thread_position_in_grid]]) {
const uint i = gid * 4u;
if (i >= n) {
return;
}
const uint remaining = n - i;
if (remaining >= 4u) {
const half4 value = half4(src[gid]);
dst[gid] = packed_float4(float4(value));
return;
}
device const half * src_scalar = (device const half *)src;
device float * dst_scalar = (device float *)dst;
for (uint lane = 0; lane < remaining; ++lane) {
dst_scalar[i + lane] = float(src_scalar[i + lane]);
}
}
// Bitwise F16 transport for cache staging. Use ushort rather than half so NaN
// payloads and every other binary16 encoding pass through unchanged.
kernel void kernel_cpy_contig_f16_f16_bits_4(
constant uint & n,
device const packed_ushort4 * src,
device packed_ushort4 * dst,
uint gid [[thread_position_in_grid]]) {
const uint i = gid * 4u;
if (i >= n) {
return;
}
const uint remaining = n - i;
if (remaining >= 4u) {
dst[gid] = src[gid];
return;
}
device const ushort * src_scalar = (device const ushort *)src;
device ushort * dst_scalar = (device ushort *)dst;
for (uint lane = 0; lane < remaining; ++lane) {
dst_scalar[i + lane] = src_scalar[i + lane];
}
}
struct ds4_metal_args_flash_kv_stage_f16 {
uint raw_cap;
uint raw_start;
uint n_raw;
uint n_comp;
uint pad_rows;
uint shared_pad;
};
// Decode-time gathered attention consumes a logical raw-cache ring followed
// by an already-F16 compressed cache. Pack both regions into the contiguous
// F16 FlashAttention scratch in one dispatch. The raw conversion expression
// and compressed ushort4 transport exactly match the standalone copy kernels.
kernel void kernel_dsv4_flash_kv_stage_f16(
constant ds4_metal_args_flash_kv_stage_f16 & args,
device const char * raw_src,
device const char * comp_src,
device char * dst,
device const char * mask_src,
device char * pad_dst,
uint gid [[thread_position_in_grid]]) {
constexpr uint row_vecs = 128;
const uint raw_vecs = args.n_raw * row_vecs;
const uint n_keys = args.n_raw + args.n_comp;
const uint total_vecs = n_keys * row_vecs;
if (gid < raw_vecs) {
const uint logical_row = gid >> 7;
const uint col = gid & 127u;
uint physical_row = args.raw_start + logical_row;
if (physical_row >= args.raw_cap) {
physical_row -= args.raw_cap;
}
device const packed_float4 *raw =
(device const packed_float4 *)raw_src;
device packed_half4 *dst_half = (device packed_half4 *)dst;
const float4 value =
float4(raw[physical_row * row_vecs + col]);
dst_half[gid] = packed_half4(half4(value));
return;
}
if (gid < total_vecs) {
device const packed_ushort4 *comp =
(device const packed_ushort4 *)comp_src;
device packed_ushort4 *dst_bits = (device packed_ushort4 *)dst;
dst_bits[gid] = comp[gid - raw_vecs];
return;
}
// The vector FlashAttention kernel redirects its final partial block to
// a compact K/V/mask buffer. When requested, append those writes to this
// dispatch so gathered decode does not need a standalone pad dispatch.
const uint pad_rows = args.pad_rows;
const uint pad_vecs = pad_rows * row_vecs;
uint pad_gid = gid - total_vecs;
if (pad_gid < pad_vecs) {
const uint row = pad_gid >> 7;
const uint col = pad_gid & 127u;
const uint valid_rows = n_keys % pad_rows;
device packed_half4 *pad_half = (device packed_half4 *)pad_dst;
device packed_ushort4 *pad_bits = (device packed_ushort4 *)pad_dst;
if (row >= valid_rows) {
const packed_half4 zero = packed_half4(half4(0.0h));
pad_half[pad_gid] = zero;
if (!args.shared_pad) {
pad_half[pad_vecs + pad_gid] = zero;
}
return;
}
const uint logical_row = n_keys - valid_rows + row;
if (logical_row < args.n_raw) {
uint physical_row = args.raw_start + logical_row;
if (physical_row >= args.raw_cap) {
physical_row -= args.raw_cap;
}
device const packed_float4 *raw =
(device const packed_float4 *)raw_src;
const float4 value =
float4(raw[physical_row * row_vecs + col]);
const packed_half4 value_half = packed_half4(half4(value));
pad_half[pad_gid] = value_half;
if (!args.shared_pad) {
pad_half[pad_vecs + pad_gid] = value_half;
}
} else {
device const packed_ushort4 *comp =
(device const packed_ushort4 *)comp_src;
const packed_ushort4 value_bits =
comp[(logical_row - args.n_raw) * row_vecs + col];
pad_bits[pad_gid] = value_bits;
if (!args.shared_pad) {
pad_bits[pad_vecs + pad_gid] = value_bits;
}
}
return;
}
pad_gid -= pad_vecs;
if (pad_gid < pad_rows) {
const uint valid_rows = n_keys % pad_rows;
device const ushort *mask_bits = (device const ushort *)mask_src;
device ushort *pad_mask_bits =
(device ushort *)pad_dst + 2u * pad_vecs * 4u;
pad_mask_bits[pad_gid] = pad_gid < valid_rows
? mask_bits[n_keys - valid_rows + pad_gid]
: 0xfbffu;
}
}
metal/dense.metal
+469 -31
View File
@@ -25,6 +25,13 @@ struct ds4_metal_args_mul_mv {
short r3;
};
struct ds4_metal_args_compressor_pair_store {
uint32_t width;
uint32_t ratio;
uint32_t pos;
uint32_t ape_type;
};
struct ds4_metal_args_mul_mm {
int32_t ne00;
int32_t ne02;
@@ -203,6 +210,141 @@ kernel void kernel_mul_mv_q8_0_f32_r4(
kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
// Output projection alias used by the optimized host dispatch.
[[host_name("kernel_mul_mv_q8_0_f32_nr4")]]
kernel void kernel_mul_mv_q8_0_f32_nr4(
constant ds4_metal_args_mul_mv & args,
device const char * src0,
device const char * src1,
device char * dst,
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(
args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
// Decode Q-A/KV pair. Both projections consume the same activation row but
// have independent weight ranges and output extents. Keep the standalone Q8_0
// lane/block traversal and two-stage reduction verbatim for each bank; only
// the activation load and threadgroup scheduling are shared.
[[host_name("kernel_mul_mv_q8_0_f32_pair")]]
kernel void kernel_mul_mv_q8_0_f32_pair(
constant ds4_metal_args_mul_mv & args0,
constant ds4_metal_args_mul_mv & args1,
device const char * src0_a,
device const char * src0_b,
device const char * src1,
device char * dst_a,
device char * dst_b,
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
const short NSG = FC_mul_mv_nsg;
constexpr short NW = N_SIMDWIDTH;
constexpr short NQ = 8;
constexpr short NR0 = N_R0_Q8_0;
const int nb = args0.ne00 / QK8_0;
const int r0 = tgpig.x * NR0;
const bool active_a = r0 < args0.ne01;
const bool active_b = r0 < args1.ne01;
device const float *y = (device const float *)src1;
device const block_q8_0 *ax_a[NR0];
device const block_q8_0 *ax_b[NR0];
FOR_UNROLL (short row = 0; row < NR0; ++row) {
const int out_row = r0 + row;
ax_a[row] = active_a && out_row < args0.ne01
? (device const block_q8_0 *)(src0_a + (uint64_t)out_row * args0.nb01)
: (device const block_q8_0 *)src0_a;
ax_b[row] = active_b && out_row < args1.ne01
? (device const block_q8_0 *)(src0_b + (uint64_t)out_row * args1.nb01)
: (device const block_q8_0 *)src0_b;
}
float suma[NR0] = { 0.f };
float sumb[NR0] = { 0.f };
const short ix = tiisg / (NW / NQ);
const short il = tiisg % (NW / NQ);
const int ib0 = sgitg * NQ + ix;
float yl[NQ];
device const float *yb = y + ib0 * QK8_0 + il * NQ;
for (int ib = ib0; ib < nb; ib += NSG * NQ) {
FOR_UNROLL (short i = 0; i < NQ; ++i) {
yl[i] = yb[i];
}
FOR_UNROLL (short row = 0; row < NR0; ++row) {
const int out_row = r0 + row;
if (active_a && out_row < args0.ne01) {
device const int8_t *qs = ax_a[row][ib].qs + il * NQ;
float sumq = 0.f;
FOR_UNROLL (short i = 0; i < NQ; ++i) {
sumq += qs[i] * yl[i];
}
suma[row] += sumq * ax_a[row][ib].d;
}
if (active_b && out_row < args1.ne01) {
device const int8_t *qs = ax_b[row][ib].qs + il * NQ;
float sumq = 0.f;
FOR_UNROLL (short i = 0; i < NQ; ++i) {
sumq += qs[i] * yl[i];
}
sumb[row] += sumq * ax_b[row][ib].d;
}
}
yb += NSG * NQ * QK8_0;
}
threadgroup float *shared = (threadgroup float *)shmem;
threadgroup float *sha[NR0];
threadgroup float *shb[NR0];
FOR_UNROLL (short row = 0; row < NR0; ++row) {
sha[row] = shared + NW * row;
shb[row] = shared + NW * (NR0 + row);
if (sgitg == 0) {
sha[row][tiisg] = 0.0f;
if (active_b) shb[row][tiisg] = 0.0f;
}
suma[row] = simd_sum(suma[row]);
if (active_b) sumb[row] = simd_sum(sumb[row]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
FOR_UNROLL (short row = 0; row < NR0; ++row) {
if (tiisg == 0) {
sha[row][sgitg] = suma[row];
if (active_b) shb[row][sgitg] = sumb[row];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
device float *out_a = (device float *)dst_a;
device float *out_b = (device float *)dst_b;
FOR_UNROLL (short row = 0; row < NR0; ++row) {
const float total_a = simd_sum(sha[row][tiisg]);
if (tiisg == 0 && sgitg == 0) {
const int out_row = r0 + row;
if (active_a && out_row < args0.ne01) out_a[out_row] = total_a;
}
if (active_b) {
const float total_b = simd_sum(shb[row][tiisg]);
if (tiisg == 0 && sgitg == 0) {
const int out_row = r0 + row;
if (out_row < args1.ne01) out_b[out_row] = total_b;
}
}
}
}
// Decode shared-expert gate/up projections followed by SwiGLU:
//
// mid = silu(min(gate, limit)) * clamp(up, -limit, limit)
@@ -779,6 +921,59 @@ kernel void kernel_mul_mv_f16_f32_pair_4(
args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg);
}
// Decode compressor projection plus recurrent-state append. The paired
// matvec remains unchanged and still materializes both F32 outputs. After a
// device-memory barrier, the first NR0 threads reload those exact stored bits
// and perform the same state write and score+APE addition as
// kernel_dsv4_compressor_store_one.
kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4(
constant ds4_metal_args_mul_mv & args,
constant ds4_metal_args_compressor_pair_store & store,
device const char * src0_a,
device const char * src0_b,
device const char * src1,
device char * dst_a,
device char * dst_b,
device const char * ape,
device float * state_kv,
device float * state_score,
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tiitg [[thread_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
kernel_mul_mv_f16_f32_pair_4_disp<constant ds4_metal_args_mul_mv &>(
args, src0_a, src0_b, src1, dst_a, dst_b,
shmem, tgpig, tiisg, sgitg);
threadgroup_barrier(mem_flags::mem_device);
if (tiitg >= args.nr0 || store.width == 0u || store.ratio == 0u) {
return;
}
const uint col = tgpig.x * (uint)args.nr0 + tiitg;
if (col >= store.width) return;
const uint pos_mod = store.pos % store.ratio;
const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod;
const uint dst = dst_row * store.width + col;
const uint ape_i = pos_mod * store.width + col;
device volatile const float * projected_kv =
(device volatile const float *)dst_a;
device volatile const float * projected_score =
(device volatile const float *)dst_b;
float ape_v;
if (store.ape_type == 1u) {
ape_v = (float)(((device const half *)ape)[ape_i]);
} else {
ape_v = ((device const float *)ape)[ape_i];
}
state_kv[dst] = projected_kv[col];
state_score[dst] = projected_score[col] + ape_v;
}
template<typename T0, typename T1, typename args_t>
void kernel_mul_mv_t_t_short_impl(
args_t args,
@@ -919,6 +1114,27 @@ void dequantize_dense_q4_K(device const ds4_dense_block_q4_K *xb, short il, thre
reg = (type4x4)reg_f;
}
/*
* Bit-identical twin of dequantize_q8_0 for the MPP staging loop: same
* half(float(qs[i]) * d) per element, but the 16 consecutive int8 lanes are
* fetched as eight aligned 16-bit loads instead of sixteen byte loads.
* xb->qs is always 2-byte aligned (block_q8_0.d is a half).
*/
void dequantize_q8_0_pairs(device const block_q8_0 *xb, short il, thread half4x4 & reg) {
device const ushort *qs16 = (device const ushort *)(xb->qs + 16*il);
const float d = xb->d;
float4x4 reg_f;
FOR_UNROLL (short i = 0; i < 8; i++) {
const ushort u = qs16[i];
reg_f[i/2][(i%2)*2 + 0] = ((float)(int8_t)(u & 0xFF)) * d;
reg_f[i/2][(i%2)*2 + 1] = ((float)(int8_t)(u >> 8)) * d;
}
reg = (half4x4) reg_f;
}
template <typename type4>
void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
device const int8_t * qs = ((device const int8_t *)xb->qs);
@@ -1386,6 +1602,11 @@ template [[host_name("kernel_mul_mm_f16_f32_mpp")]] kernel mul_mm_mpp_t kernel_
// aligned F16/Q8_0 prompt matmuls. The host selects the widest token tile that
// evenly divides the batch, with 128-token tiles retained after the 64-token
// retest was neutral or slower.
//
// The host dispatch guarantees M % NR0 == 0 and K % NK == 0, so the dequant
// stage does no bounds work. The weight tile is double-buffered: the next
// k-step's dequant overlaps the current cooperative matmul, so the k-loop
// needs one threadgroup barrier per step instead of two.
template<
short NR1,
typename SA, typename SA_4x4, typename block_q, short nl,
@@ -1419,7 +1640,8 @@ kernel void kernel_mul_mm_mpp_direct_rhs(
const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
threadgroup SA *sa = (threadgroup SA *)shmem;
auto tA = tensor(sa, dextents<int32_t, 2>(NK, NR0));
auto tA0 = tensor(sa, dextents<int32_t, 2>(NK, NR0));
auto tA1 = tensor(sa + NR0*NK, dextents<int32_t, 2>(NK, NR0));
device T1 *ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
const int strideB = args.nb11/sizeof(T1);
@@ -1430,7 +1652,7 @@ kernel void kernel_mul_mm_mpp_direct_rhs(
matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<4>> mm;
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA0), float>();
#pragma unroll
for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
@@ -1439,43 +1661,48 @@ kernel void kernel_mul_mm_mpp_direct_rhs(
}
}
for (int loop_k = 0; loop_k < K; loop_k += NK) {
for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) {
const int row = work/NL;
const int k_chunk = work%NL;
// NR0*NL/NUM_THREADS 16-value weight chunks per thread (1 at NR0=64).
auto stage_tile = [&](const int loop_k, threadgroup SA *buf) {
FOR_UNROLL (int work = tiitg; work < NR0*NL; work += NUM_THREADS) {
const int row = work / NL;
const int k_chunk = work % NL;
const int k_pos = loop_k + k_chunk*16;
const short k_base = k_chunk*16;
if (r0 + row < M) {
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0);
FOR_UNROLL (short i = 0; i < 16; i++) {
sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0;
}
} else {
const int block_idx = k_pos/(16*nl);
const short il = (k_pos/16)%nl;
device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0);
SA_4x4 temp_a;
dequantize_func(row_ptr + block_idx, il, temp_a);
FOR_UNROLL (short i = 0; i < 16; i++) {
sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
}
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
device const T0 *row_ptr_f =
(device const T0 *)(srcA + args.nb01*(r0 + row) + offset0);
FOR_UNROLL (short i = 0; i < 16; i++) {
buf[row*NK + k_base + i] = (SA)row_ptr_f[k_pos + i];
}
} else {
FOR_UNROLL (short i = 0; i < 16; i++) {
sa[row*NK + k_base + i] = (SA)0;
}
device const block_q *row_ptr =
(device const block_q *)(srcA + args.nb01*(r0 + row) + offset0);
SA_4x4 temp_a;
dequantize_func(row_ptr + k_pos/(16*nl), (k_pos/16)%nl, temp_a);
typedef vec<SA, 4> SA4;
threadgroup SA4 *dst4 = (threadgroup SA4 *)(buf + row*NK + k_base);
dst4[0] = temp_a[0];
dst4[1] = temp_a[1];
dst4[2] = temp_a[2];
dst4[3] = temp_a[3];
}
}
};
threadgroup_barrier(mem_flags::mem_threadgroup);
stage_tile(0, sa);
threadgroup_barrier(mem_flags::mem_threadgroup);
auto mA = tA.slice(0, 0);
uint buf_sel = 0;
for (int loop_k = 0; loop_k < K; loop_k += NK) {
auto mA = (buf_sel ? tA1 : tA0).slice(0, 0);
auto mB = tB.slice(loop_k, r1);
mm.run(mB, mA, cT);
const int next_k = loop_k + NK;
if (next_k < K) {
buf_sel ^= 1u;
stage_tile(next_k, buf_sel ? sa + NR0*NK : sa);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
@@ -1490,9 +1717,6 @@ typedef decltype(kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, float4x4, 1, de
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>;
@@ -1500,6 +1724,9 @@ template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs")]] kernel mul_mm_mp
template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>;
#endif
// Tiled matrix-matrix kernel used for prompt batches larger than 8. DS4 uses
@@ -1706,6 +1933,217 @@ kernel void kernel_mul_mm(
}
}
// Legacy F16-weight/F32-RHS prefill matmul with a per-row RMS scale applied at
// the existing F32-to-F16 RHS staging boundary. The tile layout, half inputs,
// float accumulators, and output path intentionally mirror kernel_mul_mm so the
// only arithmetic change versus materializing RMSNorm first is where x*scale is
// rounded from F32 to F16.
kernel void kernel_mul_mm_f16_f32_scaled(
constant ds4_metal_args_mul_mm & args,
device const char * src0,
device const char * src1,
device char * dst,
device const float * scales,
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiitg[[thread_index_in_threadgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
threadgroup half * sa = (threadgroup half *)(shmem);
threadgroup half * sb = (threadgroup half *)(shmem + 4096);
constexpr int NR0 = 64;
constexpr int NR1 = 32;
constexpr int NK = 32;
constexpr int NL0 = NK/16;
constexpr int NL1 = NK/8;
const int im = tgpig.z;
const int r0 = tgpig.y*NR0;
const int r1 = tgpig.x*NR1;
// if this block is of 64x32 shape or smaller
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
// a thread shouldn't load data outside of the matrix
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
const short il0 = (tiitg % NL0);
short il = il0;
const int i12 = im%args.ne12;
const int i13 = im/args.ne12;
const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03;
const short offset1 = il0;
device const half4x4 * x = (device const half4x4 *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
const short iy = 8*(tiitg % NL1);
device const float * y = (device const float *)(src1
+ args.nb13*i13
+ args.nb12*i12
+ args.nb11*(r1 + lr1)
+ args.nb10*iy);
const float row_scale = scales[r1 + lr1];
simdgroup_half8x8 ma[4];
simdgroup_half8x8 mb[2];
simdgroup_float8x8 mc[8];
for (short i = 0; i < 8; i++){
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
}
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
// load data and store to threadgroup memory
if (FC_mul_mm_bc_inp) {
threadgroup_barrier(mem_flags::mem_threadgroup);
// no need for dequantization
for (short i = 0; i < 16; i++) {
const short sx = 2*il0 + i/8;
const short sy = (tiitg/NL0)/8;
const short lx = (tiitg/NL0)%8;
const short ly = i%8;
const short ib = 8*sx + sy;
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device half *) x + i) : 0;
}
} else {
half4x4 temp_a;
dequantize_f16(x, il, temp_a);
threadgroup_barrier(mem_flags::mem_threadgroup);
FOR_UNROLL (short i = 0; i < 16; i++) {
const short sx = 2*il0 + i/8;
const short sy = (tiitg/NL0)/8;
const short lx = (tiitg/NL0)%8;
const short ly = i%8;
const short ib = 8*sx + sy;
// Pointer-form store matches the legacy F16 tile layout.
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
}
}
if (FC_mul_mm_bc_inp) {
for (short i = 0; i < 8; ++i) {
const short sx = (tiitg%NL1);
const short sy = (tiitg/NL1)/8;
const short lx = i;
const short ly = (tiitg/NL1)%8;
const short ib = 4*sx + sy;
const float scaled = loop_k + iy + i < args.ne00 ? y[i] * row_scale : 0.0f;
*(sb + 64*ib + 8*ly + lx) = (half)scaled;
}
} else {
const short sx = (tiitg%NL1);
const short sy = (tiitg/NL1)/8;
const short ly = (tiitg/NL1)%8;
const short ib = 4*sx + sy;
const float2x4 raw = *((device const float2x4 *) y);
float2x4 scaled;
scaled[0] = raw[0] * row_scale;
scaled[1] = raw[1] * row_scale;
*(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)scaled;
}
il = (il + 2 < 1) ? il + 2 : il % 2;
x = (il < 2) ? x + 2 : x;
y += NK;
threadgroup_barrier(mem_flags::mem_threadgroup);
// load matrices from threadgroup memory and conduct outer products
threadgroup const half * lsma = (sa + 4*64*(sgitg%2));
threadgroup const half * lsmb = (sb + 2*64*(sgitg/2));
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 4; i++) {
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 2; i++) {
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 8; i++){
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
}
lsma += 8*64;
lsmb += 4*64;
}
}
if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
// if no bounds checks on the output are needed, we can directly write to device memory
device float * C = (device float *) dst +
(r0 + 32*(sgitg & 1)) + \
(r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
for (short i = 0; i < 8; i++) {
simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
}
} else {
// block is smaller than 64x32, we should avoid writing data outside of the matrix
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
for (short i = 0; i < 8; i++) {
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sgitg == 0) {
for (int j = tiitg; j < nr1; j += NR1) {
device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
device float4 * D4 = (device float4 *) D;
threadgroup float * C = temp_str + (j*NR0);
threadgroup float4 * C4 = (threadgroup float4 *) C;
int i = 0;
for (; i < nr0/4; i++) {
*(D4 + i) = *(C4 + i);
}
i *= 4;
for (; i < nr0; i++) {
*(D + i) = *(C + i);
}
}
}
}
}
kernel void kernel_mul_mm_f16_f32_pair(
constant ds4_metal_args_mul_mm & args,
device const char * src0_a,
metal/dsv4_hc.metal
+137 -5
View File
@@ -21,6 +21,26 @@ struct ds4_metal_args_dsv4_hc_weighted_sum {
uint64_t nb1;
};
struct ds4_metal_args_dsv4_hc_weighted_sum_norm {
int64_t n_embd;
int64_t n_hc;
int64_t n_tokens;
uint64_t nb_x0;
uint64_t nb_x1;
uint64_t nb_x2;
uint64_t nb_w0;
uint64_t nb_w1;
uint64_t nb0;
uint64_t nb1;
uint64_t nb_norm1;
float norm_eps;
};
struct ds4_metal_args_dsv4_output_hc_weights4 {
float post_scale;
float eps;
};
struct ds4_metal_args_dsv4_hc_split_weighted_sum {
int64_t n_embd;
int32_t n_hc;
@@ -506,10 +526,12 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
device const float4 *x1 = (device const float4 *)(x + 1 * args.nb_x1 + (uint64_t)row * args.nb_x2);
device const float4 *x2 = (device const float4 *)(x + 2 * args.nb_x1 + (uint64_t)row * args.nb_x2);
device const float4 *x3 = (device const float4 *)(x + 3 * args.nb_x1 + (uint64_t)row * args.nb_x2);
const float4 v = x0[i] * pre_shmem[0] +
x1[i] * pre_shmem[1] +
x2[i] * pre_shmem[2] +
x3[i] * pre_shmem[3];
// Preserve the standalone HC collapse's explicit accumulation order.
float4 v = 0.0f;
v += x0[i] * pre_shmem[0];
v += x1[i] * pre_shmem[1];
v += x2[i] * pre_shmem[2];
v += x3[i] * pre_shmem[3];
row_shmem[i] = v;
sumf += dot(v, v);
}
@@ -523,7 +545,11 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4(
sumf = sum_shmem[tiisg];
sumf = simd_sum(sumf);
const float norm_scale = rsqrt(sumf / float(n_embd) + args.norm_eps);
// Batched prefill must match kernel_rms_norm_fuse_impl so enabling this
// fusion does not change its scale by an ULP. Keep the established
// single-row decode result unchanged; that path historically used rsqrt.
const float norm_arg = sumf / float(n_embd) + args.norm_eps;
const float norm_scale = args.n_rows > 1 ? 1.0f / sqrt(norm_arg) : rsqrt(norm_arg);
device float4 *dst4 = (device float4 *)(dst + (uint64_t)row * args.nb1);
device const float4 *w4 = (device const float4 *)norm_weight;
@@ -883,3 +909,109 @@ kernel void kernel_dsv4_hc_weighted_sum(
*((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc;
}
// The one-row output head immediately applies a learned RMSNorm after reducing
// its four HC streams. Preserve the standalone scalar HC accumulation, write
// the collapsed row for diagnostics, then reload its F32 values from
// threadgroup memory using the standalone RMSNorm's float4 reduction mapping.
kernel void kernel_dsv4_hc_weighted_sum_norm4(
constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args,
device const char * x,
device const char * weights,
device char * dst,
device const char * norm_weight,
device char * norm_dst,
threadgroup float * shared [[threadgroup(0)]],
ushort tid [[thread_position_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort ntg [[threads_per_threadgroup]]) {
if (args.n_tokens != 1 || args.n_hc != 4 ||
args.n_embd <= 0 || (args.n_embd & 3) != 0) {
return;
}
const uint n_embd = uint(args.n_embd);
const uint n4 = n_embd >> 2;
threadgroup float *row_shmem = shared;
threadgroup float *sum_shmem = shared + n_embd;
if (sgitg == 0) {
sum_shmem[tiisg] = 0.0f;
}
for (uint d = tid; d < n_embd; d += ntg) {
float acc = 0.0f;
for (int64_t h = 0; h < args.n_hc; ++h) {
const float xv = *((device const float *)(
x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1));
const float wv = *((device const float *)(
weights + (uint64_t)h*args.nb_w0));
acc += xv * wv;
}
row_shmem[d] = acc;
*((device float *)(dst + (uint64_t)d*args.nb0)) = acc;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup const float4 *row4 =
(threadgroup const float4 *)row_shmem;
float sumf = 0.0f;
for (uint i = tid; i < n4; i += ntg) {
sumf += dot(row4[i], row4[i]);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
sum_shmem[sgitg] = sumf;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
sumf = sum_shmem[tiisg];
sumf = simd_sum(sumf);
const float mean = sumf/args.n_embd;
const float scale = 1.0f/sqrt(mean + args.norm_eps);
device const float4 *w4 = (device const float4 *)norm_weight;
device float4 *norm4 = (device float4 *)norm_dst;
for (uint i = tid; i < n4; i += ntg) {
norm4[i] = (row4[i]*scale)*w4[i];
}
}
// The one-row HC=4 output head historically materializes four device-F32
// stages across separate launches. Collapse those launches into one tiny
// two-thread group while preserving the scalar/vector lane mapping and every
// global rounding boundary.
kernel void kernel_dsv4_output_hc_weights4(
constant ds4_metal_args_dsv4_output_hc_weights4 & args,
device const float * pre,
device const float * hc_scale,
device const float * hc_base,
device float * dst,
ushort tid [[thread_position_in_threadgroup]]) {
device volatile float *stage = (device volatile float *)dst;
for (uint i = tid; i < 4; i += 2) {
stage[i] = pre[i] * hc_scale[0];
}
threadgroup_barrier(mem_flags::mem_device);
for (uint i = tid; i < 4; i += 2) {
stage[i] = stage[i] + hc_base[i];
}
threadgroup_barrier(mem_flags::mem_device);
if (tid == 0) {
const float4 x = *((device volatile float4 *)stage);
*((device volatile float4 *)stage) = 1 / (1 + exp(-x));
}
threadgroup_barrier(mem_flags::mem_device);
if (tid == 0) {
const float4 x = *((device volatile float4 *)stage);
*((device volatile float4 *)stage) =
args.post_scale * x + args.eps;
}
}
metal/dsv4_kv.metal
+55
View File
@@ -41,6 +41,13 @@ struct ds4_metal_args_dsv4_ratio4_shift {
uint32_t width;
};
struct ds4_metal_args_dsv4_compressor_pack_ratio4 {
uint32_t head_dim;
uint32_t n_comp;
uint32_t replay;
uint32_t n_threads;
};
struct ds4_metal_args_dsv4_compressor_store_one {
uint32_t width;
uint32_t ratio;
@@ -264,6 +271,54 @@ kernel void kernel_dsv4_kv_fp8_store_f32(
}
}
// Builds the two ratio-4 softmax-pool packs together. Each output plane holds
// four previous-half rows followed by four current-half rows. Normal prefill
// seeds plane zero with 0/-inf; replay seeds it from the previous compressor
// state. Later planes take their previous half from the prior input group.
kernel void kernel_dsv4_compressor_pack_ratio4(
constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args,
device const uint * kv,
device const uint * score,
device const uint * state_kv,
device const uint * state_score,
device uint * packed_kv,
device uint * packed_score,
uint2 group [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]) {
if (group.x >= args.n_comp || group.y >= 8u ||
args.head_dim == 0u || args.n_threads == 0u) {
return;
}
const uint plane = group.x;
const uint row = group.y;
const uint64_t input_row_stride = 2ull * args.head_dim;
const uint64_t dst_row = ((uint64_t)plane * 8u + row) * args.head_dim;
for (uint col = tid; col < args.head_dim; col += args.n_threads) {
const uint64_t dst = dst_row + col;
if (row >= 4u) {
const uint token = plane * 4u + (row - 4u);
const uint64_t src = (uint64_t)token * input_row_stride +
args.head_dim + col;
packed_kv[dst] = kv[src];
packed_score[dst] = score[src];
} else if (plane != 0u) {
const uint token = (plane - 1u) * 4u + row;
const uint64_t src = (uint64_t)token * input_row_stride + col;
packed_kv[dst] = kv[src];
packed_score[dst] = score[src];
} else if (args.replay != 0u) {
const uint64_t src = (uint64_t)row * input_row_stride + col;
packed_kv[dst] = state_kv[src];
packed_score[dst] = state_score[src];
} else {
packed_kv[dst] = 0u;
packed_score[dst] = 0xff800000u;
}
}
}
// Ratio-4 compression keeps two 4-row halves of recurrent state. After an
// emitted compressed row, the second half becomes the next window's previous
// half. The old encoder expressed this as four generic copies; this DS4-specific
metal/dsv4_misc.metal
+424 -37
View File
@@ -43,6 +43,21 @@ struct ds4_metal_args_dsv4_softmax_pool {
uint64_t nb1;
};
struct ds4_metal_args_dsv4_softmax_pool_ratio4_direct {
int64_t n_rows;
uint32_t head_dim;
uint32_t n_comp;
uint32_t replay;
uint32_t pad;
};
struct ds4_metal_args_dsv4_compressor_score_ape {
uint32_t width;
uint32_t ratio;
uint32_t pos0;
uint32_t n_tokens;
};
struct ds4_metal_args_dsv4_indexed_attention {
uint32_t n_tokens;
uint32_t n_head;
@@ -4621,6 +4636,54 @@ kernel void kernel_glm_router_select_one(
}
}
// Batched Flash-router weight finalization after selection is already known.
// Six active lanes deliberately match kernel_sum_rows_f32_f32's reduction
// topology. The denominator and divided weights cross threadgroup storage
// boundaries so division cannot be reassociated with the final scale.
kernel void kernel_dsv4_router_weights_batch(
constant float &scale,
device const float *probs,
device const int32_t *selected,
device float *weights,
threadgroup volatile float *scratch [[threadgroup(0)]],
uint row [[threadgroup_position_in_grid]],
ushort tid [[thread_position_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]]) {
if (tid >= 6) return;
threadgroup volatile float *sum_scratch = scratch;
threadgroup volatile float *denom_scratch = scratch + 32;
threadgroup volatile float *div_scratch = scratch + 33;
const uint out_index = row * 6u + (uint)tid;
const int32_t expert = selected[out_index];
const float p = probs[row * 256u + (uint)expert];
// Keep this sequence identical to kernel_sum_rows_f32_f32 for width 6.
if (sgitg == 0) {
sum_scratch[tiisg] = 0.0f;
}
float sumf = 0.0f;
sumf += p;
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
sum_scratch[sgitg] = sumf;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
sumf = sum_scratch[tiisg];
sumf = simd_sum(sumf);
if (tid == 0) {
denom_scratch[0] = clamp(sumf, 6.103515625e-5f, INFINITY);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
div_scratch[tid] = p / denom_scratch[0];
threadgroup_barrier(mem_flags::mem_threadgroup);
weights[out_index] = div_scratch[tid] * scale;
}
// Decode router selection for one token after the existing
// sqrt(softplus(logit)) probability kernel has run. Bias affects only top-k
// selection. Route-weight normalization deliberately stays in the old one-token
@@ -4682,6 +4745,176 @@ kernel void kernel_dsv4_router_finalize_one(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// M3 decode specialization for the non-hash one-token router. Scores and ids
// stay in registers. Intra-SIMD bitonic stages use shuffle-xor; the six stages
// that cross 32-lane SIMD groups exchange through alternating threadgroup
// banks. The next bank's publish barrier proves every prior-bank read finished;
// by the time a bank is reused two cross stages later, no reader can remain.
kernel void kernel_dsv4_router_finalize_one_simd(
constant ds4_metal_args_dsv4_router_select_one & args,
device const float *probs,
device const float *bias,
device const int32_t *hash,
device const int32_t *tokens,
device int32_t *selected,
threadgroup float *scratch [[threadgroup(0)]],
uint tid [[thread_position_in_threadgroup]]) {
if (tid >= 256 || args.hash_mode) return;
(void)hash;
(void)tokens;
threadgroup float *score0_tg = scratch;
threadgroup int32_t *idx0_tg =
(threadgroup int32_t *)(scratch + 256);
threadgroup float *score1_tg = scratch + 512;
threadgroup int32_t *idx1_tg =
(threadgroup int32_t *)(scratch + 768);
const float p = probs[tid];
float score = args.has_bias ? p + bias[tid] : p;
int32_t idx = (int32_t)tid;
uint cross_stage = 0;
for (uint k = 2; k <= 256; k <<= 1) {
for (uint j = k >> 1; j > 0; j >>= 1) {
float peer_score;
int32_t peer_idx;
bool take_peer;
const bool lower = (tid & j) == 0;
const bool descending = (tid & k) == 0;
if (j < 32) {
peer_score = simd_shuffle_xor(score, (ushort)j);
peer_idx = simd_shuffle_xor(idx, (ushort)j);
take_peer = descending
? (lower ? score < peer_score : score > peer_score)
: (lower ? score > peer_score : score < peer_score);
if (take_peer) {
score = peer_score;
idx = peer_idx;
}
} else {
threadgroup float *score_tg =
(cross_stage & 1u) != 0u ? score1_tg : score0_tg;
threadgroup int32_t *idx_tg =
(cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
score_tg[tid] = score;
idx_tg[tid] = idx;
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint other = tid ^ j;
peer_score = score_tg[other];
peer_idx = idx_tg[other];
take_peer = descending
? (lower ? score < peer_score : score > peer_score)
: (lower ? score > peer_score : score < peer_score);
if (take_peer) {
score = peer_score;
idx = peer_idx;
}
cross_stage++;
}
}
}
if (tid < 6) {
selected[tid] = idx;
}
}
// M3 decode specialization that extends the register/TG SIMD selection above
// through the existing six-value serial weight normalization. The selected ids
// cross the same device-memory boundary as the standalone weight kernel;
// volatile TG stores pin its left-fold and scaled-reciprocal rounding points.
kernel void kernel_dsv4_router_finalize_weights_one_simd(
constant ds4_metal_args_dsv4_router_select_one & args,
device const float *probs,
device const float *bias,
device const int32_t *hash,
device const int32_t *tokens,
device int32_t *selected,
device float *weights,
threadgroup float *scratch [[threadgroup(0)]],
uint tid [[thread_position_in_threadgroup]]) {
if (tid >= 256 || args.hash_mode) return;
(void)hash;
(void)tokens;
threadgroup float *score0_tg = scratch;
threadgroup int32_t *idx0_tg =
(threadgroup int32_t *)(scratch + 256);
threadgroup float *score1_tg = scratch + 512;
threadgroup int32_t *idx1_tg =
(threadgroup int32_t *)(scratch + 768);
const float p = probs[tid];
float score = args.has_bias ? p + bias[tid] : p;
int32_t idx = (int32_t)tid;
uint cross_stage = 0;
for (uint k = 2; k <= 256; k <<= 1) {
for (uint j = k >> 1; j > 0; j >>= 1) {
float peer_score;
int32_t peer_idx;
bool take_peer;
const bool lower = (tid & j) == 0;
const bool descending = (tid & k) == 0;
if (j < 32) {
peer_score = simd_shuffle_xor(score, (ushort)j);
peer_idx = simd_shuffle_xor(idx, (ushort)j);
take_peer = descending
? (lower ? score < peer_score : score > peer_score)
: (lower ? score > peer_score : score < peer_score);
if (take_peer) {
score = peer_score;
idx = peer_idx;
}
} else {
threadgroup float *score_tg =
(cross_stage & 1u) != 0u ? score1_tg : score0_tg;
threadgroup int32_t *idx_tg =
(cross_stage & 1u) != 0u ? idx1_tg : idx0_tg;
score_tg[tid] = score;
idx_tg[tid] = idx;
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint other = tid ^ j;
peer_score = score_tg[other];
peer_idx = idx_tg[other];
take_peer = descending
? (lower ? score < peer_score : score > peer_score)
: (lower ? score > peer_score : score < peer_score);
if (take_peer) {
score = peer_score;
idx = peer_idx;
}
cross_stage++;
}
}
}
if (tid < 6) {
selected[tid] = idx;
}
threadgroup_barrier(mem_flags::mem_device);
threadgroup volatile float *norm_scratch =
(threadgroup volatile float *)scratch;
if (tid == 0) {
device const int32_t *s = selected;
norm_scratch[0] = 0.0f;
for (uint i = 0; i < 6; i++) {
norm_scratch[0] = norm_scratch[0] + probs[s[i]];
}
norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f);
norm_scratch[1] = 1.5f / norm_scratch[0];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid < 6) {
device const int32_t *s = selected;
weights[tid] = probs[s[tid]] * norm_scratch[1];
}
}
// Fills the dense compressed-attention mask with -inf. The selected top-k rows
// are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch.
kernel void kernel_dsv4_topk_mask(
@@ -5528,6 +5761,15 @@ kernel void kernel_dsv4_indexer_scores_tiled(
// compressed-row dot tile. The kernel intentionally leaves top-k selection and
// indexed attention semantics unchanged; all 512 selected rows remain available
// to the later attention kernel.
//
// Each matmul processes a pair of heads (TQ = 2 x TM q rows): the per-element
// dot is still a 128-deep reduction in 32-wide k-steps, so scores are
// bit-identical to single-head tiles while the run count halves. The q tile
// is double-buffered, so the next k-step's stage overlaps the current
// cooperative matmul and each pair needs 5 barriers instead of 10. q and k
// staging use one float4/half4 per lane (each thread covers one row of 8/32
// consecutive elements), which is the same half(float) conversion per element
// as the scalar form.
kernel void kernel_dsv4_indexer_scores_nax(
constant ds4_metal_args_dsv4_indexer_scores_fused & args,
device const char *q,
@@ -5538,6 +5780,7 @@ kernel void kernel_dsv4_indexer_scores_nax(
uint2 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]]) {
constexpr int TM = 16;
constexpr int TQ = 32;
constexpr int TN = 32;
constexpr int NK = 32;
constexpr int D = 128;
@@ -5549,9 +5792,9 @@ kernel void kernel_dsv4_indexer_scores_nax(
const uint c0 = tgpig.x * TN;
const uint t0 = tgpig.y * TM;
threadgroup half *qtg = shared; // [16][32]
threadgroup half *ktg = qtg + TM*NK; // [32][128]
threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [16][32], column-major
threadgroup half *qtg = shared; // 2 x [TQ][NK]
threadgroup half *ktg = qtg + 2*TQ*NK; // [32][128]
threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [TQ][TN], column-major
const uint last_token = min(t0 + (uint)TM, args.n_tokens);
const uint max_visible = last_token > t0 ?
@@ -5572,17 +5815,21 @@ kernel void kernel_dsv4_indexer_scores_nax(
return;
}
for (uint work = tid; work < TN*D; work += NUM_THREADS) {
const uint cc = work / D;
const uint d = work - cc*D;
{
// One compressed row per 4 threads, 32 consecutive floats per thread.
const uint cc = tid / 4;
const uint comp = c0 + cc;
half v = half(0.0f);
device const float *krow = nullptr;
if (comp < args.n_comp) {
device const float *krow = (device const float *)(index_comp +
krow = (device const float *)(index_comp +
(uint64_t)comp * args.index_row_stride);
v = half(krow[d]);
}
ktg[cc*D + d] = v;
const uint d0 = (tid % 4) * 32;
FOR_UNROLL (uint j = 0; j < 8; j++) {
const float4 kv = krow ? *(device const float4 *)(krow + d0 + 4*j)
: float4(0.0f);
*(threadgroup half4 *)(ktg + cc*D + d0 + 4*j) = half4(kv);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -5592,17 +5839,45 @@ kernel void kernel_dsv4_indexer_scores_nax(
acc[j] = 0.0f;
}
auto tq = tensor(qtg, dextents<int32_t, 2>(NK, TM));
auto tq0 = tensor(qtg, dextents<int32_t, 2>(NK, TQ));
auto tq1 = tensor(qtg + TQ*NK, dextents<int32_t, 2>(NK, TQ));
auto tk = tensor(ktg, dextents<int32_t, 2>(D, TN));
auto td = tensor(dot, dextents<int32_t, 2>(TM, TN), array<int, 2>({1, TM}));
auto td = tensor(dot, dextents<int32_t, 2>(TQ, TN), array<int, 2>({1, TQ}));
matmul2d<
matmul2d_descriptor(TN, TM, NK, false, true, false,
matmul2d_descriptor(TN, TQ, NK, false, true, false,
matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<4>> mm;
for (uint head = 0; head < args.n_head; head++) {
auto ct = mm.template get_destination_cooperative_tensor<decltype(tk), decltype(tq), float>();
// One q row per 4 threads, 8 consecutive floats per thread. Row r covers
// head (r / TM) of the pair and token row (r % TM).
const uint q_r = tid / 4;
const uint q_k4 = (tid % 4) * 8;
const uint q_hl = q_r / TM;
const uint q_tr = q_r % TM;
const uint q_token = t0 + q_tr;
device const char *q_row_base = nullptr;
if (q_token < args.n_tokens) {
q_row_base = q + (uint64_t)q_token * args.q_token_stride;
}
auto stage_q = [&](const uint head0, const uint loop_k, threadgroup half *buf) {
const uint head = head0 + q_hl;
half4 v0 = half4(0.0f);
half4 v1 = half4(0.0f);
if (q_row_base && head < args.n_head) {
device const float4 *src4 = (device const float4 *)
(q_row_base + (uint64_t)head * args.q_head_stride +
(uint64_t)(loop_k + q_k4) * sizeof(float));
v0 = half4(src4[0]);
v1 = half4(src4[1]);
}
*(threadgroup half4 *)(buf + q_r*NK + q_k4) = v0;
*(threadgroup half4 *)(buf + q_r*NK + q_k4 + 4) = v1;
};
for (uint head0 = 0; head0 < args.n_head; head0 += 2) {
auto ct = mm.template get_destination_cooperative_tensor<decltype(tk), decltype(tq0), float>();
#pragma unroll
for (uint16_t i = 0; i < ct.get_capacity(); i++) {
if (ct.is_valid_element(i)) {
@@ -5610,27 +5885,19 @@ kernel void kernel_dsv4_indexer_scores_nax(
}
}
for (uint loop_k = 0; loop_k < D; loop_k += NK) {
for (uint work = tid; work < TM*NK; work += NUM_THREADS) {
const uint r = work / NK;
const uint k = work - r*NK;
const uint token = t0 + r;
half v = half(0.0f);
if (token < args.n_tokens) {
device const float *qrow = (device const float *)(q +
(uint64_t)token * args.q_token_stride +
(uint64_t)head * args.q_head_stride);
v = half(qrow[loop_k + k]);
}
qtg[r*NK + k] = v;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
stage_q(head0, 0, qtg);
threadgroup_barrier(mem_flags::mem_threadgroup);
auto mq = tq.slice(0, 0);
auto mk = tk.slice(loop_k, 0);
uint qsel = 0;
FOR_UNROLL (uint i = 0; i < 4; i++) {
auto mk = tk.slice(i*NK, 0);
auto mq = (qsel ? tq1 : tq0).slice(0, 0);
mm.run(mk, mq, ct);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (i < 3) {
qsel ^= 1u;
stage_q(head0, (i + 1)*NK, qsel ? qtg + TQ*NK : qtg);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
ct.store(td);
@@ -5646,12 +5913,16 @@ kernel void kernel_dsv4_indexer_scores_nax(
if (token < args.n_tokens) {
device const float *w = (device const float *)(weights +
(uint64_t)token * args.weights_token_stride);
acc[j] += max(dot[cc*TM + r], 0.0f) * (w[head] * args.scale);
acc[j] += max(dot[cc*TQ + r], 0.0f) * (w[head0] * args.scale);
if (head0 + 1 < args.n_head) {
acc[j] += max(dot[cc*TQ + TM + r], 0.0f) * (w[head0 + 1] * args.scale);
}
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// No barrier here: the next pair's q stage and these dot reads touch
// different buffers, and the next q-stage barrier separates the next
// ct.store from these reads.
}
#pragma unroll
@@ -5700,6 +5971,39 @@ kernel void kernel_dsv4_indexer_weighted_sum(
*((device float *) (dst + ic*args.nb0 + it*args.nb1)) = acc;
}
// Adds the periodic compressor APE directly to projected scores. The legacy
// path materializes one repeated APE segment per period and then performs this
// same single F32 add; these kernels remove only that intermediate copy graph.
kernel void kernel_dsv4_compressor_score_ape_f32(
constant ds4_metal_args_dsv4_compressor_score_ape & args,
device const float *score,
device const float *ape,
device float *dst,
uint gid [[thread_position_in_grid]]) {
const uint64_t total = (uint64_t)args.n_tokens * args.width;
if ((uint64_t)gid >= total) return;
const uint token = gid / args.width;
const uint col = gid - token*args.width;
const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio);
dst[gid] = score[gid] + ape[(uint64_t)ape_row*args.width + col];
}
kernel void kernel_dsv4_compressor_score_ape_f16(
constant ds4_metal_args_dsv4_compressor_score_ape & args,
device const float *score,
device const half *ape,
device float *dst,
uint gid [[thread_position_in_grid]]) {
const uint64_t total = (uint64_t)args.n_tokens * args.width;
if ((uint64_t)gid >= total) return;
const uint token = gid / args.width;
const uint col = gid - token*args.width;
const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio);
dst[gid] = score[gid] + float(ape[(uint64_t)ape_row*args.width + col]);
}
// Fused softmax-weighted pooling of compressed KV rows. It is used when several
// compressor rows are present; the one-row case deliberately follows the
// unfused softmax/mul/sum graph in Objective-C to keep identical reductions.
@@ -5768,3 +6072,86 @@ kernel void kernel_dsv4_tp_flag_set(
atomic_store_explicit(&flag, value, memory_order_relaxed);
}
}
// Ratio-4 compressor pooling without materializing the [n_comp, 8, head_dim]
// KV and score packs. The row mapping and both reduction loops deliberately
// match kernel_dsv4_softmax_pool so the arithmetic order is unchanged.
kernel void kernel_dsv4_softmax_pool_ratio4_direct(
constant ds4_metal_args_dsv4_softmax_pool_ratio4_direct & args,
device const float * kv,
device const float * score,
device const float * state_kv,
device const float * state_score,
device float * dst,
uint gid [[thread_position_in_grid]]) {
const uint64_t n = (uint64_t)args.head_dim * args.n_comp;
if ((uint64_t)gid >= n || args.head_dim == 0u) {
return;
}
const uint64_t id = gid % args.head_dim;
const uint64_t ic = gid / args.head_dim;
const uint64_t input_row_stride = 2ull * args.head_dim;
float max_s = -INFINITY;
float sum = 0.0f;
float acc = 0.0f;
if (ic != 0u) {
const int64_t token_base = (int64_t)ic * 4 - 4;
for (int64_t ir = 0; ir < args.n_rows; ++ir) {
const uint64_t token = (uint64_t)(token_base + ir);
const uint64_t src = token * input_row_stride +
((uint64_t)ir >> 2u) * args.head_dim + id;
const float s = score[src];
max_s = max(max_s, s);
}
for (int64_t ir = 0; ir < args.n_rows; ++ir) {
const uint64_t token = (uint64_t)(token_base + ir);
const uint64_t src = token * input_row_stride +
((uint64_t)ir >> 2u) * args.head_dim + id;
const float s = score[src];
const float w = exp(s - max_s);
const float v = kv[src];
sum += w;
acc += v*w;
}
} else {
for (int64_t ir = 0; ir < args.n_rows; ++ir) {
float s;
if (ir >= 4) {
const uint64_t src = (uint64_t)(ir - 4) * input_row_stride +
args.head_dim + id;
s = score[src];
} else if (args.replay != 0u) {
s = state_score[(uint64_t)ir * input_row_stride + id];
} else {
s = -INFINITY;
}
max_s = max(max_s, s);
}
for (int64_t ir = 0; ir < args.n_rows; ++ir) {
float s;
float v;
if (ir >= 4) {
const uint64_t src = (uint64_t)(ir - 4) * input_row_stride +
args.head_dim + id;
s = score[src];
v = kv[src];
} else if (args.replay != 0u) {
const uint64_t src = (uint64_t)ir * input_row_stride + id;
s = state_score[src];
v = state_kv[src];
} else {
s = -INFINITY;
v = 0.0f;
}
const float w = exp(s - max_s);
sum += w;
acc += v*w;
}
}
dst[ic * args.head_dim + id] = acc/sum;
}
metal/dsv4_rope.metal
+220
View File
@@ -24,6 +24,23 @@ struct ds4_metal_args_dsv4_rope_tail {
bool src2;
};
struct ds4_metal_args_dsv4_rope_affine_pair {
uint64_t row_bytes;
uint64_t token_bytes;
int32_t head_dim;
int32_t n_dims;
int32_t n_ctx_orig;
int32_t inverse;
uint32_t pos0;
uint32_t pos_step;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
};
static float rope_yarn_ramp(const float low, const float high, const int i0) {
const float y = (i0 / 2 - low) / max(0.001f, high - low);
return 1.0f - min(1.0f, max(0.0f, y));
@@ -163,3 +180,206 @@ kernel void kernel_dsv4_rope_tail_f32(
}
}
}
// DS4 only calls the Metal RoPE helper in-place and uses the adjacent-pair
// layout (mode 0). The generic kernel above still copies the no-position
// prefix, even though source and destination alias, and dispatches enough
// lanes for the full head. This specialization maps lanes directly to the
// rotated pairs and deliberately never reads or writes the unchanged prefix.
// Keep the arithmetic below in the same order as the mode-0 branch above so
// the optimized and reference paths remain bit-identical.
kernel void kernel_dsv4_rope_tail_f32_inplace_pair(
constant ds4_metal_args_dsv4_rope_tail & args,
device const char * src0,
device const char * src1,
device const char * src2,
device char * dst,
uint tid [[thread_index_in_threadgroup]],
ushort3 ntg [[threads_per_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
if (args.mode != 0) {
return;
}
const int i1 = tgpig[0];
const int i2 = tgpig[1];
const int i3 = tgpig[2];
const int n_nope = args.ne00 - args.n_dims;
if (n_nope < 0) {
return;
}
device const int32_t * pos = (device const int32_t *) src1;
float corr_dims[2];
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
const float theta_base = (float) pos[i2];
const float inv_ndims = -1.f/args.n_dims;
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
// Keep each pair on the same SIMD lane as the generic in-place dispatch.
// n_nope is 32-aligned for the supported DS4 shapes, so logical tail index
// r ran on lane r%32 in the reference kernel. Compacting pairs would move
// them to lane (r/2)%32 and can perturb fast-math code generation.
for (int r = tid; r < args.n_dims; r += ntg.x) {
if ((r & 1) != 0) {
continue;
}
const int ic = r/2;
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
#else
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
#endif
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
if (args.inverse) {
sin_theta = -sin_theta;
}
const int j0 = n_nope + r;
const int j1 = j0 + 1;
const float x0 = *((device const float *) (src_base + j0*args.nb00));
const float x1 = *((device const float *) (src_base + j1*args.nb00));
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
}
}
// The Q/K RoPE calls use the same position and scaling parameters for every
// head. Group four heads into one threadgroup so one 64-thread cohort computes
// the 32 adjacent-pair coefficients and the other cohorts reuse them. Keeping
// r on the same r%32 SIMD lane as the per-head specialization preserves the
// fast-math instruction mapping; only the redundant coefficient work changes.
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4(
constant ds4_metal_args_dsv4_rope_tail & args,
device const char * src0,
device const char * src1,
device const char * src2,
device char * dst,
uint tid [[thread_index_in_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
if (args.mode != 0 || args.n_dims != 64) {
return;
}
const int n_nope = args.ne00 - args.n_dims;
if (n_nope < 0) {
return;
}
const uint cohort = tid >> 6;
const int r = (int)(tid & 63u);
threadgroup float cos_shared[32];
threadgroup float sin_shared[32];
device const int32_t * pos = (device const int32_t *) src1;
float corr_dims[2];
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
const int i2 = tgpig[1];
const float theta_base = (float) pos[i2];
const float inv_ndims = -1.f/args.n_dims;
if (cohort == 0 && (r & 1) == 0) {
const int ic = r/2;
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
#else
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
#endif
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
if (args.inverse) {
sin_theta = -sin_theta;
}
cos_shared[ic] = cos_theta;
sin_shared[ic] = sin_theta;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const int i1 = (int)tgpig[0]*4 + (int)cohort;
if (i1 >= args.ne01 || (r & 1) != 0) {
return;
}
const int i3 = tgpig[2];
device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01;
device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
const int j0 = n_nope + r;
const int j1 = j0 + 1;
const float x0 = *((device const float *) (src_base + j0*args.nb00));
const float x1 = *((device const float *) (src_base + j1*args.nb00));
const float cos_theta = cos_shared[r/2];
const float sin_theta = sin_shared[r/2];
*((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta;
*((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta;
}
// DS4 positions are always affine within one RoPE dispatch. This variant
// reconstructs the same wrapped int32 position in-kernel, avoiding the host
// position array and its buffer binding while preserving the pair lane mapping
// and all floating-point operations of the specialization above.
kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine(
constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]],
device const char * src0 [[buffer(1)]],
device char * dst [[buffer(4)]],
uint tid [[thread_index_in_threadgroup]],
ushort3 ntg [[threads_per_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
const int i1 = tgpig[0];
const int i2 = tgpig[1];
const int n_nope = args.head_dim - args.n_dims;
if (n_nope < 0) {
return;
}
float corr_dims[2];
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step;
const float theta_base = (float)as_type<int>(raw_pos);
const float inv_ndims = -1.f/args.n_dims;
device const char * src_base =
src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
device char * dst_base =
dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes;
for (int r = tid; r < args.n_dims; r += ntg.x) {
if ((r & 1) != 0) {
continue;
}
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
#else
const float theta = theta_base * pow(args.freq_base, inv_ndims*r);
#endif
const float freq_factor = 1.0f;
float cos_theta;
float sin_theta;
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
if (args.inverse) {
sin_theta = -sin_theta;
}
const int j0 = n_nope + r;
const int j1 = j0 + 1;
const float x0 = *((device const float *) (src_base + j0*sizeof(float)));
const float x1 = *((device const float *) (src_base + j1*sizeof(float)));
*((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta;
*((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta;
}
}
metal/flash_attn.metal
+9 -2
View File
@@ -939,6 +939,7 @@ constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_
constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]];
constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]];
constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]];
constant bool FC_flash_attn_ext_vec_shared_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]];
constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 20)]];
constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]];
constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]];
@@ -1074,8 +1075,14 @@ kernel void kernel_flash_attn_ext_vec(
if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) {
k = pad;
v = k + args.nb11*C*args.ne_12_2*args.ne_12_3;
mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3;
if (FC_flash_attn_ext_vec_shared_kvpad) {
v = k;
mask = k + args.nb11*C*args.ne_12_2*args.ne_12_3 +
args.nb21*C*args.ne_12_2*args.ne_12_3;
} else {
v = k + args.nb11*C*args.ne_12_2*args.ne_12_3;
mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3;
}
const short ikv2 = iq2/(args.ne02/args.ne_12_2);
const short ikv3 = iq3/(args.ne03/args.ne_12_3);
metal/moe.metal
+75 -73
View File
@@ -7478,8 +7478,13 @@ kernel void kernel_mul_mm_id_addr(
}
}
// Fused routed gate+up grouped matmul with the SwiGLU epilogue, generic over
// the expert quant block. Both A streams share one staged B tile per k-step;
// each output keeps the exact MMA accumulation order of the separate GEMMs,
// and the epilogue matches kernel_dsv4_moe_swiglu_weight_f16, so the fused
// result is bit-identical to the unfused path.
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread half4x4 &)>
kernel void kernel_mul_mm_id_pair_swiglu_f16(
kernel void kernel_mul_mm_id_pair_swiglu_f16_impl(
constant ds4_metal_args_mul_mm_id & args,
constant ds4_metal_dsv4_moe_swiglu_weight_args & act,
device const char * src0_gate,
@@ -7494,8 +7499,9 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16(
ushort tiitg[[thread_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
threadgroup half *sa = (threadgroup half *)(shmem);
threadgroup half *sb = (threadgroup half *)(shmem + 4096);
threadgroup half *sa_gate = (threadgroup half *)(shmem);
threadgroup half *sa_up = (threadgroup half *)(shmem + 4096);
threadgroup half *sb = (threadgroup half *)(shmem + 8192);
constexpr int NR0 = 64;
constexpr int NR1 = 32;
@@ -7547,7 +7553,8 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16(
+ args.nb11*i11
+ args.nb10*iy);
simdgroup_half8x8 ma[4];
simdgroup_half8x8 ma_g[4];
simdgroup_half8x8 ma_u[4];
simdgroup_half8x8 mb[2];
simdgroup_float8x8 mc_gate[8];
@@ -7559,83 +7566,47 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16(
}
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
const short sx_b = (tiitg%NL1);
const short sy_b = (tiitg/NL1)/8;
const short ly_b = (tiitg/NL1)%8;
const short ib_b = 4*sx_b + sy_b;
*(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) =
(half2x4)(*((device float2x4 *) y));
half4x4 temp_gate;
dequantize_func(xg, il, temp_gate);
threadgroup_barrier(mem_flags::mem_threadgroup);
FOR_UNROLL (short i = 0; i < 16; i++) {
const short sx = 2*il0 + i/8;
const short sy = (tiitg/NL0)/8;
const short lx = (tiitg/NL0)%8;
const short ly = i%8;
const short ib = 8*sx + sy;
*(sa + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup const half * lsma_gate = (sa + 4*64*(sgitg%2));
threadgroup const half * lsmb = (sb + 2*64*(sgitg/2));
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 4; i++) {
simdgroup_load(ma[i], lsma_gate + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 2; i++) {
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 8; i++) {
simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma[i%4], mc_gate[i]);
}
lsma_gate += 8*64;
lsmb += 4*64;
}
half4x4 temp_up;
dequantize_func(xu, il, temp_up);
threadgroup_barrier(mem_flags::mem_threadgroup);
// Stage the shared B tile and both expert A tiles in one pass: two
// threadgroup barriers per k-step total instead of two per matrix.
{
const short sx_b = (tiitg%NL1);
const short sy_b = (tiitg/NL1)/8;
const short ly_b = (tiitg/NL1)%8;
const short ib_b = 4*sx_b + sy_b;
*(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) =
(half2x4)(*((device float2x4 *) y));
}
FOR_UNROLL (short i = 0; i < 16; i++) {
const short sx = 2*il0 + i/8;
const short sy = (tiitg/NL0)/8;
const short lx = (tiitg/NL0)%8;
const short ly = i%8;
const short ib = 8*sx + sy;
*(sa + 64*ib + 8*ly + lx) = temp_up[i/4][i%4];
*(sa_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4];
*(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup const half * lsma_up = (sa + 4*64*(sgitg%2));
lsmb = (sb + 2*64*(sgitg/2));
threadgroup const half * lsma_gate = (sa_gate + 4*64*(sgitg%2));
threadgroup const half * lsma_up = (sa_up + 4*64*(sgitg%2));
threadgroup const half * lsmb = (sb + 2*64*(sgitg/2));
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 4; i++) {
simdgroup_load(ma[i], lsma_up + 64*i, 8, 0, false);
simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false);
simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false);
}
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 2; i++) {
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
}
@@ -7643,10 +7614,12 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16(
simdgroup_barrier(mem_flags::mem_none);
FOR_UNROLL (short i = 0; i < 8; i++) {
simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma[i%4], mc_up[i]);
simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]);
simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]);
}
lsma_up += 8*64;
lsma_gate += 8*64;
lsma_up += 8*64;
lsmb += 4*64;
}
@@ -7701,11 +7674,14 @@ kernel void kernel_mul_mm_id_pair_swiglu_f16(
}
}
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>) mul_mm_id_pair_swiglu_f16;
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>) mul_mm_id_pair_swiglu_f16_iq2;
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_q4_K, QK_NL, dequantize_q4_K>) mul_mm_id_pair_swiglu_f16_q4;
typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl<block_q5_K, QK_NL, dequantize_q5_K>) mul_mm_id_pair_swiglu_f16_q5;
template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16 kernel_mul_mm_id_pair_swiglu_f16<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16 kernel_mul_mm_id_pair_swiglu_f16<block_q4_K, QK_NL, dequantize_q4_K>;
template [[host_name("kernel_mul_mm_id_q5_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16 kernel_mul_mm_id_pair_swiglu_f16<block_q5_K, QK_NL, dequantize_q5_K>;
// Host-visible fused routed pair matmuls for the DS4 expert quant formats.
template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_iq2 kernel_mul_mm_id_pair_swiglu_f16_impl<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q4 kernel_mul_mm_id_pair_swiglu_f16_impl<block_q4_K, QK_NL, dequantize_q4_K>;
template [[host_name("kernel_mul_mm_id_q5_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q5 kernel_mul_mm_id_pair_swiglu_f16_impl<block_q5_K, QK_NL, dequantize_q5_K>;
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id;
typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs;
@@ -7741,6 +7717,11 @@ template [[host_name("kernel_mul_mm_id_addr_q4_K_f16")]] kernel mul_mm_id_add
// weights to a half tile, then let TensorOps read the dense head activations
// directly. Only the 64-token direct-RHS instantiation is exported because the
// staged-RHS and 32-token variants were benchmark-only experiments.
//
// Full tiles (the host dispatch guarantee for aligned batches) skip all bounds
// work. The weight tile is double-buffered: the next k-step's dequant
// overlaps the current cooperative matmul, so the k-loop needs one
// threadgroup barrier per step instead of two.
template<short NR1>
kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs(
constant ds4_metal_args_mul_mm_id & args,
@@ -7768,7 +7749,8 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs(
const bool full_tile = r0 + NR0 <= M && r1 + NR1 <= N && (K % NK) == 0;
threadgroup half *sa = (threadgroup half *)shmem;
auto tA = tensor(sa, dextents<int32_t, 2>(NK, NR0));
auto tA0 = tensor(sa, dextents<int32_t, 2>(NK, NR0));
auto tA1 = tensor(sa + NR0*NK, dextents<int32_t, 2>(NK, NR0));
device float *ptrB = (device float *)(srcB + args.nb11*group);
const int strideB = args.nb12/sizeof(float);
@@ -7779,7 +7761,7 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs(
matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<4>> mm;
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
auto cT = mm.template get_destination_cooperative_tensor<decltype(tB), decltype(tA0), float>();
#pragma unroll
for (uint16_t i = 0; i < cT.get_capacity(); ++i) {
@@ -7788,37 +7770,57 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs(
}
}
for (int loop_k = 0; loop_k < K; loop_k += NK) {
auto stage_tile = [&](const int loop_k, threadgroup half *buf) {
for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) {
const int row = work/NL;
const int k_chunk = work%NL;
const int k_pos = loop_k + k_chunk*16;
const short k_base = k_chunk*16;
if (full_tile || r0 + row < M) {
if (full_tile) {
device const block_q8_0 *row_ptr =
(device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02);
half4x4 temp_a;
dequantize_q8_0_pairs(row_ptr + k_pos/32, (k_pos/16)%2, temp_a);
threadgroup half4 *dst4 = (threadgroup half4 *)(buf + row*NK + k_base);
dst4[0] = temp_a[0];
dst4[1] = temp_a[1];
dst4[2] = temp_a[2];
dst4[3] = temp_a[3];
} else if (r0 + row < M) {
const int block_idx = k_pos/32;
const short il = (k_pos/16)%2;
device const block_q8_0 *row_ptr =
(device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02);
half4x4 temp_a;
dequantize_q8_0(row_ptr + block_idx, il, temp_a);
dequantize_q8_0_pairs(row_ptr + block_idx, il, temp_a);
FOR_UNROLL (short i = 0; i < 16; i++) {
sa[row*NK + k_base + i] = (full_tile || k_pos + i < K) ? temp_a[i/4][i%4] : (half)0;
buf[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (half)0;
}
} else {
FOR_UNROLL (short i = 0; i < 16; i++) {
sa[row*NK + k_base + i] = (half)0;
buf[row*NK + k_base + i] = (half)0;
}
}
}
};
threadgroup_barrier(mem_flags::mem_threadgroup);
stage_tile(0, sa);
threadgroup_barrier(mem_flags::mem_threadgroup);
auto mA = tA.slice(0, 0);
uint buf_sel = 0;
for (int loop_k = 0; loop_k < K; loop_k += NK) {
auto mA = (buf_sel ? tA1 : tA0).slice(0, 0);
auto mB = tB.slice(loop_k, r1);
mm.run(mB, mA, cT);
const int next_k = loop_k + NK;
if (next_k < K) {
buf_sel ^= 1u;
stage_tile(next_k, buf_sel ? sa + NR0*NK : sa);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
metal/norm.metal
+44 -20
View File
@@ -96,24 +96,21 @@ kernel void kernel_add_rms_norm_mul_f32_4(
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
if (sgitg == 0) {
shmem_f32[tiisg] = 0.0f;
}
ushort3 ntg[[threads_per_threadgroup]]) {
if (sgitg == 0) shmem_f32[tiisg] = 0.0f;
const int i01 = tgpig.x;
const int i02 = tgpig.y;
const int i03 = tgpig.z;
device const float4 * a = (device const float4 *)
device const float4 *a = (device const float4 *)
(src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
device const float4 * b = (device const float4 *)
device const float4 *b = (device const float4 *)
(src1 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
device const float4 * w = (device const float4 *)
device const float4 *w = (device const float4 *)
(weight + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
device float4 * sum = (device float4 *)
device float4 *sum = (device float4 *)
(sum_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
device float4 * norm = (device float4 *)
device float4 *norm = (device float4 *)
(norm_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
float sumf = 0.0f;
@@ -123,26 +120,53 @@ kernel void kernel_add_rms_norm_mul_f32_4(
sumf += dot(v, v);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) {
shmem_f32[sgitg] = sumf;
}
if (tiisg == 0) shmem_f32[sgitg] = sumf;
threadgroup_barrier(mem_flags::mem_threadgroup);
sumf = shmem_f32[tiisg];
sumf = simd_sum(sumf);
sumf = simd_sum(shmem_f32[tiisg]);
const float mean = sumf / args.ne00;
const float scale = 1.0f / sqrt(mean + args.eps);
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
norm[i00] = (sum[i00] * scale) * w[i00];
}
}
// RMSNorm reduction used when the following F16 matmul applies the row scale
// while staging its RHS tile.
kernel void kernel_rms_norm_scale_f32_4(
constant ds4_metal_args_norm & args,
device const char * src0,
device float * dst_scale,
threadgroup float * shmem_f32 [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
if (sgitg == 0) shmem_f32[tiisg] = 0.0f;
const int i01 = tgpig.x;
const int i02 = tgpig.y;
const int i03 = tgpig.z;
device const float4 *x = (device const float4 *)
(src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
float sumf = 0.0f;
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
sumf += dot(x[i00], x[i00]);
}
sumf = simd_sum(sumf);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tiisg == 0) shmem_f32[sgitg] = sumf;
threadgroup_barrier(mem_flags::mem_threadgroup);
sumf = simd_sum(shmem_f32[tiisg]);
const float mean = sumf / args.ne00;
const float scale = 1.0f / sqrt(mean + args.eps);
if (tpitg.x == 0) dst_scale[i01] = scale;
}
struct ds4_metal_args_qkv_rms_norm {
int32_t q_n;
int32_t q_n4;
rocm/ds4_rocm_compressor.cuh
+8 -5
View File
@@ -260,7 +260,8 @@ extern "C" int ds4_gpu_compressor_update_tensor(
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
float rms_eps,
bool state_already_stored) {
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
!model_map || !cuda_compressor_shape_supported(head_dim, ratio) ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
@@ -284,10 +285,12 @@ extern "C" int ds4_gpu_compressor_update_tensor(
(emit && !cuda_tensor_has_bytes(comp_cache, comp_bytes))) {
return 0;
}
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
if (!state_already_stored) {
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
}
}
if (!emit) return 1;
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
rocm/ds4_rocm_matmul.cuh
+34
View File
@@ -729,6 +729,40 @@ extern "C" int ds4_gpu_matmul_f16_pair_tensor(
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
}
extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
(void)out_kv;
(void)out_score;
(void)state_kv;
(void)state_score;
(void)model_map;
(void)model_size;
(void)weight_kv_offset;
(void)weight_score_offset;
(void)ape_offset;
(void)ape_type;
(void)in_dim;
(void)width;
(void)x;
(void)ratio;
(void)pos;
return 0;
}
extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
tests/ds4_test.c
+4275 -3
View File
File diff suppressed because it is too large. Load diff