perf(ds4): accelerate prefill via gpu-resident hc pipeline and intra-warp flash attention

Closed
CH
2mo ago
+232-109

Summary

Stacked on PR #633 (codex/perf-ds4-prefill-phase1).

This change accelerates chunked prefill on DeepSeek-V4 from 130.84 tok/s to 225.83 tok/s sustained (225.85 tok/s peak) on dual-GPU AMD ROCm target hardware (gfx1151;gfx1201), reducing 3,960-token prompt latency from 30.26s down to 17.53s ($1.73\times$ speedup over PR #633, $15.17\times$ over pre-633 dual baseline).

Key Architectural Changes

  1. GPU-Resident HC Pipeline by Default (deepseek4_graph.cpp):

    • Switched batched GPU-resident Hierarchical Controller (HC) pre/post graph compute on by default during heterogeneous prefill (use_backend_prefill_hc).
    • Eliminates 9.58 seconds of host CPU-bound Sinkhorn normalization loops across all 43 layers.
    • Retains DFLASH_DS4_DISABLE_HYBRID_PREFILL_GPU_HC as an explicit safety opt-out escape hatch.
  2. Intra-Warp Shuffle Reductions (fattn.cu):

    • Replaced multi-barrier sequential LDS tree reductions (blockDim.x / 2 loops with 16 __syncthreads()) with single-cycle __shfl_xor (HIP) / __shfl_xor_sync (CUDA) intra-warp reductions.
    • Reduced per-block LDS scratch from 1024 floats down to 32 floats, cutting 32-way LDS bank conflicts and reducing barrier synchronization count from 16 to 2.
  3. Branchless $D=512$ Attention Loop Splitting (fattn.cu):

    • Statically partitioned the 512-dimension dot-product into 448 unrotated dimensions (direct dot-product) + 64 tail dimensions (RoPE rotation with pre-computed coefficients).
    • Removes inner-loop condition branches and enables full compiler loop unrolling and memory load pipelining.
  4. 128-Bit Vectorized Global Memory Stores (fattn.cu):

    • Aligned and vectorized destination context writes into 128-bit float4 and 64-bit float2 global store instructions (global_store_dwordx4 on AMD GCN/RDNA).
  5. API Documentation Alignment (deepseek4_internal.h):

    • Clarified that deepseek4_step populates out_logits only when need_logits=true, leaving out_logits empty when need_logits=false.

Hardware Scope (Strix-Only vs Dual-GPU)

  • Strix-Only Users (gfx1151 APU): Fully enabled. Wave32 shuffle reductions compile directly to hardware permute instructions (v_permlane16_b32 / v_permlanex16_b32). Single-GPU execution uses ds4_try_layer_major_prefill which executes monolithic GPU graphs on Strix Halo with zero host roundtrips.
  • Dual-GPU Users (gfx1201 + gfx1151): Fully enabled. GPU-resident HC graphs eliminate 86 cross-device synchronization sync points per token chunk.

Measured Prefill Performance

Model: 96 GB DeepSeek V4 (DeepSeek-V4-Flash-ROCMFP2-STRIX.gguf, 43 layers)
Prompt: 3,960 tokens (long_code_audit.txt, chunk=512)
Hardware: Dual AMD GPU (Radeon AI PRO R9700 gfx1201 + Ryzen AI MAX+ 395 gfx1151, ROCm 7.2.4)

MilestoneLatency (3,960 tok)ThroughputSpeedup vs Base
Dual Baseline (PR #590)194.20s14.88 tok/s$1.00\times$
PR #633 Head30.26s130.84 tok/s$8.79\times$
PR #634 Head (Stacked)17.53s225.83 tok/s (225.85 peak)$15.17\times$

Live Multi-Condition Serial Test Results

Executed serially on live hardware with clean server lifecycle per configuration:

Test ConfigurationPrompt TokensMeasured LatencyMeasured ThroughputNotes
Long Prompt (chunk=512)3,96017.53s225.83 tok/sSustained prefill rate ($1.73\times$ speedup)
Cold Start (chunk=512)3,32817.21s193.37 tok/sIncludes initial memory pool allocation & graph build
Medium Prompt (chunk=512)1,7159.96s172.18 tok/sMonotonic scaling across chunk boundaries
Short Chunk Stress (chunk=256)3,32835.62s93.43 tok/sConfirms chunk=512 is optimal for PCIe amortization
Negative Control (GPU_HC=0)3,32837.05s89.83 tok/s$2.15\times$ slowdown, confirming real GPU compute gain

Determinism & Generation Quality

  • Bit-Exact Determinism: 3 independent greedy runs ($T=0.0$) produced bit-identical output with matching SHA-256 (8d1ca74caeecec5f2b52300735692c95dd6d7e83e6e0734d3cb8daacfc85e5a9).
  • Generation Quality: Generates syntactically correct, non-degraded code with zero repetition loops or NaNs.
  • Unit Validation: 53/53 passed (OK) on dual ROCm 7.2.4 (gfx1151;gfx1201) and local SM 8.9 CUDA.
  • git diff --check: Clean (zero formatting or whitespace errors).