| 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 | requires_accelerator = 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 | @requires_accelerator |
| 48 | @pytest.mark.parametrize( |
| 49 | "seq_lens_list", |
| 50 | [ |
| 51 | [16, 16, 16, 16], # no empty shards (the long-context decode case) |
| 52 | [16, 0, 16, 0], # interior empty shards |
| 53 | [0, 0, 0, 0], # every shard empty |
| 54 | [16, 16, 0, 0], # trailing empty shards, i.e. cudagraph padding |
| 55 | [0, 16, 16, 16], # leading empty shard |
| 56 | [16] * 52, # the profiled conc-52 decode batch |
| 57 | [7], # single request |
| 58 | ], |
| 59 | ) |
| 60 | def test_fused_mask_matches_eager( |
| 61 | dtype, world_size, h_per_rank, tokens_per_req, num_pad_rows, seq_lens_list |
| 62 | ): |
| 63 | torch.manual_seed(0) |
| 64 | device = "cuda" |
| 65 | head_dim = 512 |
| 66 | num_seqs = len(seq_lens_list) |
| 67 | # tokens_per_req > 1 is the MTP / multi-token-verify shape. |
| 68 | num_rows = num_seqs * tokens_per_req + num_pad_rows |
| 69 | num_heads = world_size * h_per_rank |
| 70 | |
| 71 | out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype) |
| 72 | lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32) |
| 73 | seq_lens = torch.tensor(seq_lens_list, device=device, dtype=torch.int32) |
| 74 | query_start_loc = torch.arange( |
| 75 | 0, |
| 76 | (num_seqs + 1) * tokens_per_req, |
| 77 | tokens_per_req, |
| 78 | device=device, |
| 79 | dtype=torch.int32, |
| 80 | ) # rows past query_start_loc[-1] are cudagraph padding |
| 81 | |
| 82 | # Reference: mask eagerly, then pack with masking disabled. |
| 83 | ref_lse = lse.clone() |
| 84 | mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc) |
| 85 | expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None) |
| 86 | |
| 87 | # Under test: pack with the mask fused in. |
| 88 | actual = _pack( |
| 89 | out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc |
| 90 | ) |
| 91 | |
| 92 | # The packed LSE slots hold halves of an fp32, not meaningful floats, so |
| 93 | # compare the send buffers bitwise: -inf's high half is a NaN in fp16. |
| 94 | assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), ( |
| 95 | "fused mask is not bit-exact with the eager mask" |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 100 | @pytest.mark.parametrize("seed", range(8)) |
| 101 | @requires_accelerator |
| 102 | def test_fused_mask_matches_eager_ragged(dtype, seed): |
| 103 | """Non-uniform per-request query lengths. |
| 104 | |
| 105 | The uniform-stride cases above never exercise an unequal gap between |
| 106 | consecutive ``query_start_loc`` entries, which is what the in-kernel |
| 107 | boundary search has to get right. MTP with partial acceptance produces |
| 108 | exactly this: some requests contribute one row, others several, in the |
| 109 | same batch. |
| 110 | """ |
| 111 | torch.manual_seed(seed) |
| 112 | device = "cuda" |
| 113 | world_size, h_per_rank, head_dim = 8, 2, 512 |
| 114 | num_seqs = 17 |
| 115 | |
| 116 | query_lens = torch.randint(1, 5, (num_seqs,), dtype=torch.int32) |
| 117 | query_start_loc = torch.zeros(num_seqs + 1, dtype=torch.int32) |
| 118 | query_start_loc[1:] = torch.cumsum(query_lens, 0) |
| 119 | # Mix real, empty, and (via padding rows) out-of-range shards. |
| 120 | seq_lens = torch.randint(0, 2, (num_seqs,), dtype=torch.int32) * 16 |
| 121 | num_rows = int(query_start_loc[-1]) + 3 |
| 122 | |
| 123 | query_start_loc = query_start_loc.to(device) |
| 124 | seq_lens = seq_lens.to(device) |
| 125 | num_heads = world_size * h_per_rank |
| 126 | out = torch.randn(num_rows, num_heads, head_dim, device=device, dtype=dtype) |
| 127 | lse = torch.randn(num_rows, num_heads, device=device, dtype=torch.float32) |
| 128 | |
| 129 | ref_lse = lse.clone() |
| 130 | mask_dcp_empty_shards_(ref_lse, seq_lens, query_start_loc) |
| 131 | expected = _pack(out, ref_lse, world_size, h_per_rank, head_dim, None, None) |
| 132 | actual = _pack( |
| 133 | out, lse.clone(), world_size, h_per_rank, head_dim, seq_lens, query_start_loc |
| 134 | ) |
| 135 | |
| 136 | assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)), ( |
| 137 | f"ragged query_start_loc mismatch: query_lens={query_lens.tolist()}, " |
| 138 | f"seq_lens={seq_lens.tolist()}" |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | @requires_accelerator |
| 143 | def test_mask_disabled_is_unmasked(): |
| 144 | """Passing no seq_lens/query_start_loc must leave the LSE untouched.""" |
| 145 | torch.manual_seed(0) |
| 146 | device = "cuda" |
| 147 | world_size, h_per_rank, head_dim = 8, 2, 512 |
| 148 | out = torch.randn( |
| 149 | 4, world_size * h_per_rank, head_dim, device=device, dtype=torch.bfloat16 |
| 150 | ) |
| 151 | lse = torch.randn(4, world_size * h_per_rank, device=device, dtype=torch.float32) |
| 152 | |
| 153 | packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None) |
| 154 | assert torch.isfinite(packed.float()).all() |
| 155 | |
| 156 | |
| 157 | def test_eager_mask_handles_rank_with_no_local_sequences(): |
| 158 | """Padded graph rows are all empty when a DCP rank has no sequences.""" |
| 159 | lse = torch.randn(8, 4, dtype=torch.float32) |
| 160 | seq_lens = torch.empty(0, dtype=torch.int32) |
| 161 | query_start_loc = torch.tensor([0], dtype=torch.int32) |
| 162 | |
| 163 | mask_dcp_empty_shards_(lse, seq_lens, query_start_loc) |
| 164 | |
| 165 | assert torch.isneginf(lse).all() |