From fe5ceee235f256c30f18e551f74a77acc2fc934e Mon Sep 17 00:00:00 2001 From: iCreil Date: Tue, 7 Jul 2026 12:46:02 +0200 Subject: [PATCH 01/10] CUDA: add multi-sequence decode attention kernel Add ds4_gpu_attention_decode_heads_multi_tensor(): one query token per sequence, each attending to its own raw + compressed caches with decode (everything visible) semantics. Sequence b reads q row b and writes heads row b; per-sequence cache pointers, ring geometry and optional indexer mask row travel in a ds4_gpu_attn_seqview array (max 4, passed by value in the kernel parameters). The kernel is a standalone adaptation of the n_tokens==1 launch of attention_decode_mixed_kernel, so the single-session hot path is untouched. Sequences whose compressed row count exceeds the shared score buffer are rejected for now: the large-context online variant is a follow-up, and callers gate batched decode on this. tests/attention_multi_smoke drives two sequences with different shapes, ring offsets and masks and checks the multi launch against two independent single-sequence launches: bitwise identical on both the head_dim==512 fast path and the generic path (RTX PRO 6000, CUDA 13). --- Makefile | 6 ++ ds4_cuda.cu | 175 ++++++++++++++++++++++++++++++++++ ds4_gpu.h | 27 ++++++ tests/attention_multi_smoke.c | 161 +++++++++++++++++++++++++++++++ 4 files changed, 369 insertions(+) create mode 100644 tests/attention_multi_smoke.c diff --git a/Makefile b/Makefile index 81d4881b2..216edfe55 100644 --- a/Makefile +++ b/Makefile @@ -344,6 +344,12 @@ test-cuda-mixed-batch: tests/test_cuda_mixed_batch DS4_TEST_MODEL="$(DS4_TEST_MODEL)" ./tests/test_cuda_mixed_batch endif +tests/attention_multi_smoke.o: tests/attention_multi_smoke.c ds4_gpu.h + $(CC) $(CFLAGS) -I. -c -o $@ tests/attention_multi_smoke.c + +tests/attention_multi_smoke: tests/attention_multi_smoke.o ds4_cuda.o + $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) + ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS) diff --git a/ds4_cuda.cu b/ds4_cuda.cu index aaa4df113..54c55656e 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -9450,6 +9450,123 @@ __global__ static void attention_static_mixed_heads8_online_kernel( } } +/* Multi-sequence decode attention. One block column per sequence: sequence b + * evaluates its single query token against its own raw + compressed caches + * with everything visible, exactly like the n_tokens==1 launch of + * attention_decode_mixed_kernel, but each sequence reads from its own + * cache pointers. Kept as a standalone kernel so the single-session hot + * path stays untouched. */ +typedef struct { + const float *raw_kv; + const float *comp_kv; + const float *comp_mask; /* per-sequence mask row, NULL = none */ + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; +} cuda_attn_seqview; + +typedef struct { + cuda_attn_seqview v[DS4_GPU_DECODE_MULTI_MAX]; +} cuda_attn_seqviews; + +__global__ static void attention_decode_multi_kernel( + float *heads, + const float *sinks, + const float *q, + cuda_attn_seqviews views, + uint32_t n_seqs, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t t = blockIdx.x; + const uint32_t h = blockIdx.y; + if (t >= n_seqs || h >= n_head) return; + const cuda_attn_seqview view = views.v[t]; + const uint32_t raw_count = view.n_raw > 256u ? 256u : view.n_raw; + const uint32_t visible_comp = view.n_comp; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + __shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP]; + __shared__ uint32_t raw_rows[256]; + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + const float scale = rsqrtf((float)head_dim); + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + raw_rows[r] = (view.raw_start + r) % view.raw_cap; + } + __syncthreads(); + const uint32_t n_score = raw_count + visible_comp; + float local_max = sinks[h]; + for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) { + const float *kvrow = view.raw_kv + (uint64_t)raw_rows[r] * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + scores[r] = dot * scale; + local_max = fmaxf(local_max, scores[r]); + } + for (uint32_t cidx = threadIdx.x; cidx < visible_comp; cidx += blockDim.x) { + float add = view.comp_mask ? view.comp_mask[cidx] : 0.0f; + float s = -INFINITY; + if (add > -1.0e20f) { + const float *kvrow = view.comp_kv + (uint64_t)cidx * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d]; + s = dot * scale + add; + } + scores[raw_count + cidx] = s; + local_max = fmaxf(local_max, s); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]); + __syncthreads(); + } + if (threadIdx.x == 0) max_s = partial[0]; + __syncthreads(); + float den_local = 0.0f; + for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) { + scores[i] = expf(scores[i] - max_s); + den_local += scores[i]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s); + __syncthreads(); + float *oh = heads + ((uint64_t)t * n_head + h) * head_dim; + if (head_dim == 512u && blockDim.x == 256u) { + const uint32_t d0 = threadIdx.x; + const uint32_t d1 = d0 + 256u; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) { + const float s = scores[r]; + const float *kv = view.raw_kv + (uint64_t)raw_rows[r] * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + for (uint32_t cidx = 0; cidx < visible_comp; cidx++) { + const float s = scores[raw_count + cidx]; + const float *kv = view.comp_kv + (uint64_t)cidx * head_dim; + acc0 += kv[d0] * s; + acc1 += kv[d1] * s; + } + oh[d0] = acc0 / denom; + oh[d1] = acc1 / denom; + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < raw_count; r++) acc += view.raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r]; + for (uint32_t cidx = 0; cidx < visible_comp; cidx++) acc += view.comp_kv[(uint64_t)cidx * head_dim + d] * scores[raw_count + cidx]; + oh[d] = acc / denom; + } + } +} + __global__ static void attention_decode_mixed_heads8_online_kernel( float *heads, const float *sinks, @@ -14867,6 +14984,64 @@ extern "C" int ds4_gpu_attention_decode_rows_rope_tensor( "attention decode rows inverse rope launch"); } +extern "C" int ds4_gpu_attention_decode_heads_multi_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_attn_seqview *seqs, + uint32_t n_seqs, + uint32_t comp_kv_f16, + uint32_t n_head, + uint32_t head_dim) { + if (comp_kv_f16 || !heads || !q || !seqs || !model_map || + n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_seqs * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_seqs * n_head * head_dim * sizeof(float)) { + return 0; + } + cuda_attn_seqviews views; + memset(&views, 0, sizeof(views)); + for (uint32_t b = 0; b < n_seqs; b++) { + const ds4_gpu_attn_seqview *sv = &seqs[b]; + if (!sv->raw_kv || sv->n_raw == 0 || sv->raw_cap < sv->n_raw || + sv->raw_start >= sv->raw_cap || + (sv->n_comp != 0 && !sv->comp_kv) || + sv->raw_kv->bytes < (uint64_t)sv->raw_cap * head_dim * sizeof(float) || + (sv->n_comp && sv->comp_kv->bytes < (uint64_t)sv->n_comp * head_dim * sizeof(float)) || + (sv->comp_mask && sv->comp_mask->bytes < (uint64_t)sv->n_comp * sizeof(float))) { + return 0; + } + if (!cuda_attention_score_buffer_fits(sv->n_comp)) { + /* The large-context online variant is not implemented for the + * multi path yet; callers gate batched decode on this. */ + return 0; + } + views.v[b].raw_kv = (const float *)sv->raw_kv->ptr; + views.v[b].comp_kv = sv->n_comp ? (const float *)sv->comp_kv->ptr : (const float *)sv->raw_kv->ptr; + views.v[b].comp_mask = sv->comp_mask ? (const float *)sv->comp_mask->ptr : NULL; + views.v[b].n_raw = sv->n_raw; + views.v[b].raw_cap = sv->raw_cap; + views.v[b].raw_start = sv->raw_start; + views.v[b].n_comp = sv->n_comp; + } + const float *sinks = (const float *)cuda_model_range_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), "attn_sinks"); + if (!sinks) return 0; + dim3 grid(n_seqs, n_head, 1); + attention_decode_multi_kernel<<>>((float *)heads->ptr, + sinks, + (const float *)q->ptr, + views, + n_seqs, + n_head, + head_dim); + return cuda_ok(cudaGetLastError(), "attention decode multi launch"); +} + extern "C" int ds4_gpu_attention_prefill_raw_heads_tensor(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, uint64_t sinks_offset, const ds4_gpu_tensor *q, const ds4_gpu_tensor *raw_kv, uint32_t n_tokens, uint32_t window, uint32_t n_head, uint32_t head_dim) { if (!heads || !q || !raw_kv || !model_map || sinks_offset > model_size || model_size - sinks_offset < (uint64_t)n_head * sizeof(float) || diff --git a/ds4_gpu.h b/ds4_gpu.h index 2000bba8b..a3b14859d 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -1755,6 +1755,33 @@ int ds4_gpu_flash_kv_stage_f16_tensor( uint32_t n_comp, uint32_t head_dim); +/* Batched multi-sequence decode attention: one query token per sequence, + * each attending to its own KV caches (all rows visible, decode semantics). + * Slot b reads q row b and writes heads row b. */ +#define DS4_GPU_DECODE_MULTI_MAX 4u + +typedef struct { + const ds4_gpu_tensor *raw_kv; + const ds4_gpu_tensor *comp_kv; /* NULL when the layer keeps no compressed cache */ + const ds4_gpu_tensor *comp_mask; /* per-sequence mask row (n_comp floats), NULL = none */ + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; +} ds4_gpu_attn_seqview; + +int ds4_gpu_attention_decode_heads_multi_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_attn_seqview *seqs, + uint32_t n_seqs, + uint32_t comp_kv_f16, + uint32_t n_head, + uint32_t head_dim); + int ds4_gpu_attention_prefill_raw_heads_tensor( ds4_gpu_tensor *heads, const void *model_map, diff --git a/tests/attention_multi_smoke.c b/tests/attention_multi_smoke.c new file mode 100644 index 000000000..511a1f54b --- /dev/null +++ b/tests/attention_multi_smoke.c @@ -0,0 +1,161 @@ +/* attention_multi_smoke.c — equivalence check for the multi-sequence decode + * attention kernel: two sequences with different raw/compressed caches must + * produce, in one ds4_gpu_attention_decode_heads_multi_tensor() call, exactly + * the heads that two independent ds4_gpu_attention_decode_heads_tensor() + * calls produce. + * + * Usage: tests/attention_multi_smoke + */ +#include "ds4_gpu.h" + +#include +#include +#include +#include +#include + +#define N_HEAD 8u +#define N_SEQS 2u + +static uint32_t lcg_state = 0x12345678u; +static float frand(void) { + lcg_state = lcg_state * 1664525u + 1013904223u; + return ((float)(lcg_state >> 8) / (float)(1u << 24)) - 0.5f; +} + +static void fill(float *p, uint64_t n) { + for (uint64_t i = 0; i < n; i++) p[i] = frand(); +} + +static ds4_gpu_tensor *upload(const float *host, uint64_t n) { + ds4_gpu_tensor *t = ds4_gpu_tensor_alloc(n * sizeof(float)); + if (!t || !ds4_gpu_tensor_write(t, 0, host, n * sizeof(float))) { + fprintf(stderr, "attention_multi_smoke: tensor upload failed\n"); + exit(1); + } + return t; +} + +static int run_case(uint32_t head_dim, const void *map, uint64_t map_size) { + /* Two sequences with deliberately different shapes and ring offsets. */ + const uint32_t n_raw[N_SEQS] = {96u, 41u}; + const uint32_t raw_cap[N_SEQS] = {128u, 64u}; + const uint32_t raw_start[N_SEQS] = {37u, 5u}; + const uint32_t n_comp[N_SEQS] = {200u, 0u}; + + float *q_host = malloc((size_t)N_SEQS * N_HEAD * head_dim * sizeof(float)); + fill(q_host, (uint64_t)N_SEQS * N_HEAD * head_dim); + ds4_gpu_tensor *q = upload(q_host, (uint64_t)N_SEQS * N_HEAD * head_dim); + + ds4_gpu_tensor *raw[N_SEQS]; + ds4_gpu_tensor *comp[N_SEQS] = {NULL, NULL}; + ds4_gpu_tensor *mask[N_SEQS] = {NULL, NULL}; + for (uint32_t b = 0; b < N_SEQS; b++) { + float *raw_host = malloc((size_t)raw_cap[b] * head_dim * sizeof(float)); + fill(raw_host, (uint64_t)raw_cap[b] * head_dim); + raw[b] = upload(raw_host, (uint64_t)raw_cap[b] * head_dim); + free(raw_host); + if (n_comp[b]) { + float *comp_host = malloc((size_t)n_comp[b] * head_dim * sizeof(float)); + fill(comp_host, (uint64_t)n_comp[b] * head_dim); + comp[b] = upload(comp_host, (uint64_t)n_comp[b] * head_dim); + free(comp_host); + float *mask_host = malloc((size_t)n_comp[b] * sizeof(float)); + for (uint32_t i = 0; i < n_comp[b]; i++) { + mask_host[i] = (i % 7u == 3u) ? -1.0e30f : frand() * 0.1f; + } + mask[b] = upload(mask_host, n_comp[b]); + free(mask_host); + } + } + + const uint64_t head_row = (uint64_t)N_HEAD * head_dim; + ds4_gpu_tensor *heads_ref = ds4_gpu_tensor_alloc((uint64_t)N_SEQS * head_row * sizeof(float)); + ds4_gpu_tensor *heads_multi = ds4_gpu_tensor_alloc((uint64_t)N_SEQS * head_row * sizeof(float)); + if (!heads_ref || !heads_multi) { + fprintf(stderr, "attention_multi_smoke: heads alloc failed\n"); + return 1; + } + + /* Reference: one single-sequence call per slot through row views. */ + for (uint32_t b = 0; b < N_SEQS; b++) { + ds4_gpu_tensor *hv = ds4_gpu_tensor_view(heads_ref, (uint64_t)b * head_row * sizeof(float), + head_row * sizeof(float)); + ds4_gpu_tensor *qv = ds4_gpu_tensor_view(q, (uint64_t)b * head_row * sizeof(float), + head_row * sizeof(float)); + if (!hv || !qv || + !ds4_gpu_attention_decode_heads_tensor(hv, map, map_size, 0, qv, + raw[b], n_raw[b], raw_cap[b], raw_start[b], + comp[b], 0, n_comp[b], + mask[b], mask[b] != NULL, + N_HEAD, head_dim)) { + fprintf(stderr, "attention_multi_smoke: single call failed (seq %u, head_dim %u)\n", + b, head_dim); + return 1; + } + } + + /* Multi: one call for both sequences. */ + ds4_gpu_attn_seqview seqs[N_SEQS]; + memset(seqs, 0, sizeof(seqs)); + for (uint32_t b = 0; b < N_SEQS; b++) { + seqs[b].raw_kv = raw[b]; + seqs[b].comp_kv = comp[b]; + seqs[b].comp_mask = mask[b]; + seqs[b].n_raw = n_raw[b]; + seqs[b].raw_cap = raw_cap[b]; + seqs[b].raw_start = raw_start[b]; + seqs[b].n_comp = n_comp[b]; + } + if (!ds4_gpu_attention_decode_heads_multi_tensor(heads_multi, map, map_size, 0, + q, seqs, N_SEQS, 0, + N_HEAD, head_dim)) { + fprintf(stderr, "attention_multi_smoke: multi call failed (head_dim %u)\n", head_dim); + return 1; + } + if (!ds4_gpu_synchronize()) { + fprintf(stderr, "attention_multi_smoke: synchronize failed\n"); + return 1; + } + + const uint64_t n_out = (uint64_t)N_SEQS * head_row; + float *ref_host = malloc(n_out * sizeof(float)); + float *multi_host = malloc(n_out * sizeof(float)); + if (!ds4_gpu_tensor_read(heads_ref, 0, ref_host, n_out * sizeof(float)) || + !ds4_gpu_tensor_read(heads_multi, 0, multi_host, n_out * sizeof(float))) { + fprintf(stderr, "attention_multi_smoke: readback failed\n"); + return 1; + } + uint64_t mismatches = 0; + float max_abs = 0.0f; + for (uint64_t i = 0; i < n_out; i++) { + const float d = fabsf(ref_host[i] - multi_host[i]); + if (d > max_abs) max_abs = d; + if (d > 1e-6f) mismatches++; + } + printf("head_dim=%u: %llu/%llu values differ (max abs diff %g) -> %s\n", + head_dim, + (unsigned long long)mismatches, + (unsigned long long)n_out, + (double)max_abs, + mismatches == 0 ? "MATCH" : "MISMATCH"); + free(ref_host); + free(multi_host); + free(q_host); + return mismatches == 0 ? 0 : 1; +} + +int main(void) { + /* Fake "model" that only carries the attention sinks at offset 0. */ + float sinks[N_HEAD]; + for (uint32_t h = 0; h < N_HEAD; h++) sinks[h] = frand(); + if (!ds4_gpu_set_model_map(sinks, sizeof(sinks))) { + fprintf(stderr, "attention_multi_smoke: set_model_map failed\n"); + return 1; + } + + int rc = run_case(512u, sinks, sizeof(sinks)); + rc |= run_case(128u, sinks, sizeof(sinks)); + printf(rc == 0 ? "OK\n" : "FAILED\n"); + return rc; +} From f43c57a4b97ad8b1a916be488c75c09e00a57720 Mon Sep 17 00:00:00 2001 From: iCreil Date: Tue, 7 Jul 2026 13:21:14 +0200 Subject: [PATCH 02/10] CUDA: add pos-vector RoPE tail variants for batched decode Batched decode evaluates one token per sequence, each at its own arbitrary position, while rope_tail_kernel and head_rms_norm_rope_tail_kernel derive the position as pos0 + t (* stride). Add _multi twins that take a small by-value position array (max DS4_GPU_DECODE_MULTI_MAX entries) indexed by token row, with per-element arithmetic kept identical so each row matches a single-token launch bitwise. tests/rope_multi_bench checks that equivalence (plain and YaRN paths) and measures why the variants exist: on an RTX PRO 6000, per-sequence loops of single-token launches over the 4 decode RoPE sites cost ~1.0 ms extra per 43-layer step at B=4 (~516 launches x ~2 us) versus batched launches. --- Makefile | 6 ++ ds4_cuda.cu | 151 +++++++++++++++++++++++++++ ds4_gpu.h | 36 +++++++ tests/rope_multi_bench.c | 215 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 408 insertions(+) create mode 100644 tests/rope_multi_bench.c diff --git a/Makefile b/Makefile index 216edfe55..3c991fd41 100644 --- a/Makefile +++ b/Makefile @@ -350,6 +350,12 @@ tests/attention_multi_smoke.o: tests/attention_multi_smoke.c ds4_gpu.h tests/attention_multi_smoke: tests/attention_multi_smoke.o ds4_cuda.o $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) +tests/rope_multi_bench.o: tests/rope_multi_bench.c ds4_gpu.h + $(CC) $(CFLAGS) -I. -c -o $@ tests/rope_multi_bench.c + +tests/rope_multi_bench: tests/rope_multi_bench.o ds4_cuda.o + $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) + ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS) diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 54c55656e..53a497a98 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -6058,6 +6058,136 @@ __global__ static void rope_tail_decode_rows_kernel( tail[i + 1u] = x0 * s + x1 * c; } +/* pos-vector twins of the two RoPE kernels above, for batched decode where + * each token belongs to a different sequence at an arbitrary position. The + * per-element arithmetic is kept identical so token t matches a single-token + * launch with pos0 = pos.v[t] bitwise. */ +typedef struct { + uint32_t v[DS4_GPU_DECODE_MULTI_MAX]; +} cuda_rope_pos; + +__global__ static void head_rms_norm_rope_tail_multi_kernel( + float *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + cuda_rope_pos pos, + uint32_t n_ctx_orig, + int inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps) { + uint32_t row = blockIdx.x; + if (row >= n_tok * n_head) return; + uint32_t t = row / n_head; + float *xr = x + (uint64_t)row * head_dim; + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) { + float v = xr[i]; + sum += v * v; + } + __shared__ float partial[256]; + partial[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride]; + __syncthreads(); + } + const float scale = rsqrtf(partial[0] / (float)head_dim + eps); + const uint32_t n_nope = head_dim - n_rot; + for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) { + xr[i] *= scale; + } + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + float denom = 2.0f * logf(freq_base); + corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom); + corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom); + corr0 = fmaxf(0.0f, corr0); + corr1 = fminf((float)(n_rot - 1), corr1); + } + for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) { + uint32_t i = pair * 2u; + float theta_extrap = (float)pos.v[t] * powf(freq_base, -((float)i) / (float)n_rot); + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = attn_factor; + if (ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + float c = cosf(theta) * mscale; + float s = sinf(theta) * mscale; + if (inverse) s = -s; + float *tail = xr + n_nope; + float x0 = tail[i] * scale; + float x1 = tail[i + 1] * scale; + tail[i] = x0 * c - x1 * s; + tail[i + 1] = x0 * s + x1 * c; + } +} + +__global__ static void rope_tail_multi_kernel( + float *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + cuda_rope_pos pos, + uint32_t n_ctx_orig, + int inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t pairs = n_tok * n_head * (n_rot / 2); + if (gid >= pairs) return; + uint32_t pair = gid % (n_rot / 2); + uint32_t tmp = gid / (n_rot / 2); + uint32_t h = tmp % n_head; + uint32_t t = tmp / n_head; + uint32_t n_nope = head_dim - n_rot; + uint32_t i = pair * 2; + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + float denom = 2.0f * logf(freq_base); + corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom); + corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom); + corr0 = fmaxf(0.0f, corr0); + corr1 = fminf((float)(n_rot - 1), corr1); + } + + float theta_extrap = (float)pos.v[t] * powf(freq_base, -((float)i) / (float)n_rot); + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = attn_factor; + if (ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + float c = cosf(theta) * mscale; + float s = sinf(theta) * mscale; + if (inverse) s = -s; + + float *tail = x + ((uint64_t)t * n_head + h) * head_dim + n_nope; + float x0 = tail[i]; + float x1 = tail[i + 1]; + tail[i] = x0 * c - x1 * s; + tail[i + 1] = x0 * s + x1 * c; +} + __device__ static float dsv4_e4m3fn_value_dev(int i) { int exp = (i >> 3) & 15; int mant = i & 7; @@ -13989,6 +14119,27 @@ extern "C" int ds4_gpu_rope_tail_decode_rows_tensor( ext_factor, attn_factor, beta_fast, beta_slow); return cuda_ok(cudaGetLastError(), "rope tail decode rows launch"); } +extern "C" int ds4_gpu_head_rms_norm_rope_tail_multi_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, const uint32_t *pos, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, float eps) { + if (!x || !pos || n_tok == 0 || n_tok > DS4_GPU_DECODE_MULTI_MAX || + n_rot > head_dim || (n_rot & 1u) || + x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; + cuda_rope_pos p; + memset(&p, 0, sizeof(p)); + for (uint32_t t = 0; t < n_tok; t++) p.v[t] = pos[t]; + head_rms_norm_rope_tail_multi_kernel<<>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, p, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, eps); + return cuda_ok(cudaGetLastError(), "head_rms_norm_rope_tail_multi launch"); +} +extern "C" int ds4_gpu_rope_tail_multi_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, const uint32_t *pos, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) { + if (!x || !pos || n_tok == 0 || n_tok > DS4_GPU_DECODE_MULTI_MAX || + n_rot > head_dim || (n_rot & 1u) || + x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0; + cuda_rope_pos p; + memset(&p, 0, sizeof(p)); + for (uint32_t t = 0; t < n_tok; t++) p.v[t] = pos[t]; + uint32_t pairs = n_tok * n_head * (n_rot / 2); + rope_tail_multi_kernel<<<(pairs + 255) / 256, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, p, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "rope_tail_multi launch"); +} extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim); extern "C" int ds4_gpu_kv_fp8_store_raw_tensor( ds4_gpu_tensor *kv, diff --git a/ds4_gpu.h b/ds4_gpu.h index a3b14859d..15f74c363 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -1519,6 +1519,42 @@ int ds4_gpu_glm_attention_flash_tensor( uint32_t value_dim, bool cache_f16); +/* pos-vector variants for batched decode: token t uses pos[t] instead of + * pos0 + t. n_tok is capped at DS4_GPU_DECODE_MULTI_MAX (positions travel + * by value in the kernel parameters). */ +int ds4_gpu_head_rms_norm_rope_tail_multi_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + const uint32_t *pos, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + +int ds4_gpu_rope_tail_multi_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + const uint32_t *pos, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + /* Release decode fused KV finalizer: after the standalone RoPE kernel, this * performs DS4's FP8 non-RoPE KV round trip and writes the F16-rounded raw * attention cache row in one dispatch. */ diff --git a/tests/rope_multi_bench.c b/tests/rope_multi_bench.c new file mode 100644 index 000000000..46b2559b9 --- /dev/null +++ b/tests/rope_multi_bench.c @@ -0,0 +1,215 @@ +/* rope_multi_bench.c — correctness + cost check for the pos-vector RoPE + * variants used by batched decode. + * + * Correctness: for every decode-layer RoPE call site, one + * ds4_gpu_*_multi_tensor() launch over B token rows must match, bitwise, B + * single-token launches on per-row views (the arithmetic per element is the + * same, so exact equality is required, not a tolerance). Both the plain and + * the YaRN (ext_factor != 0) paths are checked. + * + * Cost: times a full simulated decode step (43 layers x 4 RoPE sites) three + * ways — today's consecutive-pos batch kernel (lower bound), the pos-vector + * kernel, and a per-sequence loop of single-token launches (what batched + * decode would pay without the new kernels). + * + * Usage: tests/rope_multi_bench + */ +#include "ds4_gpu.h" + +#include +#include +#include +#include +#include + +#define N_SEQS 4u +#define N_LAYER 43u +#define WARMUP 20u +#define ITERS 400u + +/* DeepSeek V4 Flash decode-layer RoPE call sites (per token). */ +typedef struct { + const char *name; + uint32_t n_head; + uint32_t head_dim; + uint32_t n_rot; + int inverse; + int fused_rms; /* q path uses head_rms_norm_rope_tail */ +} rope_site; + +static const rope_site SITES[] = { + {"q_fused", 64u, 512u, 64u, 0, 1}, + {"kv", 1u, 512u, 64u, 0, 0}, + {"heads_inv", 64u, 512u, 64u, 1, 0}, + {"indexer_q", 64u, 128u, 64u, 0, 0}, +}; +#define N_SITES (sizeof(SITES) / sizeof(SITES[0])) + +static const uint32_t POS[N_SEQS] = {1000u, 4213u, 87u, 15991u}; + +static double now_sec(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; +} + +static uint32_t lcg_state = 0x1234567u; +static float frand(void) { + lcg_state = lcg_state * 1664525u + 1013904223u; + return ((float)(lcg_state >> 8) / (float)(1u << 24)) - 0.5f; +} + +typedef enum { LAUNCH_BATCH, LAUNCH_MULTI, LAUNCH_LOOP } launch_mode; + +static int launch_site(const rope_site *s, launch_mode mode, float ext_factor, + ds4_gpu_tensor *buf, ds4_gpu_tensor *views[N_SEQS]) { + const uint32_t n_ctx_orig = ext_factor != 0.0f ? 4096u : 0u; + const float freq_scale = ext_factor != 0.0f ? 0.25f : 1.0f; + switch (mode) { + case LAUNCH_BATCH: + if (s->fused_rms) { + return ds4_gpu_head_rms_norm_rope_tail_tensor(buf, N_SEQS, s->n_head, s->head_dim, + s->n_rot, POS[0], n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, + 32.0f, 1.0f, 1e-6f); + } + return ds4_gpu_rope_tail_tensor(buf, N_SEQS, s->n_head, s->head_dim, + s->n_rot, POS[0], n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, 32.0f, 1.0f); + case LAUNCH_MULTI: + if (s->fused_rms) { + return ds4_gpu_head_rms_norm_rope_tail_multi_tensor(buf, N_SEQS, s->n_head, s->head_dim, + s->n_rot, POS, n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, + 32.0f, 1.0f, 1e-6f); + } + return ds4_gpu_rope_tail_multi_tensor(buf, N_SEQS, s->n_head, s->head_dim, + s->n_rot, POS, n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, 32.0f, 1.0f); + case LAUNCH_LOOP: + for (uint32_t b = 0; b < N_SEQS; b++) { + int ok; + if (s->fused_rms) { + ok = ds4_gpu_head_rms_norm_rope_tail_tensor(views[b], 1u, s->n_head, s->head_dim, + s->n_rot, POS[b], n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, + 32.0f, 1.0f, 1e-6f); + } else { + ok = ds4_gpu_rope_tail_tensor(views[b], 1u, s->n_head, s->head_dim, + s->n_rot, POS[b], n_ctx_orig, s->inverse != 0, + 10000.0f, freq_scale, ext_factor, 1.0f, 32.0f, 1.0f); + } + if (!ok) return 0; + } + return 1; + } + return 0; +} + +static int check_site(const rope_site *s, float ext_factor, + ds4_gpu_tensor *buf, ds4_gpu_tensor *views[N_SEQS]) { + const uint64_t n = (uint64_t)N_SEQS * s->n_head * s->head_dim; + float *host = malloc(n * sizeof(float)); + float *ref = malloc(n * sizeof(float)); + float *got = malloc(n * sizeof(float)); + if (!host || !ref || !got) return 0; + for (uint64_t i = 0; i < n; i++) host[i] = frand(); + + int ok = ds4_gpu_tensor_write(buf, 0, host, n * sizeof(float)) && + launch_site(s, LAUNCH_LOOP, ext_factor, buf, views) && + ds4_gpu_synchronize() && + ds4_gpu_tensor_read(buf, 0, ref, n * sizeof(float)); + if (ok) { + ok = ds4_gpu_tensor_write(buf, 0, host, n * sizeof(float)) && + launch_site(s, LAUNCH_MULTI, ext_factor, buf, views) && + ds4_gpu_synchronize() && + ds4_gpu_tensor_read(buf, 0, got, n * sizeof(float)); + } + if (!ok) { + fprintf(stderr, "rope_multi_bench: launch failed (%s)\n", s->name); + return 0; + } + uint64_t bad = 0; + for (uint64_t i = 0; i < n; i++) { + if (memcmp(&ref[i], &got[i], sizeof(float)) != 0) bad++; + } + printf("%-10s ext=%g: %llu/%llu values differ -> %s\n", + s->name, (double)ext_factor, + (unsigned long long)bad, (unsigned long long)n, + bad == 0 ? "MATCH" : "MISMATCH"); + free(host); + free(ref); + free(got); + return bad == 0; +} + +/* One simulated decode step: every layer runs every RoPE site. */ +static int step(launch_mode mode, ds4_gpu_tensor *bufs[N_SITES], + ds4_gpu_tensor *views[N_SITES][N_SEQS]) { + for (uint32_t il = 0; il < N_LAYER; il++) { + for (uint32_t si = 0; si < N_SITES; si++) { + if (!launch_site(&SITES[si], mode, 0.0f, bufs[si], views[si])) return 0; + } + } + return ds4_gpu_synchronize(); +} + +static int run_variant(launch_mode mode, const char *label, + ds4_gpu_tensor *bufs[N_SITES], + ds4_gpu_tensor *views[N_SITES][N_SEQS], + double *out_ms_per_step) { + for (uint32_t i = 0; i < WARMUP; i++) { + if (!step(mode, bufs, views)) return 0; + } + const double t0 = now_sec(); + for (uint32_t i = 0; i < ITERS; i++) { + if (!step(mode, bufs, views)) return 0; + } + const double ms = (now_sec() - t0) * 1000.0 / (double)ITERS; + printf("%-30s %8.3f ms/step (%u layers x %u sites)\n", + label, ms, N_LAYER, (unsigned)N_SITES); + *out_ms_per_step = ms; + return 1; +} + +int main(void) { + ds4_gpu_tensor *bufs[N_SITES]; + ds4_gpu_tensor *views[N_SITES][N_SEQS]; + for (uint32_t si = 0; si < N_SITES; si++) { + const uint64_t row = (uint64_t)SITES[si].n_head * SITES[si].head_dim * sizeof(float); + bufs[si] = ds4_gpu_tensor_alloc((uint64_t)N_SEQS * row); + if (!bufs[si]) { + fprintf(stderr, "rope_multi_bench: alloc failed\n"); + return 1; + } + for (uint32_t b = 0; b < N_SEQS; b++) { + views[si][b] = ds4_gpu_tensor_view(bufs[si], (uint64_t)b * row, row); + if (!views[si][b]) { + fprintf(stderr, "rope_multi_bench: view failed\n"); + return 1; + } + } + } + + int rc = 0; + for (uint32_t si = 0; si < N_SITES; si++) { + if (!check_site(&SITES[si], 0.0f, bufs[si], views[si])) rc = 1; + if (!check_site(&SITES[si], 1.0f, bufs[si], views[si])) rc = 1; + } + if (rc != 0) { + printf("FAILED\n"); + return rc; + } + + double ms_batch = 0.0, ms_multi = 0.0, ms_loop = 0.0; + if (!run_variant(LAUNCH_BATCH, "consecutive batch (baseline)", bufs, views, &ms_batch) || + !run_variant(LAUNCH_MULTI, "pos-vector multi", bufs, views, &ms_multi) || + !run_variant(LAUNCH_LOOP, "per-seq loop (B=4 views)", bufs, views, &ms_loop)) { + fprintf(stderr, "rope_multi_bench: bench launch failed\n"); + return 1; + } + printf("\nper-seq loop overhead vs pos-vector: %.3f ms/step at B=%u\n", + ms_loop - ms_multi, N_SEQS); + printf("OK\n"); + return 0; +} From e3d8137d997cdc6bcb41107e33e1493255888042 Mon Sep 17 00:00:00 2001 From: iCreil Date: Wed, 8 Jul 2026 11:42:30 +0200 Subject: [PATCH 03/10] CUDA: narrow-batch q8_0 matmul for 2-4 token decode The batch warp kernel puts n_tok on the grid Y axis, so every token re-reads the same weight row from VRAM; in decode, where weights are the whole cost, that scales the time with n_tok for no benefit. The generic per-row kernel (used when blocks>32) is worse still. The new narrow kernel keeps one warp per output row and, for each 32-weight block it reads once, dots it against all n_tok activation blocks, so the weight bytes cross memory once and are reused from L1. Output layout matches the batch kernel; the n_tok==1 and prefill paths are untouched. Dispatched for n_tok in [2, DS4_CUDA_NARROW_MAX] ahead of cuBLAS and the batch warp kernel; DS4_CUDA_NO_Q8_NARROW forces the old path. Modeled on the n_tok==1 warp8 kernel, it loops blocks with a lane stride and handles any in_dim. Decode-step bench (RTX PRO 6000, V4 Flash q2), aggregate vs interleaving: B=2 1.47x, B=3 1.67x, B=4 1.83x (was B=2 0.62x through the batch path). Correctness: batched_decode_smoke shows narrow and batch-warp produce identical token streams. Batched decode still differs slightly from single decode, but that is the existing non-deterministic MoE reduction (atomicAdd over experts), not the kernel: verified by first-step logits being bit-identical between narrow and batch-warp. --- ds4_cuda.cu | 68 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 53a497a98..37538a081 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -5467,6 +5467,51 @@ static int cuda_q8_mma_try_launch( return cuda_ok(cudaGetLastError(), "matmul_q8_0 mma launch") ? 1 : -1; } +/* Narrow-batch q8_0 matmul for 2-4 tokens (batched decode). + * + * The batch kernel above puts n_tok on the grid Y axis, so every token + * re-reads the same weight row from VRAM independently: at n_tok=2 the + * weights stream through memory twice for no reason. In decode the weights + * are the whole cost, so that doubles the time. This kernel keeps the grid + * one-dimensional (one warp per output row) and, for each 32-weight block it + * reads once, dots it against all n_tok activation blocks — the weight bytes + * cross memory once and are reused from L1 for the remaining tokens. Output + * layout matches the batch kernel: out[tok * out_dim + row]. */ +#define DS4_CUDA_NARROW_MAX 4u +__global__ static void matmul_q8_0_preq_narrow_warp8_kernel( + float *out, + const unsigned char *w, + const int8_t *xq, + const float *xscale, + uint64_t in_dim, + uint64_t out_dim, + uint32_t n_tok, + uint64_t blocks, + int use_dp4a) { + const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u); + const uint32_t lane = threadIdx.x & 31u; + if (row >= out_dim) return; + const unsigned char *wr = w + row * blocks * 34; + float acc[DS4_CUDA_NARROW_MAX]; +#pragma unroll + for (uint32_t t = 0; t < DS4_CUDA_NARROW_MAX; t++) acc[t] = 0.0f; + for (uint64_t b = lane; b < blocks; b += 32u) { + const uint64_t i0 = b * 32; + const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32; + const float wscale = __half2float(*(const __half *)(wr + b * 34)); + const int8_t *qs = (const int8_t *)(wr + b * 34 + 2); + for (uint32_t t = 0; t < n_tok; t++) { + const int8_t *xqb = xq + (uint64_t)t * blocks * 32 + b * 32; + const float xs = xscale[(uint64_t)t * blocks + b]; + const int dot = dot_i8_block(qs, xqb, bn, use_dp4a); + acc[t] += wscale * xs * (float)dot; + } + } + for (uint32_t t = 0; t < n_tok; t++) { + const float a = warp_sum_f32(acc[t]); + if (lane == 0) out[(uint64_t)t * out_dim + row] = a; + } +} __global__ static void dequant_q8_0_to_f16_kernel( __half *out, @@ -12344,7 +12389,15 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode ? g_gpu[logical_tier].device_id : 0; const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "q8_0"); if (!wptr) return 0; - if (g_cublas_ready && n_tok > 1) { + /* 2-4 tokens (batched decode) go through the narrow kernel, which reads + * each weight row once and reuses it across tokens. It beats both cuBLAS + * (GEMM setup dwarfs the work at tiny n_tok) and the batch warp kernel + * (which re-reads weights per token), so it takes priority here. Like the + * n_tok==1 warp8 kernel it loops blocks with a lane stride, so it handles + * any in_dim (no blocks<=32 restriction). */ + const bool narrow = n_tok >= 2 && n_tok <= DS4_CUDA_NARROW_MAX && + getenv("DS4_CUDA_NO_Q8_NARROW") == NULL; + if (!narrow && g_cublas_ready && n_tok > 1) { const float *w_f32 = cuda_q8_f32_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label); if (w_f32) { const float alpha = 1.0f; @@ -12512,6 +12565,19 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode } const bool force_decode_warp = n_tok == 2u && g_glm_mtp_verify_mode; + if (narrow && !force_decode_warp) { + matmul_q8_0_preq_narrow_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>( + (float *)out->ptr, + reinterpret_cast(wptr), + xq, + xscale, + in_dim, + out_dim, + (uint32_t)n_tok, + blocks, + use_dp4a); + return cuda_ok(cudaGetLastError(), "matmul_q8_0 narrow launch"); + } if (n_tok > 1u && !force_decode_warp) { /* T matches the reduction width of whichever reference kernel would * have run: warp tree (32) for blocks <= 32, exact-thread tree From aa7b03577ea9d57660c9853d00beabc04a18e678 Mon Sep 17 00:00:00 2001 From: iCreil Date: Fri, 24 Jul 2026 22:23:29 +0200 Subject: [PATCH 04/10] CUDA: define multi-decode types locally (ds4_cuda.cu no longer includes ds4_gpu.h) --- ds4_cuda.cu | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 37538a081..c2b45d4bf 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -6103,6 +6103,23 @@ __global__ static void rope_tail_decode_rows_kernel( tail[i + 1u] = x0 * s + x1 * c; } +/* Batched-decode limits and per-sequence KV view. ds4_cuda.cu does not + * include ds4_gpu.h, so these mirror the definitions there and must stay in + * sync (DS4_GPU_DECODE_MULTI_MAX, ds4_gpu_attn_seqview). */ +#ifndef DS4_GPU_DECODE_MULTI_MAX +#define DS4_GPU_DECODE_MULTI_MAX 4u +#endif + +typedef struct { + const ds4_gpu_tensor *raw_kv; + const ds4_gpu_tensor *comp_kv; /* NULL when the layer keeps no compressed cache */ + const ds4_gpu_tensor *comp_mask; /* per-sequence mask row (n_comp floats), NULL = none */ + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; +} ds4_gpu_attn_seqview; + /* pos-vector twins of the two RoPE kernels above, for batched decode where * each token belongs to a different sequence at an arbitrary position. The * per-element arithmetic is kept identical so token t matches a single-token From fa659829fe45b926bf8ccf676d33166f50f9d96a Mon Sep 17 00:00:00 2001 From: iCreil Date: Fri, 24 Jul 2026 23:37:42 +0200 Subject: [PATCH 05/10] Graph: fused single-GPU session-batch decode (port of the batched-decode driver) Port of the fork's batched-decode driver onto the current main: - metal_graph_decode_compressor_indexer: per-sequence compressor/indexer helper re-extracted from the current decode layer phase (kept as a copy, single-token path untouched). - metal_graph_encode_layer_attention_decode_multi + encode_decode_multi: dense projections batched on the aliased batch workspace with n_tokens = n_seqs (narrow q8_0 kernel serves 2-4 rows), pos-vector RoPE, per-sequence KV store and compressor/indexer, multi-sequence attention with per-sequence fallbacks, FFN and output head via the existing batch functions. - ds4_sessions_eval_batch_cuda: single-tier decode batches take the fused pass instead of one-token pipelining (DS4_NO_FUSED_SESSION_BATCH=1 restores it); logits land one row per sequence in spec_logits, now allocated unless the fused path is disabled. --- ds4.c | 876 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 868 insertions(+), 8 deletions(-) diff --git a/ds4.c b/ds4.c index 1a79a01e4..20681f265 100644 --- a/ds4.c +++ b/ds4.c @@ -16752,7 +16752,11 @@ static bool metal_graph_alloc_raw_cap( const bool enable_splitkv_spec = metal_graph_cuda_splitkv_spec_requested(); const bool enable_splitkv_batch_verify = enable_splitkv_spec && metal_graph_cuda_splitkv_spec_batch_verify_requested(); - const bool enable_spec_logits = enable_mtp || enable_splitkv_batch_verify; + /* The fused single-GPU session-batch decode lands one logits row per + * sequence in spec_logits, so keep it allocated (16 rows, ~8 MiB) unless + * that path is disabled. */ + const bool enable_spec_logits = enable_mtp || enable_splitkv_batch_verify || + getenv("DS4_NO_FUSED_SESSION_BATCH") == NULL; const bool enable_prefix1_snapshot = enable_mtp || enable_splitkv_spec; const bool enable_frontier_snapshot = enable_mtp || @@ -29073,6 +29077,825 @@ static bool metal_graph_encode_layer_batch( return ok; } +/* Per-token compressor/indexer bookkeeping for one sequence, exactly as the + * single-token decode layer performs it (this mirrors the corresponding + * segment of metal_graph_encode_decode_layer_phase and must stay in sync with + * it). attn_norm and qr_norm are taken as parameters so the batched-decode + * driver can pass row views of the batch scratch; everything else is + * per-graph state, so per-sequence caches and counters advance exactly as + * independent single-token evals would. */ +static bool metal_graph_decode_compressor_indexer( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + ds4_gpu_tensor *attn_norm, + ds4_gpu_tensor *qr_norm, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + uint32_t *out_n_comp, + ds4_gpu_tensor **out_comp_cache, + ds4_gpu_tensor **out_comp_selected, + uint32_t *out_n_selected, + double *out_index_stage_t0) { + const bool compressed = ds4_layer_compress_ratio(il) != 0; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + bool ok = true; + uint32_t n_comp = 0; + ds4_gpu_tensor *comp_cache = NULL; + ds4_gpu_tensor *comp_selected = NULL; + uint32_t n_selected = 0; + double decode_index_stage_t0 = 0.0; + const bool decode_index_stage_profile = g->decode_index_stage_profile; + if (ok && compressed) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_width = coff * DS4_N_HEAD_DIM; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (!layer->attn_compressor_kv || !layer->attn_compressor_gate || + !layer->attn_compressor_ape || !layer->attn_compressor_norm || + layer->attn_compressor_kv->type != DS4_TENSOR_F16 || + layer->attn_compressor_gate->type != DS4_TENSOR_F16 || + layer->attn_compressor_kv->dim[0] != DS4_N_EMBD || + layer->attn_compressor_gate->dim[0] != DS4_N_EMBD || + layer->attn_compressor_kv->dim[1] != comp_width || + layer->attn_compressor_gate->dim[1] != comp_width) { + fprintf(stderr, "ds4: Metal graph compressor expects paired F16 compressor projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { + 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()) { + 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, + attn_norm, + 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, + attn_norm, + 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, + DS4_N_EMBD, comp_width, + attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, comp_width, + attn_norm, 1) != 0; + } + const uint32_t comp_row = g->layer_n_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_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], + metal_graph_attn_comp_update_target(g, il), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos, + metal_graph_attn_comp_update_row(comp_row), + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + comp_state_already_stored) != 0; + if (ok && emit) { + ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); + if (!comp_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; + } + ds4_gpu_tensor_free(comp_row_view); + if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); + } + if (ok && emit) g->layer_n_comp[il]++; + + if (ok && ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || + !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || + layer->indexer_compressor_kv->type != DS4_TENSOR_F16 || + layer->indexer_compressor_gate->type != DS4_TENSOR_F16 || + layer->indexer_compressor_kv->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_gate->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_kv->dim[1] != index_width || + layer->indexer_compressor_gate->dim[1] != index_width) { + fprintf(stderr, "ds4: Metal graph indexer compressor expects paired F16 projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { + 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()) { + 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, + attn_norm, + 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, + attn_norm, + 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, + DS4_N_EMBD, index_width, + attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, index_width, + attn_norm, 1) != 0; + } + const uint32_t index_row = g->layer_n_index_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_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], + g->layer_index_comp_cache[il], + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos, + index_row, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + index_state_already_stored) != 0; + if (ok && emit) { +#if defined(__APPLE__) + ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + if (!index_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + ds4_gpu_tensor_free(index_row_view); +#else + ds4_gpu_tensor index_row_view; + if (!metal_graph_borrow_tensor_view( + &index_row_view, + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float))) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(&index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } +#endif + } + if (ok && emit) g->layer_n_index_comp[il]++; + const uint32_t decode_sparse_threshold = + metal_graph_decode_indexer_sparse_threshold(g); + if (ok && + g->layer_n_comp[il] > decode_sparse_threshold && + g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K) { + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_attn_q_b || + !tensor_type_is_f16_or_q8_0(layer->indexer_attn_q_b->type) || + layer->indexer_attn_q_b->dim[0] != q_rank || + layer->indexer_attn_q_b->dim[1] != indexer_q_dim) { + fprintf(stderr, "ds4: Metal graph indexer q projection expects F16 or Q8_0 weights\n"); + ok = false; + } + if (ok && (!layer->indexer_proj || + layer->indexer_proj->type != DS4_TENSOR_F16 || + layer->indexer_proj->dim[0] != DS4_N_EMBD || + layer->indexer_proj->dim[1] != DS4_N_INDEXER_HEAD)) { + fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); + ok = false; + } + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_indexer_q(g), + model, + layer->indexer_attn_q_b, + q_rank, + indexer_q_dim, + qr_norm, + 1); + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_indexer_q(g), 1, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_indexer_q(g), + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_indexer_weights(g), model->map, model->size, + layer->indexer_proj->abs_offset, + DS4_N_EMBD, DS4_N_INDEXER_HEAD, + attn_norm, 1) != 0; + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), + metal_graph_indexer_q(g), + metal_graph_indexer_weights(g), + g->layer_index_comp_cache[il], + g->layer_n_index_comp[il], + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + index_scale) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_score", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + g->layer_n_index_comp[il], + 1, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_topk", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) { + comp_selected = metal_graph_comp_selected(g); + /* The indexer top-k is fixed by the model config; see the + * contract note in the single-token decode layer. */ + n_selected = DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] + ? DS4_N_INDEXER_TOP_K + : g->layer_n_index_comp[il]; + } + } + } + + n_comp = g->layer_n_comp[il]; + comp_cache = g->layer_attn_comp_cache[il]; + } + if (out_n_comp) *out_n_comp = n_comp; + if (out_comp_cache) *out_comp_cache = comp_cache; + if (out_comp_selected) *out_comp_selected = comp_selected; + if (out_n_selected) *out_n_selected = n_selected; + if (out_index_stage_t0) *out_index_stage_t0 = decode_index_stage_t0; + return ok; +} + +/* Batched-decode attention layer: one token per sequence, each sequence with + * its own graph (KV ring, compressor and indexer state) at its own position. + * Dense projections run once on the batch scratch with n_tokens = n_seqs and + * the pos-vector RoPE kernels; KV store and compressor/indexer bookkeeping + * run per sequence with the same single-token helpers the decode layer uses. + * Attention runs as one multi-sequence launch when every sequence is served + * by the regular decode kernel, and falls back to the per-sequence decode + * kernels (indexed or plain, which handle long contexts) otherwise. + * + * Sequencing constraint: comp_selected (and the other single-token scratch + * the compressor/indexer helper uses) is shared across sequences, so a + * sequence that takes the indexed-attention path must enqueue its attention + * before the next sequence's helper call overwrites the selection. */ +static bool metal_graph_encode_layer_attention_decode_multi( + ds4_gpu_graph *const *graphs, + uint32_t n_seqs, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + const uint32_t *pos) { + ds4_gpu_graph *g = graphs[0]; /* aliased batch workspace and config flags */ + if (n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX || n_seqs > g->prefill_cap) 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; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint32_t n_groups = DS4_N_OUT_GROUP; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = DS4_N_LORA_O; + const uint32_t ratio = ds4_layer_compress_ratio(il); + const bool compressed = ratio != 0; + const float freq_base = layer_rope_freq_base(il); + const float freq_scale = layer_rope_freq_scale(il); + const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; + float attn_factor = 1.0f; + if (ext_factor != 0.0f && freq_scale > 0.0f) { + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); + const uint32_t rope_ctx = compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0; + + ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_mix(g), 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); + ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( + metal_graph_batch_hc_split(g), 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); + ds4_gpu_tensor *attn_cur_view = ds4_gpu_tensor_view( + metal_graph_batch_attn_cur(g), 0, (uint64_t)n_seqs * DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( + metal_graph_batch_after_attn_hc(g), 0, (uint64_t)n_seqs * 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 && + !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_seqs, + 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_seqs) != 0; + if (metal_graph_use_reference_hc_decode()) { + if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, + hc_mix_view, + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(attn_cur_view, + metal_graph_batch_cur_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (fuse_hc_norm) { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, + metal_graph_batch_attn_norm(g), + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, + hc_split_view, + hc_mix_view, + metal_graph_batch_cur_hc(g), + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS) != 0; + } + if (ok && !fuse_hc_norm) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_attn_norm(g), + metal_graph_batch_attn_cur(g), + model->map, + model->size, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + n_seqs, + DS4_RMS_EPS) != 0; + } + + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_a", + il, + pos[0], + metal_graph_batch_qr(g), + model, + layer->attn_q_a, + DS4_N_EMBD, + q_rank, + metal_graph_batch_attn_norm(g), + n_seqs); + if (qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", + il, + pos[0], + metal_graph_batch_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + n_seqs); + if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(metal_graph_batch_qr_norm(g), + metal_graph_batch_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + n_seqs, + DS4_RMS_EPS) != 0; + } else { + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_qr_norm(g), + metal_graph_batch_qr(g), + model->map, + model->size, + layer->attn_q_a_norm->abs_offset, + (uint32_t)q_rank, + n_seqs, + DS4_RMS_EPS) != 0; + } + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_b", + il, + pos[0], + metal_graph_batch_q(g), + model, + layer->attn_q_b, + q_rank, + q_dim, + metal_graph_batch_qr_norm(g), + n_seqs); + if (ok) ok = ds4_gpu_head_rms_norm_rope_tail_multi_tensor(metal_graph_batch_q(g), + n_seqs, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + rope_ctx, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + if (!qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", + il, + pos[0], + metal_graph_batch_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_batch_attn_norm(g), + n_seqs); + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(metal_graph_batch_kv(g), + metal_graph_batch_kv_raw(g), + model->map, + model->size, + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, + n_seqs, + DS4_RMS_EPS) != 0; + } + if (ok) ok = ds4_gpu_rope_tail_multi_tensor(metal_graph_batch_kv(g), + n_seqs, + DS4_N_HEAD_KV, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + rope_ctx, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + + /* Per-sequence: KV store, compressor/indexer bookkeeping, and attention + * for the sequences the multi kernel cannot serve. */ + ds4_gpu_attn_seqview seqs[DS4_GPU_DECODE_MULTI_MAX]; + bool seq_done[DS4_GPU_DECODE_MULTI_MAX] = {false}; + bool all_multi = true; + memset(seqs, 0, sizeof(seqs)); + for (uint32_t b = 0; ok && b < n_seqs; b++) { + ds4_gpu_graph *gb = graphs[b]; + const uint32_t raw_row = pos[b] % gb->raw_cap; + const uint32_t n_raw = metal_graph_raw_span_for_batch(gb, pos[b], 1); + const uint32_t raw_start = metal_graph_raw_start_for_span(gb, pos[b], n_raw); + ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(metal_graph_batch_kv(g), b, DS4_N_HEAD_DIM); + ds4_gpu_tensor *attn_norm_view = metal_graph_tensor_row_view(metal_graph_batch_attn_norm(g), b, DS4_N_EMBD); + ds4_gpu_tensor *qr_norm_view = metal_graph_tensor_row_view(metal_graph_batch_qr_norm(g), b, q_rank); + ok = kv_view && attn_norm_view && qr_norm_view; + if (ok) ok = metal_graph_decode_kv_store(kv_view, gb->layer_raw_cache[il], gb->raw_cap, raw_row); + if (ok) ok = metal_graph_cuda_tp_attn_cache_sync_raw_row(gb, il, raw_row); + + uint32_t n_comp = 0; + ds4_gpu_tensor *comp_cache = NULL; + ds4_gpu_tensor *comp_selected = NULL; + uint32_t n_selected = 0; + double index_stage_t0 = 0.0; + if (ok) { + ok = metal_graph_decode_compressor_indexer(gb, model, layer, il, pos[b], + attn_norm_view, qr_norm_view, + freq_base, freq_scale, + ext_factor, attn_factor, + &n_comp, &comp_cache, + &comp_selected, &n_selected, + &index_stage_t0); + } + if (ok && comp_selected != NULL && n_selected != 0) { + /* Indexed attention consumes comp_selected: run it before the + * next sequence's compressor/indexer pass overwrites it. */ + ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(metal_graph_batch_q(g), b, q_dim); + ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(metal_graph_batch_heads(g), b, q_dim); + ok = q_view && heads_view && + ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + heads_view, + model->map, + model->size, + layer->attn_sinks->abs_offset, + q_view, + gb->layer_raw_cache[il], + gb->layer_attn_comp_cache[il], + metal_graph_attn_comp_cache_is_f16(), + comp_selected, + 1, + pos[b], + n_raw, + gb->raw_cap, + raw_start, + n_comp, + n_selected, + gb->raw_window, + ratio, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + ds4_gpu_tensor_free(heads_view); + ds4_gpu_tensor_free(q_view); + seq_done[b] = true; + all_multi = false; + } else if (ok) { + if (!ds4_gpu_attention_decode_multi_supported(n_comp, + n_comp ? metal_graph_attn_comp_cache_is_f16() : 0)) { + all_multi = false; + } + seqs[b].raw_kv = gb->layer_raw_cache[il]; + seqs[b].comp_kv = n_comp ? gb->layer_attn_comp_cache[il] : NULL; + seqs[b].comp_mask = NULL; + seqs[b].n_raw = n_raw; + seqs[b].raw_cap = gb->raw_cap; + seqs[b].raw_start = raw_start; + seqs[b].n_comp = n_comp; + } + ds4_gpu_tensor_free(qr_norm_view); + ds4_gpu_tensor_free(attn_norm_view); + ds4_gpu_tensor_free(kv_view); + } + if (ok && all_multi) { + ok = ds4_gpu_attention_decode_heads_multi_tensor(metal_graph_batch_heads(g), + model->map, + model->size, + layer->attn_sinks->abs_offset, + metal_graph_batch_q(g), + seqs, + n_seqs, + 0, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + } else if (ok) { + for (uint32_t b = 0; ok && b < n_seqs; b++) { + if (seq_done[b]) continue; + ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(metal_graph_batch_q(g), b, q_dim); + ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(metal_graph_batch_heads(g), b, q_dim); + ok = q_view && heads_view && + ds4_gpu_attention_decode_heads_tensor(heads_view, + model->map, + model->size, + layer->attn_sinks->abs_offset, + q_view, + seqs[b].raw_kv, + seqs[b].n_raw, + seqs[b].raw_cap, + seqs[b].raw_start, + seqs[b].comp_kv, + metal_graph_attn_comp_cache_is_f16(), + seqs[b].n_comp, + NULL, + 0, + DS4_N_HEAD, + DS4_N_HEAD_DIM) != 0; + ds4_gpu_tensor_free(heads_view); + ds4_gpu_tensor_free(q_view); + } + } + + if (ok) ok = ds4_gpu_rope_tail_multi_tensor(metal_graph_batch_heads(g), + n_seqs, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos, + rope_ctx, + true, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_attention_output_q8_batch_tensor(metal_graph_batch_attn_out(g), + metal_graph_batch_attn_low(g), + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + metal_graph_batch_heads(g), + n_seqs) != 0; + if (ok) ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, + metal_graph_batch_attn_out(g), + metal_graph_batch_cur_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + ds4_gpu_tensor_free(after_attn_hc_view); + ds4_gpu_tensor_free(attn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); + return ok; +} + +/* Encode one decode token for each of n_seqs sequences in a single pass over + * the layers: attention via the batched-decode layer above, FFN and output + * head via the existing batch functions (graphs[0] carries the aliased batch + * workspace; the FFN stages use no per-graph state). The per-sequence caches + * and counters advance exactly as n_seqs independent single-token evals + * would. Logits for sequence b land in spec_logits row b of graphs[0]. */ +static bool metal_graph_encode_decode_multi( + ds4_gpu_graph *const *graphs, + uint32_t n_seqs, + const ds4_model *model, + const ds4_weights *weights, + const int *tokens, + const uint32_t *pos, + bool need_logits) { + ds4_gpu_graph *g = graphs[0]; + if (n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX) return false; + if (metal_graph_directional_steering_attn_enabled(g)) return false; + if (need_logits && !g->spec_logits) return false; + for (uint32_t b = 0; b < n_seqs; b++) { + if (graphs[b]->raw_cap == 0 || graphs[b]->ssd_streaming) return false; + } + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + int32_t token_ids[DS4_GPU_DECODE_MULTI_MAX]; + for (uint32_t b = 0; b < n_seqs; b++) token_ids[b] = tokens[b]; + bool ok = ds4_gpu_tensor_write(metal_graph_prefill_tokens(g), 0, token_ids, + (uint64_t)n_seqs * sizeof(token_ids[0])) != 0; + for (uint32_t b = 0; ok && b < n_seqs; b++) { + metal_graph_dspark_capture_begin(graphs[b]); + ds4_gpu_tensor *hc_view = metal_graph_tensor_row_view(metal_graph_batch_cur_hc(g), b, hc_dim); + ok = hc_view && ds4_gpu_embed_token_hc_tensor(hc_view, + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)tokens[b], + DS4_N_EMBD, + DS4_N_HC) != 0; + ds4_gpu_tensor_free(hc_view); + } + + const uint32_t split_after_layers = metal_graph_token_split_after_layers(); + for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + ok = metal_graph_encode_layer_attention_decode_multi(graphs, n_seqs, model, + &weights->layer[il], il, pos); + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u batched decode attention failed\n", il); + } + if (ok) { + ok = metal_graph_encode_layer_ffn_batch(g, model, &weights->layer[il], + il, pos[0], n_seqs, NULL, 0); + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u batched decode ffn failed\n", il); + } + } + if (ok) { + ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); + g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); + g->batch_next_hc_by_tier[g->active_tier] = tmp; + } + if (ok && split_after_layers != 0 && il + 1u == split_after_layers) { + ok = ds4_gpu_flush_commands() != 0; + } + } + if (ok && need_logits) { + ok = metal_graph_encode_output_head_batch(g, model, weights, n_seqs, + weights->output->dim[1]); + } + return ok; +} + static bool metal_graph_eval_token_raw_swa_streaming( ds4_gpu_graph *g, const ds4_model *model, @@ -63600,10 +64423,41 @@ static int ds4_sessions_eval_batch_cuda(ds4_decode_item *items, int count, const char *interleave = getenv("DS4_CUDA_SESSION_BATCH_INTERLEAVE"); const bool use_pipeline = !interleave || !interleave[0] || strcmp(interleave, "0") != 0; - if (ok && use_pipeline && first->graph.placement) { + /* Single-GPU fused decode: run the whole decode batch as one graph + * pass (dense projections batched with n_tokens = count, per-sequence + * KV/compressor/indexer, multi-sequence attention). This is the + * single-tier counterpart of the placement pipeline below; kill + * switch: DS4_NO_FUSED_SESSION_BATCH=1 restores the one-token + * pipelining. Not bitwise-identical to serial decode (batched matmul + * reduction order), greedy token streams match. */ + bool fused = ok && !first->graph.placement && + getenv("DS4_NO_FUSED_SESSION_BATCH") == NULL && + count >= 2 && count <= (int)DS4_GPU_DECODE_MULTI_MAX && + first->graph.spec_logits != NULL; + for (int i = 0; fused && i < count; i++) { + ds4_session *s = items[i].session; + if (s->graph.ssd_streaming || s->graph.raw_cap == 0 || + s->distributed) { + fused = false; + } + } + if (ok && fused) { + ds4_gpu_graph *fused_graphs[DS4_GPU_DECODE_MULTI_MAX]; + int fused_tokens[DS4_GPU_DECODE_MULTI_MAX]; + uint32_t fused_pos[DS4_GPU_DECODE_MULTI_MAX]; + for (int i = 0; i < count; i++) { + ds4_session *s = items[i].session; + fused_graphs[i] = &s->graph; + fused_tokens[i] = items[i].token; + fused_pos[i] = (uint32_t)s->checkpoint.len; + } + ok = metal_graph_encode_decode_multi(fused_graphs, (uint32_t)count, + &e->model, &e->weights, + fused_tokens, fused_pos, true); + } else if (ok && use_pipeline && first->graph.placement) { ok = metal_graph_encode_session_pipeline_batch( items, count, &e->model, &e->weights); - } else { + } else if (ok) { for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; ok = metal_graph_encode_token_raw_swa( @@ -63621,11 +64475,17 @@ static int ds4_sessions_eval_batch_cuda(ds4_decode_item *items, int count, for (int i = 0; ok && i < count; i++) { ds4_session *s = items[i].session; - ok = ds4_gpu_tensor_read( - metal_graph_logits(&s->graph), - 0, - s->logits, - (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + if (fused) { + ok = metal_graph_read_spec_logits_row(&first->graph, + (uint32_t)i, + s->logits); + } else { + ok = ds4_gpu_tensor_read( + metal_graph_logits(&s->graph), + 0, + s->logits, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } } if (!ok) { for (int i = 0; i < count; i++) { From 96540c5f45fbe25ec6c46a4ca63f8e63151d5676 Mon Sep 17 00:00:00 2001 From: iCreil Date: Fri, 24 Jul 2026 23:39:10 +0200 Subject: [PATCH 06/10] CUDA: add ds4_gpu_attention_decode_multi_supported --- ds4_cuda.cu | 4 + ds4_gpu.h | 7 + scratchpad/current_segment.txt | 422 ++++++++++++++++++++++ scratchpad/old_encoder_commit.diff | 518 +++++++++++++++++++++++++++ scratchpad/old_ffn_multi_commit.diff | 287 +++++++++++++++ scratchpad/old_helper_full.txt | 388 ++++++++++++++++++++ 6 files changed, 1626 insertions(+) create mode 100644 scratchpad/current_segment.txt create mode 100644 scratchpad/old_encoder_commit.diff create mode 100644 scratchpad/old_ffn_multi_commit.diff create mode 100644 scratchpad/old_helper_full.txt diff --git a/ds4_cuda.cu b/ds4_cuda.cu index c2b45d4bf..694c309bf 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -15218,6 +15218,10 @@ extern "C" int ds4_gpu_attention_decode_rows_rope_tensor( "attention decode rows inverse rope launch"); } +extern "C" int ds4_gpu_attention_decode_multi_supported(uint32_t n_comp, uint32_t comp_kv_f16) { + return comp_kv_f16 == 0 && cuda_attention_score_buffer_fits(n_comp) ? 1 : 0; +} + extern "C" int ds4_gpu_attention_decode_heads_multi_tensor( ds4_gpu_tensor *heads, const void *model_map, diff --git a/ds4_gpu.h b/ds4_gpu.h index 15f74c363..eae694df7 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -1806,6 +1806,13 @@ typedef struct { uint32_t n_comp; } ds4_gpu_attn_seqview; +/* True when the multi-sequence decode attention kernel can serve a sequence + * with n_comp visible compressed rows (score buffer bound, no online variant + * yet) and the compressed cache layout it supports. */ +int ds4_gpu_attention_decode_multi_supported( + uint32_t n_comp, + uint32_t comp_kv_f16); + int ds4_gpu_attention_decode_heads_multi_tensor( ds4_gpu_tensor *heads, const void *model_map, diff --git a/scratchpad/current_segment.txt b/scratchpad/current_segment.txt new file mode 100644 index 000000000..b1de891e1 --- /dev/null +++ b/scratchpad/current_segment.txt @@ -0,0 +1,422 @@ + DS4_METAL_PROFILE_DECODE_STAGE("q_path"); + if (ok) { + metal_graph_debug_dump_tensor("Qcur", metal_graph_q(g), q_dim, il, pos); + } + if (!qkv_rms_fused) { + if (ok) ok = metal_graph_matmul_dense_quant_tensor(metal_graph_kv_raw(g), + model, + layer->attn_kv, + DS4_N_EMBD, + DS4_N_HEAD_DIM, + metal_graph_attn_norm(g), + 1); + if (ok) { + metal_graph_debug_dump_tensor("KVraw", metal_graph_kv_raw(g), DS4_N_HEAD_DIM, il, pos); + } + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(metal_graph_kv(g), metal_graph_kv_raw(g), + model->map, model->size, + layer->attn_kv_a_norm->abs_offset, + DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVnorm", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + const bool tp_ablate_kv = metal_graph_tp_ablate("kv"); + if (ok && !tp_ablate_kv && !kv_rope_fused) { + ok = ds4_gpu_rope_tail_tensor(metal_graph_kv(g), 1, + DS4_N_HEAD_KV, DS4_N_HEAD_DIM, + DS4_N_ROT, pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, freq_base, freq_scale, + ext_factor, attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor("KVrope", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + if (!resume_after_kv_store) { + /* The common no-debug path may fuse KV RMS with RoPE above. KV + * storage starts here after metal_graph_kv(g) contains the RoPE row. */ + if (ok) ok = metal_graph_decode_kv_store(metal_graph_kv(g), raw_cache, raw_cap, raw_row); + if (ok) ok = metal_graph_cuda_tp_attn_cache_sync_raw_row(g, il, raw_row); + DS4_METAL_PROFILE_DECODE_STAGE("kv_path"); + if (ok) { + metal_graph_debug_dump_tensor("KVcur", metal_graph_kv(g), DS4_N_HEAD_DIM, il, pos); + } + } + + uint32_t n_comp = 0; + ds4_gpu_tensor *comp_cache = NULL; + ds4_gpu_tensor *comp_selected = NULL; + uint32_t n_selected = 0; + double decode_index_stage_t0 = 0.0; + const bool decode_index_stage_profile = g->decode_index_stage_profile; + if (ok && compressed) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_width = coff * DS4_N_HEAD_DIM; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (!layer->attn_compressor_kv || !layer->attn_compressor_gate || + !layer->attn_compressor_ape || !layer->attn_compressor_norm || + layer->attn_compressor_kv->type != DS4_TENSOR_F16 || + layer->attn_compressor_gate->type != DS4_TENSOR_F16 || + layer->attn_compressor_kv->dim[0] != DS4_N_EMBD || + layer->attn_compressor_gate->dim[0] != DS4_N_EMBD || + layer->attn_compressor_kv->dim[1] != comp_width || + layer->attn_compressor_gate->dim[1] != comp_width) { + fprintf(stderr, "ds4: Metal graph compressor expects paired F16 compressor projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { + 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()) { + 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, + DS4_N_EMBD, comp_width, + metal_graph_attn_norm(g), 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, comp_width, + metal_graph_attn_norm(g), 1) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("compressor_proj"); + const uint32_t comp_row = g->layer_n_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_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], + metal_graph_attn_comp_update_target(g, il), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos, + metal_graph_attn_comp_update_row(comp_row), + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + 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); + if (!comp_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVcompress", comp_row_view, DS4_N_HEAD_DIM, il, pos); + } + } + ds4_gpu_tensor_free(comp_row_view); + DS4_METAL_PROFILE_DECODE_STAGE("compressor_quantize"); + if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); + DS4_METAL_PROFILE_DECODE_STAGE("compressor_commit"); + } + if (ok && emit) g->layer_n_comp[il]++; + + if (ok && ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || + !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || + layer->indexer_compressor_kv->type != DS4_TENSOR_F16 || + layer->indexer_compressor_gate->type != DS4_TENSOR_F16 || + layer->indexer_compressor_kv->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_gate->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_kv->dim[1] != index_width || + layer->indexer_compressor_gate->dim[1] != index_width) { + fprintf(stderr, "ds4: Metal graph indexer compressor expects paired F16 projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { + 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()) { + 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, + DS4_N_EMBD, index_width, + metal_graph_attn_norm(g), 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_comp_sc_cur(g), model->map, model->size, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, index_width, + metal_graph_attn_norm(g), 1) != 0; + } + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_proj"); + const uint32_t index_row = g->layer_n_index_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_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], + g->layer_index_comp_cache[il], + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos, + index_row, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS, + index_state_already_stored) != 0; + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_update"); + if (ok && emit) { +#if defined(__APPLE__) + ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + if (!index_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } + ds4_gpu_tensor_free(index_row_view); +#else + ds4_gpu_tensor index_row_view; + if (!metal_graph_borrow_tensor_view( + &index_row_view, + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float))) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(&index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + } +#endif + DS4_METAL_PROFILE_DECODE_STAGE("indexer_compressor_qat"); + } + if (ok && emit) g->layer_n_index_comp[il]++; + const uint32_t decode_sparse_threshold = + metal_graph_decode_indexer_sparse_threshold(g); + if (ok && + g->layer_n_comp[il] > decode_sparse_threshold && + g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K) { + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_attn_q_b || + !tensor_type_is_f16_or_q8_0(layer->indexer_attn_q_b->type) || + layer->indexer_attn_q_b->dim[0] != q_rank || + layer->indexer_attn_q_b->dim[1] != indexer_q_dim) { + fprintf(stderr, "ds4: Metal graph indexer q projection expects F16 or Q8_0 weights\n"); + ok = false; + } + if (ok && (!layer->indexer_proj || + layer->indexer_proj->type != DS4_TENSOR_F16 || + layer->indexer_proj->dim[0] != DS4_N_EMBD || + layer->indexer_proj->dim[1] != DS4_N_INDEXER_HEAD)) { + fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); + ok = false; + } + if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_indexer_q(g), + model, + layer->indexer_attn_q_b, + q_rank, + indexer_q_dim, + metal_graph_qr_norm(g), + 1); + if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_indexer_q(g), 1, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_indexer_q(g), + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_indexer_weights(g), model->map, model->size, + layer->indexer_proj->abs_offset, + DS4_N_EMBD, DS4_N_INDEXER_HEAD, + metal_graph_attn_norm(g), 1) != 0; + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_score_one_tensor(metal_graph_indexer_scores(g), + metal_graph_indexer_q(g), + metal_graph_indexer_weights(g), + g->layer_index_comp_cache[il], + g->layer_n_index_comp[il], + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + index_scale) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_score", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_topk_tensor(metal_graph_comp_selected(g), + metal_graph_indexer_scores(g), + g->layer_n_index_comp[il], + 1, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_topk", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + /* Decode used to materialize a dense compressed-row mask and + * call the generic gathered FlashAttention wrapper below. + * That wrapper scans every compressed row and rejects long + * contexts once raw+compressed rows exceed 8192. Ratio-4 DS4 + * attention is sparse after indexer top-k, so use the private + * indexed attention kernel instead: it scans only SWA raw rows + * plus the selected compressed rows, matching prefill and + * avoiding the long-context decode failure. */ + if (ok) { + comp_selected = metal_graph_comp_selected(g); + /* + * Contract: the indexer top-k is fixed by the model config + * and must remain the full 512 rows. Do not reduce this for + * throughput benchmarks. + * + * Why: the indexer is not just an implementation detail. It + * decides which compressed memory rows are visible to the + * attention kernel. If we keep only 128/256 rows, the later + * indexed-attention math may be perfectly computed, but it is + * computed over the wrong candidate set: rows ranked 257-512 + * are removed before softmax/PV can use them. Those rows may + * carry weak-but-necessary evidence for retrieval, name/number + * recall, or long-context disambiguation. The error is + * therefore semantic/algorithmic, not the acceptable kind of + * local numerical drift caused by a different reduction order + * or Tensor/NAX precision. + * + * Short prompt tests, first-token agreement, or even a small + * official-vector set can miss this because many prompts do + * not need the tail of the 512 selected compressed rows. The + * failure appears only when the model needs information that + * fell below the reduced cutoff. Optimizations belong inside + * the score/top-k/attention implementation while preserving + * DS4_N_INDEXER_TOP_K. + */ + n_selected = DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] + ? DS4_N_INDEXER_TOP_K + : g->layer_n_index_comp[il]; + } + } + } + + n_comp = g->layer_n_comp[il]; + comp_cache = g->layer_attn_comp_cache[il]; + } + DS4_METAL_PROFILE_DECODE_STAGE("compressor_indexer"); + + if (stop_before_attn) return ok; + if (ok) { + const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); + const bool indexed_attention = n_comp != 0 && comp_selected != NULL && n_selected != 0; diff --git a/scratchpad/old_encoder_commit.diff b/scratchpad/old_encoder_commit.diff new file mode 100644 index 000000000..d1f1fb475 --- /dev/null +++ b/scratchpad/old_encoder_commit.diff @@ -0,0 +1,518 @@ +commit 6d8fe29c6b8cbaf725b93b3df5ecb112b5eaf288 +Author: iCreil +Date: Tue Jul 7 13:52:15 2026 +0200 + + Graph: add batched-decode encoder (one token per sequence) + + metal_graph_encode_decode_multi evaluates one decode token for each of up + to DS4_GPU_DECODE_MULTI_MAX sequences in a single pass over the layers. + Dense projections run once on the batch scratch with n_tokens = n_seqs, + reusing the prefill FFN and output-head batch functions unchanged (token + ids go through scr->prefill_tokens for hash routing; logits land one row + per sequence in scr->spec_logits). RoPE uses the pos-vector kernels, and + KV append plus compressor/indexer bookkeeping run per sequence with the + same single-token helpers the decode layer uses, so per-sequence state + advances exactly as independent evals would. + + Attention runs as one multi-sequence launch when every sequence is served + by the regular decode kernel; sequences on the indexed path enqueue their + attention inside the per-sequence loop (scr->comp_selected is shared and + must be consumed before the next sequence overwrites it), and anything the + multi kernel cannot serve (compressed rows beyond the score buffer, F16 + compressed cache) falls back to the existing per-sequence kernels, so + there is no position limit. + + Additions only; nothing calls this yet — the session-level API and the + server scheduler wire it up next. + +diff --git a/ds4.c b/ds4.c +index 6229d13..f76f5f7 100644 +--- a/ds4.c ++++ b/ds4.c +@@ -19434,6 +19434,453 @@ static bool metal_graph_encode_layer_batch( + return ok; + } + ++/* Batched-decode attention layer: one token per sequence, each sequence with ++ * its own graph (KV ring, compressor and indexer state) at its own position. ++ * Dense projections run once on the batch scratch with n_tokens = n_seqs and ++ * the pos-vector RoPE kernels; KV store and compressor/indexer bookkeeping ++ * run per sequence with the same single-token helpers the decode layer uses. ++ * Attention runs as one multi-sequence launch when every sequence is served ++ * by the regular decode kernel, and falls back to the per-sequence decode ++ * kernels (indexed or plain, which handle long contexts) otherwise. ++ * ++ * Sequencing constraint: scr->comp_selected (and the other single-token ++ * scratch the compressor/indexer helper uses) is shared across sequences, so ++ * a sequence that takes the indexed-attention path must enqueue its attention ++ * before the next sequence's helper call overwrites the selection. */ ++static bool metal_graph_encode_layer_attention_decode_multi( ++ ds4_gpu_graph *const *graphs, ++ uint32_t n_seqs, ++ const ds4_model *model, ++ const ds4_layer_weights *layer, ++ uint32_t il, ++ const uint32_t *pos) { ++ ds4_gpu_graph *g = graphs[0]; /* engine-shared scratch and config flags */ ++ if (n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX || n_seqs > g->prefill_cap) 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; ++ const uint64_t q_rank = layer->attn_q_a->dim[1]; ++ const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; ++ const uint32_t n_groups = DS4_N_OUT_GROUP; ++ const uint32_t group_heads = DS4_N_HEAD / n_groups; ++ const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; ++ const uint32_t rank = DS4_N_LORA_O; ++ const uint32_t ratio = ds4_layer_compress_ratio(il); ++ const bool compressed = ratio != 0; ++ const float freq_base = layer_rope_freq_base(il); ++ const float freq_scale = layer_rope_freq_scale(il); ++ const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; ++ float attn_factor = 1.0f; ++ if (ext_factor != 0.0f && freq_scale > 0.0f) { ++ attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); ++ } ++ const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); ++ const uint32_t rope_ctx = compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0; ++ ++ ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( ++ g->scr->batch_hc_mix, 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); ++ ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( ++ g->scr->batch_hc_split, 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); ++ ds4_gpu_tensor *attn_cur_view = ds4_gpu_tensor_view( ++ g->scr->batch_attn_cur, 0, (uint64_t)n_seqs * DS4_N_EMBD * sizeof(float)); ++ ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( ++ g->scr->batch_after_attn_hc, 0, (uint64_t)n_seqs * 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 && ++ !metal_graph_use_reference_hc_decode() && ++ metal_graph_enable_batch_hc_norm_fusion(); ++ ++ if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->scr->batch_flat_hc, ++ g->scr->batch_cur_hc, ++ (uint32_t)hc_dim, ++ n_seqs, ++ 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, ++ g->scr->batch_flat_hc, ++ n_seqs) != 0; ++ if (metal_graph_use_reference_hc_decode()) { ++ if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, ++ hc_mix_view, ++ model->map, ++ model->size, ++ layer->hc_attn_scale->abs_offset, ++ layer->hc_attn_base->abs_offset, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS) != 0; ++ if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(attn_cur_view, ++ g->scr->batch_cur_hc, ++ hc_split_view, ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ } else if (fuse_hc_norm) { ++ if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, ++ g->scr->batch_attn_norm, ++ hc_split_view, ++ hc_mix_view, ++ g->scr->batch_cur_hc, ++ model->map, ++ model->size, ++ layer->hc_attn_scale->abs_offset, ++ layer->hc_attn_base->abs_offset, ++ layer->attn_norm->abs_offset, ++ DS4_N_EMBD, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS, ++ DS4_RMS_EPS) != 0; ++ } else { ++ if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, ++ hc_split_view, ++ hc_mix_view, ++ g->scr->batch_cur_hc, ++ model->map, ++ model->size, ++ layer->hc_attn_scale->abs_offset, ++ layer->hc_attn_base->abs_offset, ++ DS4_N_EMBD, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS) != 0; ++ } ++ if (ok && !fuse_hc_norm) { ++ ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_attn_norm, ++ g->scr->batch_attn_cur, ++ model->map, ++ model->size, ++ layer->attn_norm->abs_offset, ++ DS4_N_EMBD, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ } ++ ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_a", ++ il, ++ pos[0], ++ g->scr->batch_qr, ++ model, ++ layer->attn_q_a, ++ DS4_N_EMBD, ++ q_rank, ++ g->scr->batch_attn_norm, ++ n_seqs); ++ if (qkv_rms_fused) { ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", ++ il, ++ pos[0], ++ g->scr->batch_kv_raw, ++ model, ++ layer->attn_kv, ++ DS4_N_EMBD, ++ DS4_N_HEAD_DIM, ++ g->scr->batch_attn_norm, ++ n_seqs); ++ if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->scr->batch_qr_norm, ++ g->scr->batch_qr, ++ model->map, ++ model->size, ++ layer->attn_q_a_norm->abs_offset, ++ (uint32_t)q_rank, ++ g->scr->batch_kv, ++ g->scr->batch_kv_raw, ++ layer->attn_kv_a_norm->abs_offset, ++ DS4_N_HEAD_DIM, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ } else { ++ if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_qr_norm, ++ g->scr->batch_qr, ++ model->map, ++ model->size, ++ layer->attn_q_a_norm->abs_offset, ++ (uint32_t)q_rank, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ } ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_b", ++ il, ++ pos[0], ++ g->scr->batch_q, ++ model, ++ layer->attn_q_b, ++ q_rank, ++ q_dim, ++ g->scr->batch_qr_norm, ++ n_seqs); ++ if (ok) ok = ds4_gpu_head_rms_norm_rope_tail_multi_tensor(g->scr->batch_q, ++ n_seqs, ++ DS4_N_HEAD, ++ DS4_N_HEAD_DIM, ++ DS4_N_ROT, ++ pos, ++ rope_ctx, ++ false, ++ freq_base, ++ freq_scale, ++ ext_factor, ++ attn_factor, ++ DS4_ROPE_YARN_BETA_FAST, ++ DS4_ROPE_YARN_BETA_SLOW, ++ DS4_RMS_EPS) != 0; ++ if (!qkv_rms_fused) { ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", ++ il, ++ pos[0], ++ g->scr->batch_kv_raw, ++ model, ++ layer->attn_kv, ++ DS4_N_EMBD, ++ DS4_N_HEAD_DIM, ++ g->scr->batch_attn_norm, ++ n_seqs); ++ if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_kv, ++ g->scr->batch_kv_raw, ++ model->map, ++ model->size, ++ layer->attn_kv_a_norm->abs_offset, ++ DS4_N_HEAD_DIM, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ } ++ if (ok) ok = ds4_gpu_rope_tail_multi_tensor(g->scr->batch_kv, ++ n_seqs, ++ DS4_N_HEAD_KV, ++ DS4_N_HEAD_DIM, ++ DS4_N_ROT, ++ pos, ++ rope_ctx, ++ false, ++ freq_base, ++ freq_scale, ++ ext_factor, ++ attn_factor, ++ DS4_ROPE_YARN_BETA_FAST, ++ DS4_ROPE_YARN_BETA_SLOW) != 0; ++ ++ /* Per-sequence: KV store, compressor/indexer bookkeeping, and attention ++ * for the sequences the multi kernel cannot serve. */ ++ ds4_gpu_attn_seqview seqs[DS4_GPU_DECODE_MULTI_MAX]; ++ bool seq_done[DS4_GPU_DECODE_MULTI_MAX] = {false}; ++ bool all_multi = true; ++ memset(seqs, 0, sizeof(seqs)); ++ for (uint32_t b = 0; ok && b < n_seqs; b++) { ++ ds4_gpu_graph *gb = graphs[b]; ++ const uint32_t raw_row = pos[b] % gb->raw_cap; ++ const uint32_t n_raw = metal_graph_raw_span_for_batch(gb, pos[b], 1); ++ const uint32_t raw_start = metal_graph_raw_start_for_span(gb, pos[b], n_raw); ++ ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(g->scr->batch_kv, b, DS4_N_HEAD_DIM); ++ ds4_gpu_tensor *attn_norm_view = metal_graph_tensor_row_view(g->scr->batch_attn_norm, b, DS4_N_EMBD); ++ ds4_gpu_tensor *qr_norm_view = metal_graph_tensor_row_view(g->scr->batch_qr_norm, b, q_rank); ++ ok = kv_view && attn_norm_view && qr_norm_view; ++ if (ok) ok = metal_graph_decode_kv_store(kv_view, gb->layer_raw_cache[il], gb->raw_cap, raw_row); ++ ++ uint32_t n_comp = 0; ++ ds4_gpu_tensor *comp_cache = NULL; ++ ds4_gpu_tensor *comp_selected = NULL; ++ uint32_t n_selected = 0; ++ double index_stage_t0 = 0.0; ++ if (ok) { ++ ok = metal_graph_decode_compressor_indexer(gb, model, layer, il, pos[b], ++ attn_norm_view, qr_norm_view, ++ freq_base, freq_scale, ++ ext_factor, attn_factor, ++ &n_comp, &comp_cache, ++ &comp_selected, &n_selected, ++ &index_stage_t0); ++ } ++ if (ok && comp_selected != NULL && n_selected != 0) { ++ /* Indexed attention consumes scr->comp_selected: run it before the ++ * next sequence's compressor/indexer pass overwrites it. */ ++ ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(g->scr->batch_q, b, q_dim); ++ ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(g->scr->batch_heads, b, q_dim); ++ ok = q_view && heads_view && ++ ds4_gpu_attention_indexed_mixed_batch_heads_tensor( ++ heads_view, ++ model->map, ++ model->size, ++ layer->attn_sinks->abs_offset, ++ q_view, ++ gb->layer_raw_cache[il], ++ gb->layer_attn_comp_cache[il], ++ metal_graph_attn_comp_cache_is_f16(), ++ comp_selected, ++ 1, ++ pos[b], ++ n_raw, ++ gb->raw_cap, ++ raw_start, ++ n_comp, ++ n_selected, ++ gb->raw_window, ++ ratio, ++ DS4_N_HEAD, ++ DS4_N_HEAD_DIM) != 0; ++ ds4_gpu_tensor_free(heads_view); ++ ds4_gpu_tensor_free(q_view); ++ seq_done[b] = true; ++ all_multi = false; ++ } else if (ok) { ++ if (!ds4_gpu_attention_decode_multi_supported(n_comp, ++ n_comp ? metal_graph_attn_comp_cache_is_f16() : 0)) { ++ all_multi = false; ++ } ++ seqs[b].raw_kv = gb->layer_raw_cache[il]; ++ seqs[b].comp_kv = n_comp ? gb->layer_attn_comp_cache[il] : NULL; ++ seqs[b].comp_mask = NULL; ++ seqs[b].n_raw = n_raw; ++ seqs[b].raw_cap = gb->raw_cap; ++ seqs[b].raw_start = raw_start; ++ seqs[b].n_comp = n_comp; ++ } ++ ds4_gpu_tensor_free(qr_norm_view); ++ ds4_gpu_tensor_free(attn_norm_view); ++ ds4_gpu_tensor_free(kv_view); ++ } ++ if (ok && all_multi) { ++ ok = ds4_gpu_attention_decode_heads_multi_tensor(g->scr->batch_heads, ++ model->map, ++ model->size, ++ layer->attn_sinks->abs_offset, ++ g->scr->batch_q, ++ seqs, ++ n_seqs, ++ 0, ++ DS4_N_HEAD, ++ DS4_N_HEAD_DIM) != 0; ++ } else if (ok) { ++ for (uint32_t b = 0; ok && b < n_seqs; b++) { ++ if (seq_done[b]) continue; ++ ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(g->scr->batch_q, b, q_dim); ++ ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(g->scr->batch_heads, b, q_dim); ++ ok = q_view && heads_view && ++ ds4_gpu_attention_decode_heads_tensor(heads_view, ++ model->map, ++ model->size, ++ layer->attn_sinks->abs_offset, ++ q_view, ++ seqs[b].raw_kv, ++ seqs[b].n_raw, ++ seqs[b].raw_cap, ++ seqs[b].raw_start, ++ seqs[b].comp_kv, ++ metal_graph_attn_comp_cache_is_f16(), ++ seqs[b].n_comp, ++ NULL, ++ 0, ++ DS4_N_HEAD, ++ DS4_N_HEAD_DIM) != 0; ++ ds4_gpu_tensor_free(heads_view); ++ ds4_gpu_tensor_free(q_view); ++ } ++ } ++ ++ if (ok) ok = ds4_gpu_rope_tail_multi_tensor(g->scr->batch_heads, ++ n_seqs, ++ DS4_N_HEAD, ++ DS4_N_HEAD_DIM, ++ DS4_N_ROT, ++ pos, ++ rope_ctx, ++ true, ++ freq_base, ++ freq_scale, ++ ext_factor, ++ attn_factor, ++ DS4_ROPE_YARN_BETA_FAST, ++ DS4_ROPE_YARN_BETA_SLOW) != 0; ++ if (ok) ok = ds4_gpu_attention_output_q8_batch_tensor(g->scr->batch_attn_out, ++ g->scr->batch_attn_low, ++ g->scr->batch_group_tmp, ++ g->scr->batch_low_tmp, ++ model->map, ++ model->size, ++ layer->attn_output_a->abs_offset, ++ layer->attn_output_b->abs_offset, ++ group_dim, ++ rank, ++ n_groups, ++ DS4_N_EMBD, ++ g->scr->batch_heads, ++ n_seqs) != 0; ++ if (ok) ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, ++ g->scr->batch_attn_out, ++ g->scr->batch_cur_hc, ++ hc_split_view, ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ ds4_gpu_tensor_free(after_attn_hc_view); ++ ds4_gpu_tensor_free(attn_cur_view); ++ ds4_gpu_tensor_free(hc_split_view); ++ ds4_gpu_tensor_free(hc_mix_view); ++ return ok; ++} ++ ++/* Encode one decode token for each of n_seqs sequences in a single pass over ++ * the layers: attention via the batched-decode layer above, FFN and output ++ * head via the existing batch functions (graphs[0] carries the shared engine ++ * scratch; the FFN stages use no per-graph state). The per-sequence caches ++ * and counters advance exactly as n_seqs independent single-token evals ++ * would. Logits for sequence b land in scr->spec_logits row b. */ ++DS4_MAYBE_UNUSED static bool metal_graph_encode_decode_multi( ++ ds4_gpu_graph *const *graphs, ++ uint32_t n_seqs, ++ const ds4_model *model, ++ const ds4_weights *weights, ++ const int *tokens, ++ const uint32_t *pos, ++ bool need_logits) { ++ ds4_gpu_graph *g = graphs[0]; ++ if (n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX) return false; ++ if (metal_graph_directional_steering_attn_enabled(g)) return false; ++ for (uint32_t b = 0; b < n_seqs; b++) { ++ if (graphs[b]->raw_cap == 0 || graphs[b]->ssd_streaming) return false; ++ } ++ ++ const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; ++ int32_t token_ids[DS4_GPU_DECODE_MULTI_MAX]; ++ for (uint32_t b = 0; b < n_seqs; b++) token_ids[b] = tokens[b]; ++ bool ok = ds4_gpu_tensor_write(g->scr->prefill_tokens, 0, token_ids, ++ (uint64_t)n_seqs * sizeof(token_ids[0])) != 0; ++ for (uint32_t b = 0; ok && b < n_seqs; b++) { ++ ds4_gpu_tensor *hc_view = metal_graph_tensor_row_view(g->scr->batch_cur_hc, b, hc_dim); ++ ok = hc_view && ds4_gpu_embed_token_hc_tensor(hc_view, ++ model->map, ++ model->size, ++ weights->token_embd->abs_offset, ++ (uint32_t)weights->token_embd->dim[1], ++ (uint32_t)tokens[b], ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ ds4_gpu_tensor_free(hc_view); ++ } ++ ++ const uint32_t split_after_layers = metal_graph_token_split_after_layers(); ++ for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { ++ ok = metal_graph_encode_layer_attention_decode_multi(graphs, n_seqs, model, ++ &weights->layer[il], il, pos); ++ if (ok) ok = metal_graph_encode_layer_ffn_batch(g, model, &weights->layer[il], ++ il, pos[0], n_seqs); ++ if (ok) { ++ ds4_gpu_tensor *tmp = g->scr->batch_cur_hc; ++ g->scr->batch_cur_hc = g->scr->batch_next_hc; ++ g->scr->batch_next_hc = tmp; ++ } ++ if (ok && split_after_layers != 0 && il + 1u == split_after_layers) { ++ ok = ds4_gpu_flush_commands() != 0; ++ } ++ } ++ if (ok && need_logits) { ++ ok = metal_graph_encode_output_head_batch(g, model, weights, n_seqs, ++ weights->output->dim[1]); ++ } ++ return ok; ++} ++ + static bool metal_graph_eval_token_raw_swa_streaming( + ds4_gpu_graph *g, + const ds4_model *model, +diff --git a/ds4_cuda.cu b/ds4_cuda.cu +index c8a9e5a..ba609f7 100644 +--- a/ds4_cuda.cu ++++ b/ds4_cuda.cu +@@ -8970,6 +8970,10 @@ extern "C" int ds4_gpu_attention_decode_heads_tensor( + 0, 0, n_head, head_dim); + return cuda_ok(cudaGetLastError(), "attention decode launch"); + } ++extern "C" int ds4_gpu_attention_decode_multi_supported(uint32_t n_comp, uint32_t comp_kv_f16) { ++ return comp_kv_f16 == 0 && cuda_attention_score_buffer_fits(n_comp) ? 1 : 0; ++} ++ + extern "C" int ds4_gpu_attention_decode_heads_multi_tensor( + ds4_gpu_tensor *heads, + const void *model_map, +diff --git a/ds4_gpu.h b/ds4_gpu.h +index 7e41c4b..cab0cb2 100644 +--- a/ds4_gpu.h ++++ b/ds4_gpu.h +@@ -646,6 +646,13 @@ typedef struct { + uint32_t n_comp; + } ds4_gpu_attn_seqview; + ++/* True when the multi-sequence decode attention kernel can serve a sequence ++ * with n_comp visible compressed rows (score buffer bound, no online variant ++ * yet) and the compressed cache layout it supports. */ ++int ds4_gpu_attention_decode_multi_supported( ++ uint32_t n_comp, ++ uint32_t comp_kv_f16); ++ + int ds4_gpu_attention_decode_heads_multi_tensor( + ds4_gpu_tensor *heads, + const void *model_map, diff --git a/scratchpad/old_ffn_multi_commit.diff b/scratchpad/old_ffn_multi_commit.diff new file mode 100644 index 000000000..4a1adfc7e --- /dev/null +++ b/scratchpad/old_ffn_multi_commit.diff @@ -0,0 +1,287 @@ +commit e4e415dc646b22feb377a88fa849b088c4e65142 +Author: iCreil +Date: Tue Jul 7 20:51:56 2026 +0200 + + Graph: run batched-decode routed experts per sequence + + metal_graph_encode_layer_ffn_decode_multi replaces the prefill FFN batch + in the batched-decode driver. The dense stages (HC pre/norm, router, + shared expert, HC post) stay batched with n_tokens = n_seqs, where the + same weights serve every sequence. The routed experts now run one + sequence at a time through the single-token decode MoE kernel + (ds4_gpu_routed_moe_one_tensor) on batch-scratch row views: at batch 2-4 + the selected experts of different sequences are almost always distinct, + so the generic batch expert path reads them all without sharing, costing + more than n_seqs single-token passes. Token streams stay identical to + single decode (batched_decode_smoke MATCH). + +diff --git a/ds4.c b/ds4.c +index ba5b781..5dbf252 100644 +--- a/ds4.c ++++ b/ds4.c +@@ -19822,12 +19822,248 @@ static bool metal_graph_encode_layer_attention_decode_multi( + return ok; + } + ++/* Batched-decode FFN layer. The dense stages — HC pre/norm, router, shared ++ * expert, HC post — run once on the batch scratch with n_tokens = n_seqs, ++ * which is where batching pays (the same weights serve every sequence). The ++ * routed experts run per sequence through the single-token decode MoE kernel: ++ * at batch 2-4 the selected experts of different sequences are almost always ++ * distinct, so there is nothing to share, and the generic batch expert path ++ * costs far more than n_seqs single-token passes. ffn_out materialization ++ * (steering/debug) is not supported here; the driver gates steering off. */ ++static bool metal_graph_encode_layer_ffn_decode_multi( ++ ds4_gpu_graph *g, ++ uint32_t n_seqs, ++ const ds4_model *model, ++ const ds4_layer_weights *layer, ++ uint32_t il, ++ uint32_t pos0) { ++ if (n_seqs == 0 || n_seqs > DS4_GPU_DECODE_MULTI_MAX || n_seqs > g->prefill_cap) 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; ++ const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; ++ const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; ++ const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; ++ const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; ++ const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; ++ const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); ++ const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; ++ const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); ++ const uint64_t down_expert_bytes = routed_out_dim * down_row_bytes; ++ ++ ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( ++ g->scr->batch_hc_mix, 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); ++ ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( ++ g->scr->batch_hc_split, 0, (uint64_t)n_seqs * mix_hc * sizeof(float)); ++ ds4_gpu_tensor *ffn_cur_view = ds4_gpu_tensor_view( ++ g->scr->batch_ffn_cur, 0, (uint64_t)n_seqs * DS4_N_EMBD * sizeof(float)); ++ ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view( ++ g->scr->batch_next_hc, 0, (uint64_t)n_seqs * 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 && ++ !metal_graph_use_reference_hc_decode() && ++ metal_graph_enable_batch_hc_norm_fusion(); ++ ++ if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->scr->batch_flat_hc, ++ g->scr->batch_after_attn_hc, ++ (uint32_t)hc_dim, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ if (ok) ok = ds4_gpu_matmul_f16_tensor(hc_mix_view, ++ model->map, ++ model->size, ++ layer->hc_ffn_fn->abs_offset, ++ hc_dim, ++ mix_hc, ++ g->scr->batch_flat_hc, ++ n_seqs) != 0; ++ if (metal_graph_use_reference_hc_decode()) { ++ if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, ++ hc_mix_view, ++ model->map, ++ model->size, ++ layer->hc_ffn_scale->abs_offset, ++ layer->hc_ffn_base->abs_offset, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS) != 0; ++ if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(ffn_cur_view, ++ g->scr->batch_after_attn_hc, ++ hc_split_view, ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ } else if (fuse_hc_norm) { ++ if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(ffn_cur_view, ++ g->scr->batch_ffn_norm, ++ hc_split_view, ++ hc_mix_view, ++ g->scr->batch_after_attn_hc, ++ model->map, ++ model->size, ++ layer->hc_ffn_scale->abs_offset, ++ layer->hc_ffn_base->abs_offset, ++ layer->ffn_norm->abs_offset, ++ DS4_N_EMBD, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS, ++ DS4_RMS_EPS) != 0; ++ } else { ++ if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(ffn_cur_view, ++ hc_split_view, ++ hc_mix_view, ++ g->scr->batch_after_attn_hc, ++ model->map, ++ model->size, ++ layer->hc_ffn_scale->abs_offset, ++ layer->hc_ffn_base->abs_offset, ++ DS4_N_EMBD, ++ DS4_N_HC, ++ DS4_N_HC_SINKHORN_ITER, ++ DS4_HC_EPS) != 0; ++ } ++ if (ok && !fuse_hc_norm) { ++ ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_ffn_norm, ++ g->scr->batch_ffn_cur, ++ model->map, ++ model->size, ++ layer->ffn_norm->abs_offset, ++ DS4_N_EMBD, ++ n_seqs, ++ DS4_RMS_EPS) != 0; ++ } ++ ++ if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_router_logits, ++ model->map, ++ model->size, ++ layer->ffn_gate_inp->abs_offset, ++ DS4_N_EMBD, ++ DS4_N_EXPERT, ++ g->scr->batch_ffn_norm, ++ n_seqs) != 0; ++ if (ok) ok = ds4_gpu_router_select_batch_tensor(g->scr->batch_router_selected, ++ g->scr->batch_router_weights, ++ g->scr->batch_router_probs, ++ model->map, ++ model->size, ++ layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, ++ layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, ++ layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, ++ 0, ++ 0, ++ layer->ffn_exp_probs_b != NULL, ++ layer->ffn_gate_tid2eid != NULL, ++ g->scr->batch_router_logits, ++ g->scr->prefill_tokens, ++ DS4_N_EXPERT, ++ DS4_N_EXPERT_USED, ++ DS4_EXPERT_WEIGHT_SCALE, ++ n_seqs) != 0; ++ ++ /* Shared expert, batched (same weights for every sequence). */ ++ bool shared_down_f16 = false; ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_gate", il, pos0, ++ g->scr->batch_shared_gate, ++ model, layer->ffn_gate_shexp, ++ DS4_N_EMBD, shared_dim, ++ g->scr->batch_ffn_norm, n_seqs); ++ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_up", il, pos0, ++ g->scr->batch_shared_up, ++ model, layer->ffn_up_shexp, ++ DS4_N_EMBD, shared_dim, ++ g->scr->batch_ffn_norm, n_seqs); ++ if (ok) ok = ds4_gpu_swiglu_tensor(g->scr->batch_shared_mid, ++ g->scr->batch_shared_gate, ++ g->scr->batch_shared_up, ++ (uint32_t)((uint64_t)n_seqs * shared_dim), ++ DS4_SWIGLU_CLAMP_EXP, ++ 1.0f) != 0; ++ if (ok) { ++ shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->scr->batch_q_half, ++ model->map, ++ model->size, ++ layer->ffn_down_shexp->abs_offset, ++ shared_dim, ++ DS4_N_EMBD, ++ g->scr->batch_shared_mid, ++ n_seqs) != 0; ++ } ++ if (ok && !shared_down_f16) { ++ ok = metal_graph_matmul_q8_0_named_tensor("shared_down", il, pos0, ++ g->scr->batch_shared_out, ++ model, layer->ffn_down_shexp, ++ shared_dim, DS4_N_EMBD, ++ g->scr->batch_shared_mid, n_seqs); ++ } ++ ++ /* Routed experts, one sequence at a time through the decode kernel. */ ++ for (uint32_t b = 0; ok && b < n_seqs; b++) { ++ ds4_gpu_tensor *routed_out_row = metal_graph_tensor_row_view( ++ g->scr->batch_routed_out, b, DS4_N_EMBD); ++ ds4_gpu_tensor *ffn_norm_row = metal_graph_tensor_row_view( ++ g->scr->batch_ffn_norm, b, DS4_N_EMBD); ++ ds4_gpu_tensor *selected_row = metal_graph_tensor_row_view( ++ g->scr->batch_router_selected, b, DS4_N_EXPERT_USED); ++ ds4_gpu_tensor *weights_row = metal_graph_tensor_row_view( ++ g->scr->batch_router_weights, b, DS4_N_EXPERT_USED); ++ ok = routed_out_row && ffn_norm_row && selected_row && weights_row && ++ ds4_gpu_routed_moe_one_tensor(routed_out_row, ++ g->scr->routed_gate, ++ g->scr->routed_up, ++ g->scr->routed_mid, ++ g->scr->routed_down, ++ model->map, model->size, ++ layer->ffn_gate_exps->abs_offset, ++ layer->ffn_up_exps->abs_offset, ++ layer->ffn_down_exps->abs_offset, ++ layer->ffn_gate_exps->type, ++ layer->ffn_down_exps->type, ++ gate_expert_bytes, gate_row_bytes, ++ down_expert_bytes, down_row_bytes, ++ (uint32_t)expert_in_dim, ++ (uint32_t)down_in_dim, ++ (uint32_t)routed_out_dim, ++ selected_row, weights_row, ++ DS4_N_EXPERT, ++ DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, ++ ffn_norm_row, ++ il) != 0; ++ ds4_gpu_tensor_free(weights_row); ++ ds4_gpu_tensor_free(selected_row); ++ ds4_gpu_tensor_free(ffn_norm_row); ++ ds4_gpu_tensor_free(routed_out_row); ++ } ++ ++ if (ok && shared_down_f16) { ++ ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, ++ g->scr->batch_routed_out, ++ g->scr->batch_q_half, ++ g->scr->batch_after_attn_hc, ++ hc_split_view, ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ } else if (ok) { ++ ok = ds4_gpu_hc_expand_add_split_tensor(next_hc_view, ++ g->scr->batch_routed_out, ++ g->scr->batch_shared_out, ++ g->scr->batch_after_attn_hc, ++ hc_split_view, ++ DS4_N_EMBD, ++ DS4_N_HC) != 0; ++ } ++ ds4_gpu_tensor_free(next_hc_view); ++ ds4_gpu_tensor_free(ffn_cur_view); ++ ds4_gpu_tensor_free(hc_split_view); ++ ds4_gpu_tensor_free(hc_mix_view); ++ return ok; ++} ++ + /* Encode one decode token for each of n_seqs sequences in a single pass over +- * the layers: attention via the batched-decode layer above, FFN and output +- * head via the existing batch functions (graphs[0] carries the shared engine +- * scratch; the FFN stages use no per-graph state). The per-sequence caches +- * and counters advance exactly as n_seqs independent single-token evals +- * would. Logits for sequence b land in scr->spec_logits row b. */ ++ * the layers: attention and FFN via the batched-decode layer halves above, ++ * output head via the existing batch function (graphs[0] carries the shared ++ * engine scratch; the FFN stages use no per-graph state). The per-sequence ++ * caches and counters advance exactly as n_seqs independent single-token ++ * evals would. Logits for sequence b land in scr->spec_logits row b. */ + static bool metal_graph_encode_decode_multi( + ds4_gpu_graph *const *graphs, + uint32_t n_seqs, +@@ -19874,8 +20110,9 @@ static bool metal_graph_encode_decode_multi( + fprintf(stderr, "ds4: gpu layer %u batched decode attention failed\n", il); + } + if (ok) { +- ok = metal_graph_encode_layer_ffn_batch(g, model, &weights->layer[il], +- il, pos[0], n_seqs); ++ ok = metal_graph_encode_layer_ffn_decode_multi(g, n_seqs, model, ++ &weights->layer[il], ++ il, pos[0]); + if (!ok) { + fprintf(stderr, "ds4: gpu layer %u batched decode ffn failed\n", il); + } diff --git a/scratchpad/old_helper_full.txt b/scratchpad/old_helper_full.txt new file mode 100644 index 000000000..c62203687 --- /dev/null +++ b/scratchpad/old_helper_full.txt @@ -0,0 +1,388 @@ +static bool metal_graph_decode_compressor_indexer( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + ds4_gpu_tensor *attn_norm, + ds4_gpu_tensor *qr_norm, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + uint32_t *out_n_comp, + ds4_gpu_tensor **out_comp_cache, + ds4_gpu_tensor **out_comp_selected, + uint32_t *out_n_selected, + double *out_index_stage_t0) { + const bool compressed = ds4_layer_compress_ratio(il) != 0; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + bool ok = true; + uint32_t n_comp = 0; + ds4_gpu_tensor *comp_cache = NULL; + ds4_gpu_tensor *comp_selected = NULL; + uint32_t n_selected = 0; + double decode_index_stage_t0 = 0.0; + const bool decode_index_stage_profile = getenv("DS4_METAL_INDEXER_STAGE_PROFILE") != NULL; + if (compressed) { + const uint32_t ratio = ds4_layer_compress_ratio(il); + const uint32_t coff = ratio == 4 ? 2u : 1u; + const uint32_t comp_width = coff * DS4_N_HEAD_DIM; + const bool emit = ((pos + 1u) % ratio) == 0u; + if (!layer->attn_compressor_kv || !layer->attn_compressor_gate || + !layer->attn_compressor_ape || !layer->attn_compressor_norm || + layer->attn_compressor_kv->type != DS4_TENSOR_F16 || + layer->attn_compressor_gate->type != DS4_TENSOR_F16 || + layer->attn_compressor_kv->dim[0] != DS4_N_EMBD || + layer->attn_compressor_gate->dim[0] != DS4_N_EMBD || + layer->attn_compressor_kv->dim[1] != comp_width || + layer->attn_compressor_gate->dim[1] != comp_width) { + fprintf(stderr, "ds4: Metal graph compressor expects paired F16 compressor projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + if (ok && !metal_graph_use_reference_compressor_pair_proj()) { + ok = ds4_gpu_matmul_f16_pair_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, + comp_width, + attn_norm, + 1) != 0; + } else { + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_kv_cur, model->map, model->size, + layer->attn_compressor_kv->abs_offset, + DS4_N_EMBD, comp_width, + attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_sc_cur, model->map, model->size, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, comp_width, + attn_norm, 1) != 0; + } + const uint32_t comp_row = g->layer_n_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + metal_graph_attn_comp_update_target(g, il), + model->map, + model->size, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + layer->attn_compressor_norm->abs_offset, + layer->attn_compressor_norm->type, + DS4_N_HEAD_DIM, + ratio, + pos, + metal_graph_attn_comp_update_row(comp_row), + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + if (ok && emit) { + ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); + if (!comp_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_row_view, 1, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; + if (ok) { + metal_graph_debug_dump_tensor("KVcompress", comp_row_view, DS4_N_HEAD_DIM, il, pos); + } + ds4_gpu_tensor_free(comp_row_view); + } + if (ok) ok = metal_graph_commit_attn_comp_stage(g, il, comp_row, 1); + } + if (ok && emit) g->layer_n_comp[il]++; + + if (ok && ratio == 4) { + const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || + !layer->indexer_compressor_ape || !layer->indexer_compressor_norm || + layer->indexer_compressor_kv->type != DS4_TENSOR_F16 || + layer->indexer_compressor_gate->type != DS4_TENSOR_F16 || + layer->indexer_compressor_kv->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_gate->dim[0] != DS4_N_EMBD || + layer->indexer_compressor_kv->dim[1] != index_width || + layer->indexer_compressor_gate->dim[1] != index_width) { + fprintf(stderr, "ds4: Metal graph indexer compressor expects paired F16 projections\n"); + ok = false; + } + if (ok && emit && g->layer_n_index_comp[il] >= g->layer_comp_cap[il]) { + fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); + ok = false; + } + if (ok && !metal_graph_use_reference_compressor_pair_proj()) { + ok = ds4_gpu_matmul_f16_pair_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, + index_width, + attn_norm, + 1) != 0; + } else { + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_kv_cur, model->map, model->size, + layer->indexer_compressor_kv->abs_offset, + DS4_N_EMBD, index_width, + attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_sc_cur, model->map, model->size, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, index_width, + attn_norm, 1) != 0; + } + const uint32_t index_row = g->layer_n_index_comp[il]; + if (ok) ok = ds4_gpu_compressor_update_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + g->layer_index_comp_cache[il], + model->map, + model->size, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + layer->indexer_compressor_norm->abs_offset, + layer->indexer_compressor_norm->type, + DS4_N_INDEXER_HEAD_DIM, + ratio, + pos, + index_row, + DS4_N_ROT, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + if (ok && emit) { + ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( + g->layer_index_comp_cache[il], + (uint64_t)index_row * DS4_N_INDEXER_HEAD_DIM * sizeof(float), + (uint64_t)DS4_N_INDEXER_HEAD_DIM * sizeof(float)); + if (!index_row_view) { + ok = false; + } else { + ok = ds4_gpu_dsv4_indexer_qat_tensor(index_row_view, + 1, + DS4_N_INDEXER_HEAD_DIM) != 0; + ds4_gpu_tensor_free(index_row_view); + } + } + if (ok && emit) g->layer_n_index_comp[il]++; + const uint32_t decode_sparse_threshold = + metal_graph_decode_indexer_sparse_threshold(g); + if (ok && + g->layer_n_comp[il] > decode_sparse_threshold && + g->layer_n_index_comp[il] > DS4_N_INDEXER_TOP_K) { + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + if (!layer->indexer_attn_q_b || + !tensor_type_is_f16_or_q8_0(layer->indexer_attn_q_b->type) || + layer->indexer_attn_q_b->dim[0] != q_rank || + layer->indexer_attn_q_b->dim[1] != indexer_q_dim) { + fprintf(stderr, "ds4: Metal graph indexer q projection expects F16 or Q8_0 weights\n"); + ok = false; + } + if (ok && (!layer->indexer_proj || + layer->indexer_proj->type != DS4_TENSOR_F16 || + layer->indexer_proj->dim[0] != DS4_N_EMBD || + layer->indexer_proj->dim[1] != DS4_N_INDEXER_HEAD)) { + fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); + ok = false; + } + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->indexer_q, + model, + layer->indexer_attn_q_b, + q_rank, + indexer_q_dim, + qr_norm, + 1); + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->indexer_q, 1, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + DS4_N_ROT, + pos, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW) != 0; + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(g->scr->indexer_q, + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->indexer_weights, model->map, model->size, + layer->indexer_proj->abs_offset, + DS4_N_EMBD, DS4_N_INDEXER_HEAD, + attn_norm, 1) != 0; + const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary(NULL, + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_score_one_tensor(g->scr->indexer_scores, + g->scr->indexer_q, + g->scr->indexer_weights, + g->layer_index_comp_cache[il], + g->layer_n_index_comp[il], + DS4_N_INDEXER_HEAD, + DS4_N_INDEXER_HEAD_DIM, + index_scale) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_score", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + if (ok) ok = ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->indexer_scores, + g->layer_n_index_comp[il], + 1, + DS4_N_INDEXER_TOP_K) != 0; + if (ok && decode_index_stage_profile) { + ok = metal_graph_indexer_stage_profile_boundary("decode_topk", + il, + pos, + 1, + g->layer_n_index_comp[il], + &decode_index_stage_t0); + } + /* Decode used to materialize a dense compressed-row mask and + * call the generic gathered FlashAttention wrapper below. + * That wrapper scans every compressed row and rejects long + * contexts once raw+compressed rows exceed 8192. Ratio-4 DS4 + * attention is sparse after indexer top-k, so use the private + * indexed attention kernel instead: it scans only SWA raw rows + * plus the selected compressed rows, matching prefill and + * avoiding the long-context decode failure. */ + if (ok) { + comp_selected = g->scr->comp_selected; + /* + * Contract: the indexer top-k is fixed by the model config + * and must remain the full 512 rows. Do not reduce this for + * throughput benchmarks. + * + * Why: the indexer is not just an implementation detail. It + * decides which compressed memory rows are visible to the + * attention kernel. If we keep only 128/256 rows, the later + * indexed-attention math may be perfectly computed, but it is + * computed over the wrong candidate set: rows ranked 257-512 + * are removed before softmax/PV can use them. Those rows may + * carry weak-but-necessary evidence for retrieval, name/number + * recall, or long-context disambiguation. The error is + * therefore semantic/algorithmic, not the acceptable kind of + * local numerical drift caused by a different reduction order + * or Tensor/NAX precision. + * + * Short prompt tests, first-token agreement, or even a small + * official-vector set can miss this because many prompts do + * not need the tail of the 512 selected compressed rows. The + * failure appears only when the model needs information that + * fell below the reduced cutoff. Optimizations belong inside + * the score/top-k/attention implementation while preserving + * DS4_N_INDEXER_TOP_K. + */ + n_selected = DS4_N_INDEXER_TOP_K < g->layer_n_index_comp[il] + ? DS4_N_INDEXER_TOP_K + : g->layer_n_index_comp[il]; + } + } + } + + n_comp = g->layer_n_comp[il]; + comp_cache = g->layer_attn_comp_cache[il]; + } + *out_n_comp = n_comp; + *out_comp_cache = comp_cache; + *out_comp_selected = comp_selected; + *out_n_selected = n_selected; + *out_index_stage_t0 = decode_index_stage_t0; + return ok; +} + +static bool metal_graph_encode_decode_layer( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t raw_row, + uint32_t n_raw, + int token) { + 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; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint32_t n_groups = DS4_N_OUT_GROUP; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = DS4_N_LORA_O; + const uint32_t shared_dim = (uint32_t)layer->ffn_gate_shexp->dim[1]; + const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0]; + const uint64_t expert_mid_dim = layer->ffn_gate_exps->dim[1]; + const uint64_t down_in_dim = layer->ffn_down_exps->dim[0]; + const uint64_t routed_out_dim = layer->ffn_down_exps->dim[1]; + const bool compressed = ds4_layer_compress_ratio(il) != 0; + const float freq_base = layer_rope_freq_base(il); + const float freq_scale = layer_rope_freq_scale(il); + const float ext_factor = compressed && DS4_ROPE_SCALE_FACTOR > 1.0f ? 1.0f : 0.0f; + float attn_factor = 1.0f; + if (ext_factor != 0.0f && freq_scale > 0.0f) { + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); + + bool ok = true; + const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); + double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; +#define DS4_METAL_PROFILE_DECODE_STAGE(name) do { \ + if (ok && decode_stage_profile) { \ + ok = metal_graph_layer_stage_profile_boundary("decode", (name), il, pos, 1, &decode_stage_t0); \ + } \ + } while (0) + if (ok) ok = ds4_gpu_rms_norm_plain_tensor(g->scr->flat_hc, g->scr->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->hc_mix, model, layer->hc_attn_fn, + hc_dim, mix_hc, g->scr->flat_hc, 1); + const bool fuse_hc_norm = + DS4_N_HC == 4 && + !metal_graph_use_reference_hc_decode() && + !metal_graph_use_reference_hc_norm_decode(); + if (ok && fuse_hc_norm) { + ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(g->scr->attn_cur, + g->scr->attn_norm, + g->scr->hc_split, + g->scr->hc_mix, + g->scr->cur_hc, + model->map, + model->size, + layer->hc_attn_scale->abs_offset, + layer->hc_attn_base->abs_offset, + layer->attn_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_N_HC_SINKHORN_ITER, + DS4_HC_EPS, + DS4_RMS_EPS) != 0; From 1efb3f7e234158a3900953eed4b884137e532504 Mon Sep 17 00:00:00 2001 From: iCreil Date: Fri, 24 Jul 2026 23:47:17 +0200 Subject: [PATCH 07/10] Graph: run fused-decode routed experts per sequence In the fused session-batch decode each FFN batch row is one token of a distinct sequence, so the selected experts are almost always disjoint and the generic batch expert path costs more than n_rows single-token passes. Route the routed-MoE dispatch of metal_graph_encode_layer_ffn_batch through ds4_gpu_routed_moe_one_tensor per row while the fused driver is on the stack (file-scope flag; the graph worker serializes encodes). DS4_FUSED_BATCH_MOE=1 keeps the batch expert path for comparison. --- ds4.c | 63 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/ds4.c b/ds4.c index 20681f265..67dbe897f 100644 --- a/ds4.c +++ b/ds4.c @@ -28417,6 +28417,13 @@ static bool metal_graph_encode_mixed_routed_rows( * experts, sum, and HC post. A non-empty decode tail has already been * prepared in rows [n_tokens, n_tokens + decode_count); only the routed * expert dispatch is shared between the two arithmetic paths. */ +/* Set by the fused session-batch decode driver around its FFN calls: each + * batch row is one decode token of a distinct sequence, so the routed experts + * of different rows are almost always disjoint and the generic batch expert + * path costs more than n_rows single-token passes. The graph worker + * serializes every encode, so a file-scope flag is safe. */ +static bool metal_graph_fused_decode_moe_per_seq = false; + static bool metal_graph_encode_layer_ffn_batch( ds4_gpu_graph *g, const ds4_model *model, @@ -28878,6 +28885,59 @@ static bool metal_graph_encode_layer_ffn_batch( g->tp_batch_in[il], (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; } + } else if (ok && metal_graph_fused_decode_moe_per_seq && !g->ssd_streaming) { + /* Fused session-batch decode: routed experts one row (= one sequence) + * at a time through the single-token decode MoE kernel. At 2-4 rows + * the selected experts are almost always distinct, so there is + * nothing to share and the batch expert path costs far more than + * n_tokens single-token passes. */ + const uint64_t vec_bytes = (uint64_t)DS4_N_EMBD * sizeof(float); + for (uint32_t r = 0; ok && r < n_tokens; r++) { + ds4_gpu_tensor *out_row = ds4_gpu_tensor_view( + metal_graph_batch_routed_out(g), (uint64_t)r * vec_bytes, vec_bytes); + ds4_gpu_tensor *x_row = ds4_gpu_tensor_view( + metal_graph_batch_ffn_norm(g), (uint64_t)r * vec_bytes, vec_bytes); + ds4_gpu_tensor *sel_row = ds4_gpu_tensor_view( + metal_graph_batch_router_selected(g), + (uint64_t)r * DS4_N_EXPERT_USED * sizeof(int32_t), + (uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + ds4_gpu_tensor *w_row = ds4_gpu_tensor_view( + metal_graph_batch_router_weights(g), + (uint64_t)r * DS4_N_EXPERT_USED * sizeof(float), + (uint64_t)DS4_N_EXPERT_USED * sizeof(float)); + ok = out_row && x_row && sel_row && w_row && + ds4_gpu_routed_moe_one_tensor(out_row, + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + (uint32_t)expert_in_dim, + (uint32_t)down_in_dim, + (uint32_t)routed_out_dim, + sel_row, w_row, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_SWIGLU_CLAMP_EXP, + x_row, + NULL, + il, + false) != 0; + ds4_gpu_tensor_free(w_row); + ds4_gpu_tensor_free(sel_row); + ds4_gpu_tensor_free(x_row); + ds4_gpu_tensor_free(out_row); + } + g->batch_routed_mid_is_f16 = false; } else if (ok) { ok = ds4_gpu_routed_moe_batch_tensor(metal_graph_batch_routed_out(g), metal_graph_batch_routed_gate(g), @@ -29874,8 +29934,11 @@ static bool metal_graph_encode_decode_multi( fprintf(stderr, "ds4: gpu layer %u batched decode attention failed\n", il); } if (ok) { + metal_graph_fused_decode_moe_per_seq = + getenv("DS4_FUSED_BATCH_MOE") == NULL; ok = metal_graph_encode_layer_ffn_batch(g, model, &weights->layer[il], il, pos[0], n_seqs, NULL, 0); + metal_graph_fused_decode_moe_per_seq = false; if (!ok) { fprintf(stderr, "ds4: gpu layer %u batched decode ffn failed\n", il); } From a31dd223cbdaf0fab0fb3d33e3fa2b7dcf74627b Mon Sep 17 00:00:00 2001 From: iCreil Date: Sat, 25 Jul 2026 00:40:22 +0200 Subject: [PATCH 08/10] Server: reserve weight-cache VRAM for resident sessions With --batched-session N the lazy Q8->F16 weight caches warm up during engine creation and grow until only the default reserve (~1% of VRAM) stays free, so the resident sessions allocated afterwards fail even at modest context sizes (on a 96 GB card even one 32k session could not be created). Derive the reserve from the per-session context-buffer estimate times N plus the shared prefill workspace, and export it before engine creation; an explicit DS4_CUDA_Q8_F16_CACHE_RESERVE_MB wins. --- ds4_server.c | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/ds4_server.c b/ds4_server.c index f4f563ecb..291b51358 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -12901,6 +12901,34 @@ int main(int argc, char **argv) { cfg.engine.context_size = cfg.ctx_size; cfg.engine.placement_ctx_hint = cfg.ctx_size; cfg.engine.share_session_prefill_workspace = cfg.batched_sessions > 0; + /* Resident sessions allocate their context buffers after the engine has + * warmed its lazy Q8->F16 weight caches, and those caches grow until only + * the default reserve (~1% of VRAM) stays free — leaving no room for the + * sessions and failing creation even at modest ctx sizes. Raise the + * reserve so cache growth leaves space for every resident session, the + * shared prefill workspace and some slack. An explicit + * DS4_CUDA_Q8_F16_CACHE_RESERVE_MB in the environment always wins. */ + if (cfg.batched_sessions > 0) { + const uint32_t reserve_pc = + cfg.engine.prefill_chunk ? cfg.engine.prefill_chunk : 4096u; + const ds4_context_memory rm = + ds4_context_memory_estimate_with_prefill_mode(cfg.engine.backend, + cfg.ctx_size, + reserve_pc, + cfg.engine.ssd_streaming); + const uint64_t reserve_mib = 768ull + + (rm.total_bytes / (1024ull * 1024ull) + 64ull) * + (uint64_t)cfg.batched_sessions + + (uint64_t)reserve_pc + 256ull; + char reserve_buf[32]; + snprintf(reserve_buf, sizeof(reserve_buf), "%llu", + (unsigned long long)reserve_mib); + setenv("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", reserve_buf, 0); + server_log(DS4_LOG_DEFAULT, + "ds4-server: reserving %llu MiB of VRAM from the weight cache for %d resident sessions", + (unsigned long long)reserve_mib, + cfg.batched_sessions); + } ds4_engine *engine = NULL; if (cfg.gpu_vram_arg || cfg.gpu_devices_arg) { ds4_gpu_config gpu_cfg = {0}; From 4b623998349b340a37e7ca7dcaca2cb8286df0b6 Mon Sep 17 00:00:00 2001 From: iCreil Date: Sat, 25 Jul 2026 00:44:41 +0200 Subject: [PATCH 09/10] Server: estimate the session reserve with the CUDA layout and wider margins --- ds4_server.c | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/ds4_server.c b/ds4_server.c index 291b51358..df57adc5d 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -12911,15 +12911,19 @@ int main(int argc, char **argv) { if (cfg.batched_sessions > 0) { const uint32_t reserve_pc = cfg.engine.prefill_chunk ? cfg.engine.prefill_chunk : 4096u; + /* Estimate with the CUDA layout regardless of the backend configured + * so far (the final backend may be assigned later during engine + * setup): only the CUDA cache reads this env, and the CUDA estimate + * is the largest. */ const ds4_context_memory rm = - ds4_context_memory_estimate_with_prefill_mode(cfg.engine.backend, + ds4_context_memory_estimate_with_prefill_mode(DS4_BACKEND_CUDA, cfg.ctx_size, reserve_pc, cfg.engine.ssd_streaming); const uint64_t reserve_mib = 768ull + - (rm.total_bytes / (1024ull * 1024ull) + 64ull) * + (rm.total_bytes / (1024ull * 1024ull) + 128ull) * (uint64_t)cfg.batched_sessions + - (uint64_t)reserve_pc + 256ull; + 512ull; char reserve_buf[32]; snprintf(reserve_buf, sizeof(reserve_buf), "%llu", (unsigned long long)reserve_mib); From ee269d54fa060664bf49352bc4bc0397258b37b3 Mon Sep 17 00:00:00 2001 From: iCreil Date: Sat, 25 Jul 2026 00:50:39 +0200 Subject: [PATCH 10/10] Engine: reserve weight-cache VRAM for resident sessions at open Move the reserve computation into engine open, right before the startup cache warm: the context-memory estimate needs the model shape, which is only loaded here (the earlier server-side attempt underestimated it and still failed session creation). The server passes the resident session count via the new resident_sessions config field. --- ds4.c | 26 ++++++++++++++++++++++++++ ds4.h | 4 ++++ ds4_server.c | 33 +-------------------------------- 3 files changed, 31 insertions(+), 32 deletions(-) diff --git a/ds4.c b/ds4.c index 67dbe897f..b9768fd92 100644 --- a/ds4.c +++ b/ds4.c @@ -56838,6 +56838,32 @@ static int ds4_engine_open_internal(ds4_engine **out, return 1; } (void)ds4_gpu_set_model_fd_for_map(e->model.fd, e->model.map); + /* The lazy Q8->F16 weight caches warmed below grow until only the + * default reserve (~1% of VRAM) stays free, but the batched server + * creates its resident sessions after this point and their context + * buffers must still fit. The model shape is loaded here, so the + * per-session estimate is exact. An explicit + * DS4_CUDA_Q8_F16_CACHE_RESERVE_MB in the environment always wins. */ + if (opt->resident_sessions > 0) { + const uint32_t reserve_pc = + opt->prefill_chunk ? opt->prefill_chunk : 4096u; + const ds4_context_memory rm = + ds4_context_memory_estimate_with_prefill_mode( + e->backend, opt->context_size, reserve_pc, + opt->ssd_streaming); + const uint64_t reserve_mib = 768ull + + (rm.total_bytes / (1024ull * 1024ull) + 128ull) * + (uint64_t)opt->resident_sessions + + 512ull; + char reserve_buf[32]; + snprintf(reserve_buf, sizeof(reserve_buf), "%llu", + (unsigned long long)reserve_mib); + setenv("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", reserve_buf, 0); + fprintf(stderr, + "ds4: reserving %llu MiB of VRAM from the weight cache for %d resident sessions\n", + (unsigned long long)reserve_mib, + opt->resident_sessions); + } if (!accelerator_cache_model_tensors(e->backend, &e->model, load_offsets, load_sizes, load_span_count)) { diff --git a/ds4.h b/ds4.h index 42bc11660..bea8b0382 100644 --- a/ds4.h +++ b/ds4.h @@ -160,6 +160,10 @@ typedef struct { int placement_ctx_hint; /* Server batch mode serializes execution and can share prefill scratch. */ bool share_session_prefill_workspace; + /* Number of resident sessions the batched server will create after the + * engine opens; used to keep that much VRAM out of the lazy weight + * caches so session creation cannot fail. 0 = single-session server. */ + int resident_sessions; bool first_token_test; bool metal_graph_test; bool load_slice; diff --git a/ds4_server.c b/ds4_server.c index df57adc5d..3940959fa 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -12901,38 +12901,7 @@ int main(int argc, char **argv) { cfg.engine.context_size = cfg.ctx_size; cfg.engine.placement_ctx_hint = cfg.ctx_size; cfg.engine.share_session_prefill_workspace = cfg.batched_sessions > 0; - /* Resident sessions allocate their context buffers after the engine has - * warmed its lazy Q8->F16 weight caches, and those caches grow until only - * the default reserve (~1% of VRAM) stays free — leaving no room for the - * sessions and failing creation even at modest ctx sizes. Raise the - * reserve so cache growth leaves space for every resident session, the - * shared prefill workspace and some slack. An explicit - * DS4_CUDA_Q8_F16_CACHE_RESERVE_MB in the environment always wins. */ - if (cfg.batched_sessions > 0) { - const uint32_t reserve_pc = - cfg.engine.prefill_chunk ? cfg.engine.prefill_chunk : 4096u; - /* Estimate with the CUDA layout regardless of the backend configured - * so far (the final backend may be assigned later during engine - * setup): only the CUDA cache reads this env, and the CUDA estimate - * is the largest. */ - const ds4_context_memory rm = - ds4_context_memory_estimate_with_prefill_mode(DS4_BACKEND_CUDA, - cfg.ctx_size, - reserve_pc, - cfg.engine.ssd_streaming); - const uint64_t reserve_mib = 768ull + - (rm.total_bytes / (1024ull * 1024ull) + 128ull) * - (uint64_t)cfg.batched_sessions + - 512ull; - char reserve_buf[32]; - snprintf(reserve_buf, sizeof(reserve_buf), "%llu", - (unsigned long long)reserve_mib); - setenv("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", reserve_buf, 0); - server_log(DS4_LOG_DEFAULT, - "ds4-server: reserving %llu MiB of VRAM from the weight cache for %d resident sessions", - (unsigned long long)reserve_mib, - cfg.batched_sessions); - } + cfg.engine.resident_sessions = cfg.batched_sessions; ds4_engine *engine = NULL; if (cfg.gpu_vram_arg || cfg.gpu_devices_arg) { ds4_gpu_config gpu_cfg = {0};