| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the fused EAGLE slot mapping kernel.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.v1.spec_decode.utils import ( |
| 10 | PADDING_SLOT_ID, |
| 11 | eagle_step_update_slot_mapping_and_metadata, |
| 12 | ) |
| 13 | |
| 14 | DEVICE_TYPE = current_platform.device_type |
| 15 | |
| 16 | # Skip if no CUDA - Triton kernel requires GPU |
| 17 | pytest.importorskip("triton") |
| 18 | if not current_platform.is_cuda_alike() and not current_platform.is_xpu(): |
| 19 | pytest.skip("CUDA/XPU required for EAGLE kernel tests", allow_module_level=True) |
| 20 | |
| 21 | |
| 22 | def _reference_eagle_step_slot_mapping( |
| 23 | positions_1d: torch.Tensor, |
| 24 | block_table_tensor: torch.Tensor, |
| 25 | seq_lens: torch.Tensor, |
| 26 | block_size: int, |
| 27 | max_model_len: int, |
| 28 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 29 | """Python reference for eagle_step_update_slot_mapping_and_metadata.""" |
| 30 | new_positions = positions_1d + 1 |
| 31 | exceeds_max = new_positions >= max_model_len |
| 32 | clamped_positions = torch.where( |
| 33 | exceeds_max, torch.zeros_like(positions_1d), new_positions |
| 34 | ) |
| 35 | block_numbers = (clamped_positions // block_size).clamp( |
| 36 | max=block_table_tensor.shape[1] - 1 |
| 37 | ) |
| 38 | block_ids = block_table_tensor[ |
| 39 | torch.arange(positions_1d.shape[0], device=positions_1d.device), |
| 40 | block_numbers.long(), |
| 41 | ].long() |
| 42 | slot_mapping = block_ids * block_size + (clamped_positions % block_size) |
| 43 | slot_mapping = torch.where( |
| 44 | exceeds_max, torch.full_like(slot_mapping, PADDING_SLOT_ID), slot_mapping |
| 45 | ) |
| 46 | new_seq_lens = torch.where(exceeds_max, torch.ones_like(seq_lens), seq_lens + 1) |
| 47 | new_seq_lens = new_seq_lens.clamp(max=max_model_len) |
| 48 | return clamped_positions, slot_mapping, new_seq_lens |
| 49 | |
| 50 | |
| 51 | def test_eagle_step_slot_mapping_kernel(): |
| 52 | """Test fused kernel matches Python reference for slot mapping and metadata.""" |
| 53 | device = torch.device(DEVICE_TYPE) |
| 54 | batch_size = 32 |
| 55 | block_size = 16 |
| 56 | max_model_len = 4096 |
| 57 | n_blocks_per_req = (max_model_len + block_size - 1) // block_size |
| 58 | |
| 59 | positions_1d = torch.randint( |
| 60 | 0, max_model_len - 10, (batch_size,), dtype=torch.int64, device=device |
| 61 | ) |
| 62 | block_table_tensor = torch.randint( |
| 63 | 0, 1000, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device |
| 64 | ) |
| 65 | seq_lens = torch.randint(1, 100, (batch_size,), dtype=torch.int32, device=device) |
| 66 | |
| 67 | ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping( |
| 68 | positions_1d.clone(), |
| 69 | block_table_tensor, |
| 70 | seq_lens.clone(), |
| 71 | block_size, |
| 72 | max_model_len, |
| 73 | ) |
| 74 | |
| 75 | out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device) |
| 76 | out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device) |
| 77 | seq_lens_copy = seq_lens.clone() |
| 78 | eagle_step_update_slot_mapping_and_metadata( |
| 79 | positions_1d=positions_1d, |
| 80 | block_table_tensor=block_table_tensor, |
| 81 | seq_lens=seq_lens_copy, |
| 82 | block_size=block_size, |
| 83 | max_model_len=max_model_len, |
| 84 | out_clamped_positions=out_clamped, |
| 85 | out_slot_mapping=out_slot, |
| 86 | ) |
| 87 | |
| 88 | assert torch.equal(out_clamped, ref_clamped), ( |
| 89 | f"clamped: {out_clamped} vs {ref_clamped}" |
| 90 | ) |
| 91 | assert torch.equal(out_slot, ref_slot), f"slot: {out_slot} vs {ref_slot}" |
| 92 | assert torch.equal(seq_lens_copy, ref_seq_lens), ( |
| 93 | f"seq_lens: {seq_lens_copy} vs {ref_seq_lens}" |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def test_eagle_step_slot_mapping_kernel_exceeds_max(): |
| 98 | """Test fused kernel when position exceeds max_model_len.""" |
| 99 | device = torch.device(DEVICE_TYPE) |
| 100 | batch_size = 4 |
| 101 | block_size = 16 |
| 102 | max_model_len = 100 |
| 103 | n_blocks_per_req = (max_model_len + block_size - 1) // block_size |
| 104 | |
| 105 | positions_1d = torch.tensor([50, 98, 99, 100], dtype=torch.int64, device=device) |
| 106 | block_table_tensor = torch.randint( |
| 107 | 0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device |
| 108 | ) |
| 109 | seq_lens = torch.tensor([51, 99, 100, 101], dtype=torch.int32, device=device) |
| 110 | |
| 111 | out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device) |
| 112 | out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device) |
| 113 | eagle_step_update_slot_mapping_and_metadata( |
| 114 | positions_1d=positions_1d, |
| 115 | block_table_tensor=block_table_tensor, |
| 116 | seq_lens=seq_lens, |
| 117 | block_size=block_size, |
| 118 | max_model_len=max_model_len, |
| 119 | out_clamped_positions=out_clamped, |
| 120 | out_slot_mapping=out_slot, |
| 121 | ) |
| 122 | |
| 123 | assert out_clamped[0].item() == 51 |
| 124 | assert out_clamped[1].item() == 99 |
| 125 | assert out_clamped[2].item() == 0 |
| 126 | assert out_clamped[3].item() == 0 |
| 127 | assert out_slot[2].item() == PADDING_SLOT_ID |
| 128 | assert out_slot[3].item() == PADDING_SLOT_ID |
| 129 | assert seq_lens[2].item() == 1 |
| 130 | assert seq_lens[3].item() == 1 |
| 131 | |
| 132 | |
| 133 | def test_eagle_step_slot_mapping_kernel_cudagraph_padding(): |
| 134 | """Test that padding threads write PADDING_SLOT_ID when |
| 135 | input_batch_size > batch_size (cudagraph padding).""" |
| 136 | device = torch.device(DEVICE_TYPE) |
| 137 | batch_size = 4 |
| 138 | input_batch_size = 8 |
| 139 | block_size = 16 |
| 140 | max_model_len = 4096 |
| 141 | n_blocks_per_req = (max_model_len + block_size - 1) // block_size |
| 142 | |
| 143 | positions_1d = torch.tensor([10, 20, 30, 40], dtype=torch.int64, device=device) |
| 144 | block_table_tensor = torch.randint( |
| 145 | 0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device |
| 146 | ) |
| 147 | seq_lens = torch.tensor([11, 21, 31, 41], dtype=torch.int32, device=device) |
| 148 | |
| 149 | ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping( |
| 150 | positions_1d.clone(), |
| 151 | block_table_tensor, |
| 152 | seq_lens.clone(), |
| 153 | block_size, |
| 154 | max_model_len, |
| 155 | ) |
| 156 | |
| 157 | out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device) |
| 158 | out_slot = torch.full((input_batch_size,), -999, dtype=torch.int64, device=device) |
| 159 | seq_lens_copy = seq_lens.clone() |
| 160 | eagle_step_update_slot_mapping_and_metadata( |
| 161 | positions_1d=positions_1d, |
| 162 | block_table_tensor=block_table_tensor, |
| 163 | seq_lens=seq_lens_copy, |
| 164 | block_size=block_size, |
| 165 | max_model_len=max_model_len, |
| 166 | out_clamped_positions=out_clamped, |
| 167 | out_slot_mapping=out_slot, |
| 168 | input_batch_size=input_batch_size, |
| 169 | ) |
| 170 | |
| 171 | # Real slots should match the reference |
| 172 | assert torch.equal(out_clamped, ref_clamped) |
| 173 | assert torch.equal(out_slot[:batch_size], ref_slot) |
| 174 | assert torch.equal(seq_lens_copy, ref_seq_lens) |
| 175 | |
| 176 | # Padding slots should be PADDING_SLOT_ID |
| 177 | for i in range(batch_size, input_batch_size): |
| 178 | assert out_slot[i].item() == PADDING_SLOT_ID |