[ROCm][Perf] TP-shard Lightning indexer prefill - #46103
Concurrensee wants to merge 7 commits into
Conversation
|
This pull request has merge conflicts that must be resolved before it can be |
6c5d9f8 to
ae386c7
Compare
f665980 to
0551372
Compare
0551372 to
b97bbd9
Compare
0a0718c to
3035d36
Compare
Signed-off-by: Yida Wu <yida.wu@amd.com> Signed-off-by: Claude Opus 4.8
Signed-off-by: Yida Wu <yida.wu@amd.com>
3035d36 to
ae68614
Compare
Signed-off-by: Yida Wu <yida.wu@amd.com>
d482721 to
f514565
Compare
|
Thanks @Concurrensee ! What's the command to run your test here? I saw you have added two envs, what's the best practice to use? |
|
Hi @dllehr-amd can you please spare some time to look at this PR? |
|
Hi @Fangzhou-Ai, Default config: keep the defaults — TP_SHARD on, FOLD off. TP_SHARD self-gates (only active on gfx950 when tp_size > 1 and num_prefill_tokens >= 512 * tp_size), so short prefills fall back to the replicated path and long-context prefills get the speedup automatically. When tuning chunk size for long-context runs, turn on FOLD (VLLM_ROCM_USE_AITER_MLA_SPARSE_TP_SHARD_FOLD=1). Larger prefill chunks span a wider causal range per step, where the work-balanced fold partition pays off the most. Since VLLM_ROCM_USE_AITER_MLA_SPARSE_TP_SHARD is on by default, if you observe a regression on a specific case, turn it off (=0) to fall back to the original replicated path. |
|
Hi @Fangzhou-Ai, |
Check whether pynccl_comm disabled or not Signed-off-by: Yida Wu <yida.wu@amd.com>
4d943be to
b313953
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
4 participants
Purpose
On ROCm/gfx950 the DeepSeek-V3.2-style Lightning indexer prefill (the FP8 MQA
logits+ per-rowtop_k) is computed redundantly on every tensor-parallel rank: each rank runs the full indexer over all prefill query rows and produces the identicaltopk_indices_buffer. Unlike the main attention/MLP, this work is not split by TP, so it is replicatedtp_sizetimes. Its cost grows with the causal range each query attends to, so at long context the replicated indexer prefill becomes a dominant, TP-non-scaling chunk of TTFT.This PR shards that prefill indexer work across the TP group on the query (M / token) dimension. Each rank computes
rocm_fp8_mqa_logits+top_k_per_row_prefillfor only its slice of prefill rows, then the ranks exchange their slices so every rank ends up with the fulltopk_indices_bufferthat the subsequent sparse attention needs. The KV-cache quant/gather and the decode path are unchanged.Two partition schemes are provided:
TP_SHARD(default ON): split thenum_prefill_tokensrows intotp_sizecontiguous shards. The remainder rows are assigned to rank 0 (the cheapest, lowest-index shard). Reassembly is hybrid NCCL: a ringall_gatherover the evenly-divisible block plus a singlebroadcastfor the remainder rows.TP_SHARD_FOLD(default OFF): a work-balanced fold for the case where a step spans a very large causal range. Because per-row indexer cost rises with token index (~ ke - ks), a plain contiguous split overloads the last rank. Fold cuts the rows into2*tpgroups and pairs a cheap low-index group with an expensive high-index one on each rank. The split is fixed symmetric (M ≈ P/2, rounded to a multiple oftpso the first half can ringall_gather), which already balances the ~linear per-row cost. Reassembly is the rank-ordered first half via ringall_gatherand the reversed second half via per-groupbroadcasts (batched in a singlegroup_start/group_end).Sharding only kicks in on gfx950 with
tp_size > 1andnum_prefill_tokens >= 512 * tp_size, so short prefills (where the indexer is negligible and the collective would just add latency) fall back to the original replicated path.Minor cleanups in the same change:
cp_gather_indexer_k_quant_cachewhenchunk.skip_kv_gatheris set (avoids re-gathering K for chunks that don't need it).topk_indices_buffer[...] = -1memset on gfx950 (the kernel fully writes the buffer there); keep it as a safety net on other archs.Test Plan
vllm bench servesweeps over ISL × concurrency for both models (baseline vsTP_SHARDvsTP_SHARD+FOLD), reported below.Test Result
Accuracy — GSM8K, 5-shot, full set
TP_SHARDTP_SHARD+FOLDTP_SHARDTP_SHARD+FOLDPerformance
Setup: 8×MI355 (gfx950), TP=8,
rocm/vllm-dev:nightly(vllm0.23.1rc1.dev354) with this PR's two files overlaid;vllm bench serve, random dataset, OSL=1000,--ignore-eos, request-rateinf. All numbers are vs thebaseline(replicated) column and+always means better (higher throughput / lower latency).The first three columns (
baseline/TP_SHARD/TP_SHARD+FOLD) run at vLLM default gpu-mem-util and default chunked-prefill size (8192). The twochunk16kcolumns combine the shard with a larger prefill chunk (--max-num-batched-tokens 16384) at--gpu-memory-utilization 0.88; this is a serving-side tuning knob (not part of this PR's code) that trades a little short/mid-context TTFT for better long-context TPOT/throughput.Short context (≤ 8K) is neutral within ±2–3% (the indexer is a tiny fraction of work, and the gate skips sharding for the smallest prefills); the win grows with context length, reaching 2.1× throughput for GLM-5.1 at 199K and +40% for DeepSeek-V4-Pro at 512K.
GLM-5.1-FP8
Output throughput (tok/s)
Mean TTFT (s)
Mean TPOT (ms)
DeepSeek-V4-Pro
Output throughput (tok/s)
Mean TTFT (s)
Mean TPOT (ms)
Notes
Both env vars are read on gfx950 only; sharding additionally requires
tp_size > 1andnum_prefill_tokens >= 512 * tp_size.VLLM_ROCM_USE_AITER_MLA_SPARSE_TP_SHARD(default1): shard the indexer prefill across the TP group on the token dimension. Set0to fall back to the original replicated path.VLLM_ROCM_USE_AITER_MLA_SPARSE_TP_SHARD_FOLD(default0): use the work-balanced fold partition instead of contiguous shards (the fold split point is fixed atP/2).