diff --git a/ds4.c b/ds4.c index 2496319ace..c392ed56bc 100644 --- a/ds4.c +++ b/ds4.c @@ -32157,6 +32157,10 @@ static bool dspark_apply_markov_confidence_lazy_runtime( (uint64_t)draft * logits_bytes, logits_bytes); uint64_t gpu_key = 0; + /* reverse the w2 scan per draft so consecutive steps hit the + * L3-resident tail of the previous step (measured -8% on + * gfx1151; argmax output is direction-independent) */ + ds4_gpu_dspark_markov_set_reverse((int)(draft & 1u)); bool gpu_ok = row_view && ds4_gpu_dspark_markov_argmax_tensor( g->dspark_draft_tokens, diff --git a/ds4_cuda.cu b/ds4_cuda.cu index aaa4df1134..aff274f2b7 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -11697,6 +11697,12 @@ extern "C" int ds4_gpu_indexer_scores_decode_batch_tensor( n_head, head_dim, ratio, scale, 1); } +static int g_cuda_markov_reverse = 0; + +extern "C" void ds4_gpu_dspark_markov_set_reverse(int rev) { + g_cuda_markov_reverse = rev; +} + extern "C" int ds4_gpu_dspark_markov_argmax_tensor( ds4_gpu_tensor *out_idx, const ds4_gpu_tensor *logits_row, diff --git a/ds4_gpu.h b/ds4_gpu.h index 2000bba8bd..f2849da2d7 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -444,6 +444,10 @@ int ds4_gpu_indexer_scores_decode_batch_tensor( uint32_t ratio, float scale); +/* Alternates the w2 scan direction between chain steps (L3 tail reuse; + * oracle-safe: argmax is direction-independent). Default 0 = forward. */ +void ds4_gpu_dspark_markov_set_reverse(int rev); + int ds4_gpu_dspark_markov_argmax_tensor(ds4_gpu_tensor *out_idx, const ds4_gpu_tensor *logits_row, const void *model_map, diff --git a/ds4_rocm_unavailable.cu b/ds4_rocm_unavailable.cu index 7169eabb52..311c8004da 100644 --- a/ds4_rocm_unavailable.cu +++ b/ds4_rocm_unavailable.cu @@ -9,7 +9,6 @@ ROCM_UNAVAILABLE_INT(ds4_gpu_add_xdev_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_attention_decode_rows_rope_tensor) -ROCM_UNAVAILABLE_INT(ds4_gpu_attention_noncausal_raw_batch_heads_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_low_q4_K_slice_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_low_q8_rows_exact_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_attention_output_q4_K_batch_tensor) @@ -17,7 +16,6 @@ ROCM_UNAVAILABLE_INT(ds4_gpu_attention_prefill_raw_heads_range_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_attention_prefill_static_mixed_heads_range_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_device_cache_support_tensors) ROCM_UNAVAILABLE_INT(ds4_gpu_device_cache_tensors) -ROCM_UNAVAILABLE_INT(ds4_gpu_dspark_markov_argmax_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_indexer_top1_value_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_kv_fp8_store_raw_decode_rows_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor) diff --git a/rocm/ds4_rocm_attention_launch.cuh b/rocm/ds4_rocm_attention_launch.cuh index 56563b4506..3a0d806f4c 100644 --- a/rocm/ds4_rocm_attention_launch.cuh +++ b/rocm/ds4_rocm_attention_launch.cuh @@ -324,6 +324,187 @@ extern "C" int ds4_gpu_attention_decode_raw_batch_heads_tensor( n_head, head_dim); } +/* Non-causal batch attention over a raw KV ring for the DSpark draft block + * (ROCm port of the CUDA kernel). Every query row attends over all n_raw + * visible rows plus the per-head sink, with the same one-block + * max/denominator/value accumulation order as the reference decode + * attention (scores in shared, sequential value pass). */ +__global__ static void attention_noncausal_raw_batch_heads_kernel( + float *heads, + const float *sinks, + const float *q, + const float *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t tok = blockIdx.x; + const uint32_t h = blockIdx.y; + if (tok >= n_tokens || h >= n_head) return; + extern __shared__ float sh_scores[]; /* n_raw floats */ + const float *qh = q + ((uint64_t)tok * n_head + h) * head_dim; + const float scale = rsqrtf((float)head_dim); + for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { + const uint32_t row = (raw_start + r) % raw_cap; + const float *kv = raw_kv + (uint64_t)row * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d]; + sh_scores[r] = dot * scale; + } + __syncthreads(); + __shared__ float partial[256]; + __shared__ float max_s; + __shared__ float denom; + float local_max = sinks[h]; + for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) { + local_max = fmaxf(local_max, sh_scores[r]); + } + partial[threadIdx.x] = local_max; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + 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 r = threadIdx.x; r < n_raw; r += blockDim.x) { + sh_scores[r] = expf(sh_scores[r] - max_s); + den_local += sh_scores[r]; + } + partial[threadIdx.x] = den_local; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + 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)tok * n_head + h) * head_dim; + if ((head_dim & 3u) == 0u) { + for (uint32_t d = threadIdx.x * 4u; d < head_dim; + d += blockDim.x * 4u) { + float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + const float4 kv = ((const float4 *)(raw_kv + + (uint64_t)row * head_dim + d))[0]; + const float weight = sh_scores[r]; + acc.x += kv.x * weight; + acc.y += kv.y * weight; + acc.z += kv.z * weight; + acc.w += kv.w * weight; + } + ((float4 *)oh)[d >> 2u] = make_float4( + acc.x / denom, acc.y / denom, + acc.z / denom, acc.w / denom); + } + } else { + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + float acc = 0.0f; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + acc += raw_kv[(uint64_t)row * head_dim + d] * sh_scores[r]; + } + oh[d] = acc / denom; + } + } +} + +extern "C" int ds4_gpu_attention_noncausal_raw_batch_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 n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (!heads || !q || !raw_kv || !model_map || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || n_head == 0 || head_dim == 0 || + sinks_offset > model_size || + (uint64_t)n_head * sizeof(float) > model_size - sinks_offset || + heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) || + raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) { + return 0; + } + const float *sinks = (const float *)cuda_model_range_ptr( + model_map, sinks_offset, (uint64_t)n_head * sizeof(float), "dspark_attn_sinks"); + if (!sinks) return 0; + const size_t shmem = (size_t)n_raw * sizeof(float); + if (shmem > 32768) return 0; /* draft blocks are tiny; guard anyway */ + dim3 grid(n_tokens, n_head, 1); + attention_noncausal_raw_batch_heads_kernel<<>>( + (float *)heads->ptr, + sinks, + (const float *)q->ptr, + (const float *)raw_kv->ptr, + n_tokens, n_raw, raw_cap, raw_start, n_head, head_dim); + if (!cuda_ok(cudaGetLastError(), "attention noncausal raw batch heads launch")) return 0; + static int verify_left = -1; + if (verify_left < 0) { + verify_left = getenv("DS4_DSPARK_VERIFY_NONCAUSAL") != NULL ? 3 : 0; + } + if (verify_left > 0) { + verify_left--; + (void)cudaDeviceSynchronize(); + const uint64_t qn = (uint64_t)n_tokens * n_head * head_dim; + const uint64_t kn = (uint64_t)raw_cap * head_dim; + std::vector hq(qn), hkv(kn), hout(qn), hsink(n_head); + (void)cudaMemcpy(hq.data(), q->ptr, qn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hkv.data(), raw_kv->ptr, kn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hout.data(), heads->ptr, qn * 4, cudaMemcpyDeviceToHost); + (void)cudaMemcpy(hsink.data(), sinks, (uint64_t)n_head * 4, cudaMemcpyDeviceToHost); + double max_abs = 0.0, max_rel = 0.0; + const double scale = 1.0 / sqrt((double)head_dim); + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t h = 0; h < n_head; h++) { + std::vector sc(n_raw); + double mx = (double)hsink[h]; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + double dot = 0.0; + for (uint32_t d = 0; d < head_dim; d++) { + dot += (double)hq[((uint64_t)t * n_head + h) * head_dim + d] * + (double)hkv[(uint64_t)row * head_dim + d]; + } + sc[r] = dot * scale; + if (sc[r] > mx) mx = sc[r]; + } + double den = exp((double)hsink[h] - mx); + for (uint32_t r = 0; r < n_raw; r++) den += exp(sc[r] - mx); + for (uint32_t d = 0; d < head_dim; d++) { + double acc = 0.0; + for (uint32_t r = 0; r < n_raw; r++) { + const uint32_t row = (raw_start + r) % raw_cap; + acc += exp(sc[r] - mx) * (double)hkv[(uint64_t)row * head_dim + d]; + } + const double ref = acc / den; + const double got = (double)hout[((uint64_t)t * n_head + h) * head_dim + d]; + const double ad = fabs(ref - got); + if (ad > max_abs) max_abs = ad; + if (fabs(ref) > 1e-3 && ad / fabs(ref) > max_rel) max_rel = ad / fabs(ref); + } + } + } + fprintf(stderr, + "ds4: DSpark noncausal verify n_tok=%u n_raw=%u start=%u cap=%u " + "max_abs=%.3e max_rel=%.3e\n", + n_tokens, n_raw, raw_start, raw_cap, max_abs, max_rel); + } + return 1; +} + extern "C" int ds4_gpu_attention_decode_mixed_batch_heads_tensor( ds4_gpu_tensor *heads, const void *model_map, @@ -996,7 +1177,10 @@ extern "C" int ds4_gpu_attention_output_q8_batch_tensor( if (!attn_output_cublas) { if ((group_dim & 31u) == 0u && rank <= UINT32_MAX && n_tokens <= UINT32_MAX) { const uint32_t rows_per_block = 32u; - const uint32_t tile = 32u; + /* Small batches (DSpark verify n<=16) waste up to 6x ALU on the + * fixed 32-wide token tile; pick the narrowest tile that fits. + * Per-token accumulation order is unchanged (bit-identical). */ + const uint32_t tile = n_tokens <= 8u ? 8u : (n_tokens <= 16u ? 16u : 32u); const uint32_t block_tile = 16u; cuda_launch_grouped_q8_a_sharedx((float *)low->ptr, out_a, diff --git a/rocm/ds4_rocm_indexer.cuh b/rocm/ds4_rocm_indexer.cuh index a5e4eb93dc..9a96509fe8 100644 --- a/rocm/ds4_rocm_indexer.cuh +++ b/rocm/ds4_rocm_indexer.cuh @@ -287,6 +287,43 @@ __global__ static void argmax_kernel(int32_t *out_idx, const float *logits, uint if (tid == 0u) *out_idx = sm_idx[0]; } +/* Batched argmax: one block per row, same first-max-wins tie-break as + * indexer_topk_kernel with top_k=1 (used by DSpark verify row tops). */ +__global__ static void argmax_rows_kernel(uint32_t *selected, const float *scores, + uint32_t n_comp, uint32_t n_tokens) { + enum { THREADS = 1024 }; + __shared__ float sm_val[THREADS]; + __shared__ uint32_t sm_idx[THREADS]; + const uint32_t t = blockIdx.x; + if (t >= n_tokens) return; + const float *row = scores + (uint64_t)t * n_comp; + const uint32_t tid = threadIdx.x; + float local_v = -INFINITY; + uint32_t local_i = 0; + for (uint32_t i = tid; i < n_comp; i += THREADS) { + const float v = row[i]; + if (v > local_v) { + local_v = v; + local_i = i; + } + } + sm_val[tid] = local_v; + sm_idx[tid] = local_i; + __syncthreads(); + for (uint32_t s = THREADS / 2u; s > 0u; s >>= 1u) { + if (tid < s) { + const float vr = sm_val[tid + s]; + const uint32_t ir = sm_idx[tid + s]; + if (vr > sm_val[tid] || (vr == sm_val[tid] && ir < sm_idx[tid])) { + sm_val[tid] = vr; + sm_idx[tid] = ir; + } + } + __syncthreads(); + } + if (tid == 0u) selected[t] = sm_idx[0]; +} + __global__ static void indexer_topk_kernel(uint32_t *selected, const float *scores, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) { uint32_t t = blockIdx.x; if (t >= n_tokens || threadIdx.x != 0) return; @@ -865,6 +902,14 @@ extern "C" int ds4_gpu_indexer_topk_tensor( selected->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) { return 0; } + if (top_k == 1u) { + /* DSpark verify row tops: a batched parallel argmax, not the + * single-threaded insertion-sort fallback (was ~13 ms per call). */ + argmax_rows_kernel<<>>((uint32_t *)selected->ptr, + (const float *)scores->ptr, + n_comp, n_tokens); + return cuda_ok(cudaGetLastError(), "indexer topk argmax-rows launch"); + } if (top_k == 512u && n_comp <= 1024u) { indexer_topk_1024_kernel<<>>((uint32_t *)selected->ptr, (const float *)scores->ptr, @@ -1130,3 +1175,138 @@ extern "C" int ds4_gpu_dsv4_indexer_qat_tensor(ds4_gpu_tensor *x, uint32_t n_row indexer_hadamard_fp4_kernel<<>>((float *)x->ptr, n_rows, head_dim); return cuda_ok(cudaGetLastError(), "indexer_hadamard_fp4 launch"); } + +/* DSpark markov chain step (ROCm port of the CUDA kernel): + * out = argmax_i(logits[i] + dot(w2[i], w1[prev])) over the vocab, entirely + * on-device. w1/w2 are q8_0 with 34-byte blocks (half scale + 32 int8). */ +__global__ static void dspark_markov_argmax_kernel( + unsigned long long *out_key, + const float *logits, + const unsigned char *w1_row, + const unsigned char *w2, + uint32_t vocab, + uint32_t rank_blocks, + uint32_t reverse) { + __shared__ float state[256]; + const uint32_t tid = threadIdx.x; + for (uint32_t t = tid * 2u; t < rank_blocks * 32u; t += blockDim.x * 2u) { + const uint32_t b = t >> 5, k = t & 31u; + const unsigned char *blk = w1_row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + const uint16_t packed = *(const uint16_t *)(blk + 2 + k); + state[t] = d * (float)(int8_t)(packed & 0xffu); + state[t + 1u] = d * (float)(int8_t)(packed >> 8u); + } + __syncthreads(); + + float best_v = -INFINITY; + uint32_t best_i = 0; + for (uint32_t i0 = blockIdx.x * blockDim.x + tid; i0 < vocab; + i0 += gridDim.x * blockDim.x) { + const uint32_t i = reverse ? (vocab - 1u - i0) : i0; + const unsigned char *row = w2 + (uint64_t)i * rank_blocks * 34u; + float acc = 0.0f; + for (uint32_t b = 0; b < rank_blocks; b++) { + const unsigned char *blk = row + (uint64_t)b * 34u; + const float d = __half2float(*(const __half *)blk); + const int8_t *q = (const int8_t *)(blk + 2); + float s = 0.0f; + #pragma unroll + for (uint32_t k = 0; k < 32u; k += 2u) { + const uint16_t packed = *(const uint16_t *)(q + k); + s += (float)(int8_t)(packed & 0xffu) * + state[b * 32u + k]; + s += (float)(int8_t)(packed >> 8u) * + state[b * 32u + k + 1u]; + } + acc += d * s; + } + const float v = logits[i] + acc; + if (topk_score_better(v, i, best_v, best_i)) { + best_v = v; + best_i = i; + } + } + + __shared__ float vals[768]; + __shared__ uint32_t idxs[768]; + vals[tid] = best_v; + idxs[tid] = best_i; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1; stride > 0u; stride >>= 1u) { + if (tid < stride) { + if (topk_score_better(vals[tid + stride], idxs[tid + stride], + vals[tid], idxs[tid])) { + vals[tid] = vals[tid + stride]; + idxs[tid] = idxs[tid + stride]; + } + } + __syncthreads(); + } + if (tid == 0u) { + /* Monotonic float key; ~idx in the low bits makes ties resolve to + * the smaller index under atomicMax (matches topk_score_better). */ + const unsigned int f = __float_as_uint(vals[0]); + const unsigned int fkey = (f & 0x80000000u) ? ~f : (f | 0x80000000u); + const unsigned long long key = + ((unsigned long long)fkey << 32) | (unsigned int)(~idxs[0]); + atomicMax(out_key, key); + } +} + +static int g_markov_reverse = 0; +/* 1 row per thread: 4096 warps x 32 threads sweeps the vocab in one pass + * (measured 110 us vs 183 us at 128x256 on gfx1151; more blocks = better + * latency hiding, small blocks pack 12/CU via LDS). */ +static int g_markov_grid = 4096; +static int g_markov_block = 32; + +extern "C" void ds4_gpu_dspark_markov_set_reverse(int rev) { + g_markov_reverse = rev; +} + +extern "C" void ds4_gpu_dspark_markov_set_shape(int grid, int block) { + g_markov_grid = grid; + g_markov_block = block; +} + +int ds4_gpu_dspark_markov_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits_row, + const void *model_map, + uint64_t model_size, + uint64_t w1_offset, + uint64_t w2_offset, + uint32_t prev_token, + uint32_t vocab, + uint32_t rank) { + if (!out_idx || !logits_row || !model_map || vocab == 0 || + rank == 0 || (rank & 31u) != 0u || rank > 256u || + out_idx->bytes < sizeof(unsigned long long) || + logits_row->bytes < (uint64_t)vocab * sizeof(float)) { + return 0; + } + const uint32_t rank_blocks = rank / 32u; + const uint64_t row_bytes = (uint64_t)rank_blocks * 34u; + if (w1_offset > model_size || + (uint64_t)prev_token * row_bytes + row_bytes > model_size - w1_offset || + w2_offset > model_size || + (uint64_t)vocab * row_bytes > model_size - w2_offset) { + return 0; + } + const unsigned char *w1_row = (const unsigned char *)cuda_model_range_ptr( + model_map, w1_offset + (uint64_t)prev_token * row_bytes, + row_bytes, "markov_w1_row"); + const unsigned char *w2 = (const unsigned char *)cuda_model_range_ptr( + model_map, w2_offset, (uint64_t)vocab * row_bytes, "markov_w2"); + if (!w1_row || !w2) return 0; + if (!cuda_ok(cudaMemsetAsync(out_idx->ptr, 0, sizeof(unsigned long long)), + "dspark markov out clear")) { + return 0; + } + dspark_markov_argmax_kernel<<>>( + (unsigned long long *)out_idx->ptr, + (const float *)logits_row->ptr, + w1_row, w2, vocab, rank_blocks, g_markov_reverse); + return cuda_ok(cudaGetLastError(), "dspark markov argmax launch"); +} diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 8c05433b99..fc3f204e90 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -408,7 +408,10 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode #endif if ((in_dim & 31u) == 0u && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) { const uint32_t rows_per_block = 32u; - const uint32_t tile = 32u; + /* Small batches (DSpark verify n<=16) waste up to 6x ALU on the + * fixed 32-wide token tile; pick the narrowest tile that fits. + * Per-token accumulation order is unchanged (bit-identical). */ + const uint32_t tile = n_tok <= 8u ? 8u : (n_tok <= 16u ? 16u : 32u); const uint32_t block_tile = 16u; cuda_launch_q8_batch_sharedx((float *)out->ptr, reinterpret_cast(wptr), diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 51fafe3d38..52654dc5ce 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -228,7 +228,7 @@ static int routed_moe_q2_float_down_launch( } const uint32_t down_tile = 4u; - const uint32_t down_rpb = 16u; + const uint32_t down_rpb = 32u; const uint32_t down_threads = down_rpb * 32u; const size_t down_shmem = (size_t)down_tile * 256u * sizeof(float); const int use_f16_down = (out_dim & 1u) == 0u; @@ -750,7 +750,7 @@ static int routed_moe_launch( (!q4k_path || n_tokens >= 32u) && !disable_resident_iq2_sorted; const uint32_t use_expert_tiles = use_sorted_pairs; - const uint32_t expert_tile_m = 8u; + const uint32_t expert_tile_m = 4u; const uint32_t write_gate_up = 0u; const uint32_t use_p2_sorted = 0u; const uint32_t use_atomic_down = use_expert_tiles && n_tokens >= 128u;