Rocm prefill performance improvement on 16.3% - #616
Draft
Farenheith wants to merge 6 commits into
Draft
Conversation
Replace scalar inner dot-product loops in attention_prefill_raw_kernel and attention_prefill_mixed_kernel with 128-bit float4 vectorized loads. For head_dim=512 this cuts memory transactions from 512 scalar (4B) to 128 float4 (16B) loads per dot product — ~4x fewer requests on the unified memory bus. On Strix Halo's shared CPU/GPU memory controller this reduces bus contention and improves cache-line utilization. Measured per-token speedup on 19.5k-token prefill (DeepSeek V4 Flash, ROCm, gfx1151): 149.46 t/s -> 163.49 t/s (+9.4%).
…educe All three prefill softmax kernels (attention_prefill_raw_softmax, attention_prefill_mixed_softmax, and attention_prefill_mixed_softmax_tile) used an O(log2(N)) stride-based reduction with ~8 __syncthreads() per pass over 256 threads. Replace with block_reduce_f32_max / block_reduce_f32_sum from ds4_rocm_common.cuh which do the same reduction in a single barrier by using intra-warp shfl_down_sync + one shared-memory round. Measured gain on 19.5k-token prefill (DeepSeek V4 Flash, ROCm, gfx1151): rocm-perf v1 (float4 prefill, barrier softmax): 163.49 t/s rocm-perf v2 (float4 prefill + block_reduce): 166.49 t/s (+1.8%) Cumulative gain from baseline (working branch, 149.46 t/s): +11.4%.
Add FP32 strided-batched plan infrastructure to ds4_rocm_hipblaslt.cuh: plan cache by (m,n,k,batch_count,strides,trans), auto-heuristic algo selection, transparent fallback to cublas on failure. Replace the two cublasSgemmStridedBatched calls in both the raw and mixed prefill cublas paths with hipBLASLt-first dispatch. Measured on 19.5k-token prefill (DeepSeek V4 Flash, gfx1151): v2 (float4 + block_reduce): 166.49 t/s v3 (+ hipBLASLt FP32 sb): 166.82 t/s (+0.2%) Modest gain — hipBLASLt heuristic found similar kernels to default rocBLAS for these FP32 shapes with zero workspace.
New attention_prefill_tiled_raw_online_kernel processes KV rows in 384-row tiles with two-pass softmax per tile (scores → max → softmax → weighted sum), keeping all intermediate data in shared memory. Replaces the 5-stage cublas pipeline (pack + SGEMM + softmax + SGEMM + unpack) with a single kernel for n_tokens>1, head_dim=512, !quality. Measured 19.5k-token prefill: v3 (hipBLASLt): 166.82 t/s v4 (tiled fused): 166.41 t/s (~same, compute-bound) No regression — the tiled kernel is a foundation for FP16 tensor-core and mixed-KV variants.
New attention_prefill_sp_online_kernel combines score computation, numerically stable online softmax (double-precision correction factors), and weighted value accumulation in a single pass over KV rows. Loads KV into shared memory once per 4-row group — no separate score and value passes, no shared-memory score buffer. Measured on 19.5k-token prefill: ~165 t/s (no regression vs prior two-pass tiled kernel). Coherence verified.
New hipblaslt_gemm_tn_f16_out_f32: FP16 weights × FP16 activations → FP32 output via hipBLASLt autotuned kernels. Wired into both matmul_f16_tensor and matmul_q8_0_tensor as a try-first alternative to cublasGemmEx for n_tok > 1 (prefill batch). Measured on 19.5k-token prefill (DeepSeek V4 Flash, gfx1151): v4 (tiled fused, attention only): 166.82 t/s v5 (+ hipBLASLt FP16→FP32 matmul): 173.85 t/s (+4.2%) Coherence verified. The gain confirms MoE FFN GEMMs are the dominant prefill time cost on this model, and hipBLASLt finds better tensor-core algorithms than default cublas for the prefill batch shapes.
Farenheith
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a series of prefill optimizations DeepSeek did, focusing on my machine:
My tests showed a 16.3% improvement, from 149.46 t/s to 173.85 t/s, and I haven't identified any bugs in inference so far.
These changes were made over commit 80ebbc3 since a regression happened after this commit (I've found the same problem I had been reported here).
I'm opening this PR to share the results with the community, although I haven't checked yet whether these changes will only work for hardware similar to mine.
It'd be nice if someone more experienced than me could take a look.