[DCP][Kernel][Perf] Fuse the empty-shard LSE mask into the A2A pack kernel - #54889
Conversation
Under DCP, every MLA layer calls mask_dcp_empty_shards_ just before packing the partial attention output and LSE for the all-to-all combine. It sets the LSE of rows whose request holds no local KV shard -- and of CUDA-graph padding rows -- to -inf so they carry no weight in the reduction. Written in eager PyTorch it expands to eight kernels (arange, searchsorted, clamp, two compares, an index, an or, and the masked_fill), six of which touch only ~100 elements. All eight are captured inside the full CUDA graph, so each pays the ~4.3us dispatch floor on every replay: ~35us per MLA layer. The mask is a pure function of seq_lens and query_start_loc, both identical across every MLA layer in a step, and the pack kernel already loads the very LSE value that needs masking. Fold the mask in as a Triton-side search over query_start_loc plus a tl.where, deleting all eight kernels for no extra memory traffic. This brings the Triton path to parity with the direct symmetric-memory path, which already does exactly this inside its kernel -- see find_sequence() and `empty_kv ? -CUDART_INF_F : ...` in csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu. num_seqs is a runtime argument rather than a constexpr so the kernel specializes only on the arange bound: 7 variants across batch sizes 2..104 instead of one per distinct batch size. Measured on Kimi-K3, MI355X, TP8/DCP8, --dcp-comm-backend a2a, fp8 KV, 8k/1k conc-52 decode, full CUDA graphs: the eight kernels drop from ~1.0 to 0.0 invocations per MLA layer, removing 35.0us/MLA layer against a 1.2us increase in the pack kernel -- a net 33.8us per MLA layer, 0.81ms per step over 24 MLA layers (2.3% of a 35.7ms step). mask_dcp_empty_shards_ is retained for the AG+RS combine path, which is not exercised by --dcp-comm-backend a2a. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
The existing cases all build query_start_loc with a uniform stride, so the gap between consecutive entries is constant and the in-kernel boundary search is never asked to handle unequal spans. MTP with partial acceptance produces exactly that: some requests contribute one row and others several, in the same batch. Add eight seeded cases per dtype with random per-request query lengths in [1, 4] over 17 requests, a random mix of empty and non-empty shards, and three trailing padding rows. Each seed exercises all four query lengths and masks 50-65% of rows, so the comparison is not vacuous. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@dllehr-amd @AndreasKaratzas can you please review and add the ready/ci-run plz? |
dllehr-amd
left a comment
dllehr-amd
left a comment
There was a problem hiding this comment.
Looks okay to me. I'll wait for additional guidance from Matt or Lukas
|
✅ @rbrugaro-amd, CI is now available for this PR.
|
AndreasKaratzas
left a comment
AndreasKaratzas
left a comment
There was a problem hiding this comment.
not sure of the numerics here, mayube i m wrong, just if you could please provide some info for that. also for the decorator on top of test file.
| ) | ||
|
|
||
| pytestmark = pytest.mark.skipif( | ||
| not torch.cuda.is_available(), reason="needs a GPU for the Triton pack kernel" |
There was a problem hiding this comment.
Wait does this evaluate to true or false for us? I thought we should be using cuda_alike
There was a problem hiding this comment.
It does evaluate correctly to True, but it is the wrong predicate to express the intent, and I have switched it.
torch.cuda.is_available() returns True on ROCm, because the PyTorch ROCm build maps HIP onto the torch.cuda API. So the file was running on our hardware rather than being skipped. Measured in the ROCm image:
torch.cuda.is_available() = True
torch.version.hip = 7.2.53211
current_platform.is_cuda() = False
current_platform.is_cuda_alike() = True
The trap is the neighbouring predicate: current_platform.is_cuda() is False on ROCm and would have silently skipped the whole file.
Switched to [current_platform.is_cuda_alike()] which is what the rest of the tree uses for "CUDA or ROCm"
No behavioural change; 353 tests still pass on MI355X.
| num_heads = world_size * h_per_rank | ||
|
|
||
| out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype) | ||
| lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32) |
There was a problem hiding this comment.
we do we need 32 bit numerics here?
There was a problem hiding this comment.
The LSE is fp32 in production — the test is mirroring the real dtype, not choosing one. All of the following predate this PR:
_dcp_a2a_lse_pack_dimexists solely to pack an fp32 LSE into a narrower output dtype, and raises"Cannot pack fp32 LSE into output dtype {...}"out_lseis allocatedtorch.float32- the direct symmetric-memory path allocates
received_lseastorch.float32
It has to be fp32 for range — LSE is a log of summed exponentials, and this mask writes -inf into it. A 16-bit LSE would lose that.
|
|
||
| # Compare bit patterns, not values: for a 2-element LSE pack the high half | ||
| # of -inf bitcasts to a NaN payload, which never compares equal to itself. | ||
| assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), ( |
There was a problem hiding this comment.
so is this correct? (i.e. using 32 bit numerics in the beginning and then for the comparison only converting them to 16 bit ones)
There was a problem hiding this comment.
I think my comment was misleading, and I have rewritten it in db07f2ab.
Two different tensors are in play. The LSE is fp32 and stays fp32. The send buffer is in the output dtype (16-bit), and _dcp_a2a_lse_pack_dim returns 2 for any 16-bit output, so the fp32 LSE is bit-split across two of its slots. .view(torch.int16) bitcasts the send buffer — reinterpreting 16-bit storage as 16-bit integers. Nothing is narrowed.
Those slots are not meaningful floats at all; they are an opaque container for half of an fp32. A masked row stores -inf = 0xFF800000, and the high half 0xFF80 reads differently depending on the format:
bf16 (1|8|7): 1 | 11111111 | 0000000 exp all-ones, mantissa 0 -> -inf
fp16 (1|5|10): 1 | 11111 | 1110000000 exp all-ones, mantissa 896 -> NaN
bf16 is the top 16 bits of fp32 (same 8-bit exponent), so the truncated half of -inf is still -inf. fp16 has only a 5-bit exponent, so three of fp32's exponent bits spill into fp16's mantissa — and an all-ones exponent with a non-zero mantissa is NaN, which never compares equal to itself.
Bit equality is also strictly stronger than allclose, which is what a bit-exactness claim should be asserting.
torch.cuda.is_available() is True on ROCm -- the PyTorch ROCm build maps HIP onto the torch.cuda API -- so the guard was behaving correctly. But it is the wrong predicate to express the intent, and the neighbouring one is a trap: current_platform.is_cuda() is False on ROCm and would silently skip the whole file. Use current_platform.is_cuda_alike(), which reads as intended and covers both backends. No behavioural change; 353 tests still pass on MI355X. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
…k test The comment claimed the high half of -inf 'bitcasts to a NaN payload' for any 2-element LSE pack. That holds for fp16 but not bf16, and it invited the reading that the fp32 LSE is being narrowed to 16 bits for the comparison. Neither is what happens. The LSE stays fp32; it is the send buffer -- already in the 16-bit output dtype -- that is bitcast, because its LSE slots hold halves of an fp32 rather than meaningful floats. bf16 shares fp32's 8-bit exponent, so the truncated half is still -inf; fp16's 5-bit exponent pushes three of those bits into its mantissa, and an all-ones exponent with a non-zero mantissa is NaN, which never compares equal to itself. Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
rbrugaro-amd
changed the title
|
/ci run |
|
✅ Triggered Buildkite CI #88094 for commit |
|
@dllehr-amd my comments have been addressed |
…ernel (vllm-project#54889) Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
4 participants
[DCP][Kernel] Fuse the empty-shard LSE mask into the A2A pack kernel
Purpose
Under DCP, every MLA layer calls
mask_dcp_empty_shards_(vllm/v1/attention/ops/dcp.py)immediately before packing the partial attention output and LSE for the all-to-all
combine. It sets the LSE of rows whose request holds no local KV shard — and of
CUDA-graph padding rows — to
-inf, so those rows carry no weight in the reduction.Written in eager PyTorch, that one line expands to eight kernels:
Six of the eight touch only ~100 elements. All eight are captured inside the full
CUDA graph, so each one pays the ~4.3 µs kernel-dispatch floor on every single
replay — roughly 35 µs per MLA layer, entirely independent of context length or
batch size.
The mask is a pure function of
seq_lensandquery_start_loc, both of which areidentical for every MLA layer within a step, and the pack kernel already loads the
very LSE value that needs masking. So the mask can be folded into
_dcp_a2a_pack_send_kernelas a Triton-side search overquery_start_locplus atl.where, deleting all eight kernels for no additional memory traffic.This is a parity fix, not a new optimization
vLLM's direct symmetric-memory DCP path already does exactly this inside its
kernel —
csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu:This PR brings the Triton fallback path to the same behaviour. The direct path is
gated on
current_platform.is_cuda()(cp_common.py) and its.cufiles are onlybuilt under
VLLM_GPU_LANG STREQUAL "CUDA", so the Triton path is what every ROCmdeployment runs, plus any CUDA deployment without symmetric memory. CUDA installs
that do take the direct path are unaffected — they never reach this code.
Design notes
kernel IDs to
hipGraphLaunch), so a data-dependent "skip when nothing is empty"branch is not a legal alternative — it would be resolved once at capture time and
baked in. The mask has to be made free, not conditional.
num_seqsis a runtime argument, not aconstexpr. Only thetl.arangeboundis
constexpr, so the kernel specializes on at mostlog2(max_num_seqs)variants.Measured across every CUDA-graph capture size 2–104: 7 compiled variants, versus
103 if
num_seqswere aconstexpr.row_indices >= query_start_loc[-1]padding term. That term is arguably redundant (padding slots are already zero-filled
in the DCP local
seq_lensbuffer), but it is retained so the fused form isbit-identical rather than merely equivalent-in-practice.
mask_dcp_empty_shards_is retained for the AG+RS combine path(
_cp_lse_common), which is not exercised by--dcp-comm-backend a2a. Fusing thatpath is left for a follow-up.
Test Plan
tests/v1/attention/test_dcp_a2a_pack_mask.pycompares the fused path against theeager
mask_dcp_empty_shards_it replaces. The comparison is on bit patterns(
torch.equaloverint16views), not values: for a 2-element LSE pack the highhalf of
-infbitcasts to a NaN payload, which never compares equal to itself.Both sides of every comparison run the same pack kernel: the reference is the new
kernel with
HAS_MASK=Falsefed by the old eagermask_dcp_empty_shards_, so theonly thing that differs is where the mask came from and everything else cancels.
Three tests, 353 cases:
test_fused_mask_matches_eagertest_fused_mask_matches_eager_ragged(new test)query_start_loc. All the cases above use a constant stride, so the in-kernel boundary search is never asked to handle unequal spans — which is exactly what MTP produces under partial acceptance. 8 seeds × 2 dtypes with random per-request query lengths in [1, 4] over 17 requests, a random mix of empty/non-empty shards, and 3 trailing padding rows. Every seed exercises all four query lengths and masks 50–65% of rows, so the comparison is not vacuoustest_mask_disabled_is_unmaskedseq_lens=None, query_start_loc=Noneleaves the LSE untouchedExisting DCP coverage is unchanged —
tests/distributed/test_dcp_a2a.pygives theidentical result with and without this patch (the 7 failures are pre-existing on the
unpatched base and are dtype assertions in the test's own reference):
Measured effect
Kimi-K3, MI355X ×8, TP8 / DCP8,
--dcp-comm-backend a2a,ROCM_AITER_MLA, fp8 KVcache,
FULL_AND_PIECEWISECUDA graphs, 8k/1k ISL/OSL at concurrency 52, pure decode.A and B are consecutive runs on the same node with identical configuration, differing
only in this patch. Values are averaged over the per-rank traces (A: 5 ranks, B: 6)
and normalised per MLA layer.
before:
after:
The eight kernels between
kn_mla_reduce_v1_psand_dcp_a2a_pack_send_kerneldisappear. Identifying them needed no name matching: selecting every kernel whose
invocations-per-MLA-layer drop by ~1.0 returns exactly eight distinct kernels,
matching the eight aten ops in the source function.
−35.00 µs/MLA layer removed, against +1.23 µs in the pack kernel
_dcp_a2a_pack_send_kernelthat absorbsthe work (6.81 → 8.04 µs) — a net −33.77 µs per MLA layer, or 0.81 ms/step
across 24 MLA layers, against a 35.7 ms decode step (2.3%). Because the saving is
dispatch-floor time it is roughly constant per step, so it is a larger fraction at
shorter contexts and smaller batches.
Nothing else moves. Per-invocation durations for unrelated kernels, A → B:
mla_a8w8_qh16_qseqlen1_gqaratio16_lse_ps)mscclKernel_Sum)kn_mla_reduce_v1_ps)End-to-end
Serving A/B at 128k/1k, TP8/DCP8, a2a, fp8 KV, decode attention on
ROCM_AITER_MLA,baseline arm run immediately before the fused arm in each pair. Metric is
median_itl_ms— the per-decode-step figure this change moves; throughput and TTFT are excluded because
at 128k both are dominated by prefill, which is untouched.
All six runs completed with full token counts and no worker deaths, and the C=52 baseline
reproduces an independent earlier measurement of the same configuration (44.70 ms).
Accuracy
GSM8K, 5-shot, same configuration, with the patch applied:
(Full 1319-sample test split. No speculative decoding in this configuration.)