Repository navigation
[HIP] [JIT] Add a DSv4 paged K gather+dequant kernel - #5680
Conversation
Add `dsv4_dequantize_and_gather_k`, a HIP kernel that gathers the DeepSeek V4 paged K record into a bf16 workspace, replacing a Triton gather that does not coalesce. One wavefront handles one gathered token so the store is a single coalesced line, and the fp8 decode is branchless bit arithmetic parameterized by the exponent bias, serving both e4m3fnuz and OCP e4m3. Matches the Triton reference bit for bit on all representable inputs, except that a product landing in the bf16 subnormal range flushes to zero: fp8 0x3f and 0xbf at scale 2^-126 under e4m3fnuz. NaN encodings are not represented, since the V4 encoder clamps to FP8_MAX. New op only: nothing existing changes and no caller is switched over. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
Re-benchmarked against vllm-project/vllm#56720vllm-project/vllm#56720 rewrites this Triton gather for ROCm (one pass, plus a max_gather_len-sized worker grid). It gates on is_rocm(), not on arch, so gfx942 will inherit it but regardless I re-measured against that rather than the old two-pass kernel. MI325X (gfx942), CUDA-graph replay, fragmented block table, 5 repeats, median. Both kernels in one process on identical inputs.
This kernel holds 3.2-3.5 TB/s at every shape. The one-pass grid caps at _GATHER_MAX_WORKERS, sized against a 256-CU gfx950; on gfx942 that cap binds past ~256k gathered tokens, each program loops, and Triton drops to a flat ~1.7-1.8 TB/s. So the gap is 1.1x on short prefills and 1.8-1.9x from 128k up! |
2 participants
Summary
Add a HIP kernel for the DeepSeek V4 paged K gather+dequant. The existing Triton kernel does not coalesce and holds ~0.34 TB/s on a 4.2 TB/s device, and it runs twice per chunk in the prefill loop.
Suggested fix
The new HIP kernel maps one wavefront to one gathered token, with lane L owning output dims [L8, L8+8). The 64 lanes cover all 512 dims so the store is a single 1024 B line. Lanes 0-55 read the fp8 payload and dequantize, lanes 56-63 pass the bf16 RoPE tail through. The fp8->f32 decode is branchless and parameterized by the exponent bias, so one path serves both e4m3fnuz and OCP e4m3.
Perf
MI325X (gfx942), CUDA-graph replay, k_cache[*, 64, 584]:
Triton sits at ~0.34 TB/s across all three; this reaches 3.3 TB/s at the larger shapes, against ~4.2 TB/s measured streaming rate on this arch.
Correctness