[DCP][Kernel][Perf] Fuse the empty-shard LSE mask into the A2A pack kernel (#54889)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Rita Brugarolas committed
Sep 10, 2026 at 08:43 UTC
6ee5bb0a0b3e32dd1a6d9fddb61b50905ccdd6e0
2 files changed
+213
-5
tests/v1/attention/test_dcp_a2a_pack_mask.py
new
+151
@@ -0,0 +1,151 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The empty-shard mask fused into the DCP A2A pack kernel must be bit-exact
4
+with the eager ``mask_dcp_empty_shards_`` pass it replaces."""
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.platforms import current_platform
10
+from vllm.v1.attention.ops.dcp import (
11
+ _dcp_a2a_lse_pack_dim,
12
+ _dcp_a2a_pack_send,
13
+ mask_dcp_empty_shards_,
14
+)
15
+
16
+pytestmark = pytest.mark.skipif(
17
+ not current_platform.is_cuda_alike(),
18
+ reason="needs a CUDA or ROCm device for the Triton pack kernel",
19
+)
20
+
21
+
22
+def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc):
23
+ lse_pack_dim = _dcp_a2a_lse_pack_dim(out.dtype)
24
+ send = torch.zeros(
25
+ (world_size, out.shape[0], h_per_rank, head_dim + lse_pack_dim),
26
+ device=out.device,
27
+ dtype=out.dtype,
28
+ )
29
+ _dcp_a2a_pack_send(
30
+ out,
31
+ lse,
32
+ send,
33
+ world_size,
34
+ h_per_rank,
35
+ head_dim,
36
+ lse_pack_dim,
37
+ seq_lens=seq_lens,
38
+ query_start_loc=query_start_loc,
39
+ )
40
+ return send
41
+
42
+
43
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
44
+@pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)])
45
+@pytest.mark.parametrize("tokens_per_req", [1, 2, 3])
46
+@pytest.mark.parametrize("num_pad_rows", [0, 2])
47
+@pytest.mark.parametrize(
48
+ "seq_lens_list",
49
+ [
50
+ [16, 16, 16, 16], # no empty shards (the long-context decode case)
51
+ [16, 0, 16, 0], # interior empty shards
52
+ [0, 0, 0, 0], # every shard empty
53
+ [16, 16, 0, 0], # trailing empty shards, i.e. cudagraph padding
54
+ [0, 16, 16, 16], # leading empty shard
55
+ [16] * 52, # the profiled conc-52 decode batch
56
+ [7], # single request
57
+ ],
58
+)
59
+def test_fused_mask_matches_eager(
60
+ dtype, world_size, h_per_rank, tokens_per_req, num_pad_rows, seq_lens_list
61
+):
62
+ torch.manual_seed(0)
63
+ device = "cuda"
64
+ head_dim = 512
65
+ num_seqs = len(seq_lens_list)
66
+ # tokens_per_req > 1 is the MTP / multi-token-verify shape.
67
+ num_rows = num_seqs * tokens_per_req + num_pad_rows
68
+ num_heads = world_size * h_per_rank
69
+
70
+ out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype)
71
+ lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32)
72
+ seq_lens = torch.tensor(seq_lens_list, device=device, dtype=torch.int32)
73
+ query_start_loc = torch.arange(
74
+ 0,
75
+ (num_seqs + 1) * tokens_per_req,
76
+ tokens_per_req,
77
+ device=device,
78
+ dtype=torch.int32,
79
+ ) # rows past query_start_loc[-1] are cudagraph padding
80
+
81
+ # Reference: mask eagerly, then pack with masking disabled.
82
+ ref_lse = lse.clone()
83
+ mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc)
84
+ expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None)
85
+
86
+ # Under test: pack with the mask fused in.
87
+ actual = _pack(
88
+ out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc
89
+ )
90
+
91
+ # The packed LSE slots hold halves of an fp32, not meaningful floats, so
92
+ # compare the send buffers bitwise: -inf's high half is a NaN in fp16.
93
+ assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), (
94
+ "fused mask is not bit-exact with the eager mask"
95
+ )
96
+
97
+
98
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
99
+@pytest.mark.parametrize("seed", range(8))
100
+def test_fused_mask_matches_eager_ragged(dtype, seed):
101
+ """Non-uniform per-request query lengths.
102
+
103
+ The uniform-stride cases above never exercise an unequal gap between
104
+ consecutive ``query_start_loc`` entries, which is what the in-kernel
105
+ boundary search has to get right. MTP with partial acceptance produces
106
+ exactly this: some requests contribute one row, others several, in the
107
+ same batch.
108
+ """
109
+ torch.manual_seed(seed)
110
+ device = "cuda"
111
+ world_size, h_per_rank, head_dim = 8, 2, 512
112
+ num_seqs = 17
113
+
114
+ query_lens = torch.randint(1, 5, (num_seqs,), dtype=torch.int32)
115
+ query_start_loc = torch.zeros(num_seqs + 1, dtype=torch.int32)
116
+ query_start_loc[1:] = torch.cumsum(query_lens, 0)
117
+ # Mix real, empty, and (via padding rows) out-of-range shards.
118
+ seq_lens = torch.randint(0, 2, (num_seqs,), dtype=torch.int32) * 16
119
+ num_rows = int(query_start_loc[-1]) + 3
120
+
121
+ query_start_loc = query_start_loc.to(device)
122
+ seq_lens = seq_lens.to(device)
123
+ num_heads = world_size * h_per_rank
124
+ out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype)
125
+ lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32)
126
+
127
+ ref_lse = lse.clone()
128
+ mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc)
129
+ expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None)
130
+ actual = _pack(
131
+ out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc
132
+ )
133
+
134
+ assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), (
135
+ f"ragged query_start_loc mismatch: query_lens={query_lens.tolist()}, "
136
+ f"seq_lens={seq_lens.tolist()}"
137
+ )
138
+
139
+
140
+def test_mask_disabled_is_unmasked():
141
+ """Passing no seq_lens/query_start_loc must leave the LSE untouched."""
142
+ torch.manual_seed(0)
143
+ device = "cuda"
144
+ world_size, h_per_rank, head_dim = 8, 2, 512
145
+ out = torch.randn(
146
+ 4, world_size * h_per_rank, head_dim, device=device, dtype=torch.bfloat16
147
+ )
148
+ lse = torch.randn(4, world_size * h_per_rank, device=device, dtype=torch.float32)
149
+
150
+ packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None)
151
+ assert torch.isfinite(packed.float()).all()
vllm/v1/attention/ops/dcp.py
+62
-5
@@ -45,13 +45,18 @@ if TYPE_CHECKING:
45
# LSE/output combine
46
47
48
-def mask_dcp_empty_shards_(
49
- lse: torch.Tensor,
48
+def _validate_dcp_empty_shard_args(
49
seq_lens: torch.Tensor | None,
50
query_start_loc: torch.Tensor | None,
52
-) -> None:
51
+) -> bool:
52
+ """Validate the empty-shard mask inputs.
53
+
54
+ Returns True when masking is requested, False when it is disabled. Raises on
55
+ an inconsistent pair. Performs no device work so it is safe to call from a
56
+ CUDA-graph-captured region.
57
+ """
58
if seq_lens is None and query_start_loc is None:
54
- return
59
+ return False
60
if seq_lens is None or query_start_loc is None:
61
raise ValueError("seq_lens and query_start_loc must be provided together")
62
if (
@@ -60,6 +65,17 @@ def mask_dcp_empty_shards_(
65
or query_start_loc.shape[0] != seq_lens.shape[0] + 1
66
):
67
raise ValueError("query_start_loc must contain one boundary per sequence")
68
+ return True
69
+
70
+
71
+def mask_dcp_empty_shards_(
72
+ lse: torch.Tensor,
73
+ seq_lens: torch.Tensor | None,
74
+ query_start_loc: torch.Tensor | None,
75
+) -> None:
76
+ if not _validate_dcp_empty_shard_args(seq_lens, query_start_loc):
77
+ return
78
+ assert seq_lens is not None and query_start_loc is not None
79
80
row_indices = torch.arange(
81
lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
@@ -603,6 +619,9 @@ def _dcp_a2a_pack_send_kernel(
619
out_ptr,
620
lse_ptr,
621
send_ptr,
622
+ seq_lens_ptr,
623
+ query_start_loc_ptr,
624
+ num_seqs,
625
out_stride_B,
626
out_stride_H,
627
out_stride_D,
@@ -616,11 +635,33 @@ def _dcp_a2a_pack_send_kernel(
635
HEAD_DIM: tl.constexpr,
636
H_PER_RANK: tl.constexpr,
637
LSE_PACK_DIM: tl.constexpr,
638
+ HAS_MASK: tl.constexpr,
639
+ NUM_SEQS_POW2: tl.constexpr,
640
):
641
batch_idx = tl.program_id(0).to(tl.int64)
642
local_head_idx = tl.program_id(1).to(tl.int64)
643
d_offsets = tl.arange(0, HEAD_DIM)
644
645
+ # Empty-shard masking, fused in to avoid a separate per-layer pass over the
646
+ # LSE. Rows that belong to a request with no local KV shard, and rows past
647
+ # the end of the batch (CUDA-graph padding), must not contribute to the
648
+ # combine, which the reduction achieves by giving them an LSE of -inf.
649
+ row_is_empty = False
650
+ if HAS_MASK:
651
+ # searchsorted(query_start_loc[1:], batch_idx, right=True), i.e. the
652
+ # number of request boundaries at or below this row.
653
+ seq_offsets = tl.arange(0, NUM_SEQS_POW2)
654
+ boundaries = tl.load(
655
+ query_start_loc_ptr + 1 + seq_offsets,
656
+ mask=seq_offsets < num_seqs,
657
+ other=2**31 - 1,
658
+ )
659
+ seq_idx = tl.sum((boundaries <= batch_idx).to(tl.int32), axis=0)
660
+ seq_idx = tl.minimum(seq_idx, num_seqs - 1)
661
+ num_rows = tl.load(query_start_loc_ptr + num_seqs)
662
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
663
+ row_is_empty = (batch_idx >= num_rows) | (seq_len == 0)
664
+
665
for rank_idx in tl.static_range(N):
666
src_head_idx = rank_idx * H_PER_RANK + local_head_idx
667
send_base = (
@@ -642,6 +683,8 @@ def _dcp_a2a_pack_send_kernel(
683
lse_val = tl.load(
684
lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
685
).to(tl.float32)
686
+ if HAS_MASK:
687
+ lse_val = tl.where(row_is_empty, -float("inf"), lse_val)
688
if LSE_PACK_DIM == 1:
689
tl.store(
690
send_ptr + send_base + HEAD_DIM * send_stride_D,
@@ -797,12 +840,21 @@ def _dcp_a2a_pack_send(
840
seq_lens: torch.Tensor | None = None,
841
query_start_loc: torch.Tensor | None = None,
842
) -> None:
800
- mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
843
+ # The empty-shard mask is fused into the pack kernel below rather than
844
+ # applied as a separate pass: the mask is a pure function of seq_lens and
845
+ # query_start_loc, which are identical for every MLA layer in a step, while
846
+ # the eager form costs eight launch-bound kernels per layer.
847
+ has_mask = _validate_dcp_empty_shard_args(seq_lens, query_start_loc)
848
+ num_seqs = seq_lens.shape[0] if has_mask and seq_lens is not None else 0
849
+
850
grid = (cp_attn_out.shape[0], h_per_rank, 1)
851
_dcp_a2a_pack_send_kernel[grid](
852
cp_attn_out,
853
cp_attn_lse,
854
send_buffer,
855
+ seq_lens,
856
+ query_start_loc,
857
+ num_seqs,
858
cp_attn_out.stride(0),
859
cp_attn_out.stride(1),
860
cp_attn_out.stride(2),
@@ -816,6 +868,11 @@ def _dcp_a2a_pack_send(
868
HEAD_DIM=head_dim,
869
H_PER_RANK=h_per_rank,
870
LSE_PACK_DIM=lse_pack_dim,
871
+ HAS_MASK=has_mask,
872
+ # Only the arange bound is a constexpr: num_seqs itself is a runtime
873
+ # argument so the kernel specializes on at most log2(max_num_seqs)
874
+ # variants instead of one per distinct batch size.
875
+ NUM_SEQS_POW2=triton.next_power_of_2(max(num_seqs, 1)),
876
)
877
878