| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU tests for the kpool tail slot mapping (no GPU required). |
| 4 | |
| 5 | The kpool tail cache is a 1-block-per-request circular ring addressed by |
| 6 | ``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block |
| 7 | allocated per request, never grown, so only column 0 of its block table is |
| 8 | ever written; the rest stays zero-initialized). |
| 9 | |
| 10 | The generic per-group slot kernel cannot express that layout: it maps |
| 11 | ``pos -> bt[req][pos // bs] * bs + pos % bs`` (``_compute_slot_mappings_kernel`` |
| 12 | in vllm/v1/worker/gpu/block_table.py), so every token at ``pos >= kpool`` |
| 13 | reads a zero column and collapses onto physical tail block 0. All concurrent |
| 14 | requests then share one ``kpool``-slot ring and corrupt each other's pool |
| 15 | compression. |
| 16 | |
| 17 | These tests pin that defect's arithmetic, verify the circular replacement |
| 18 | (``compute_kpool_tail_slot_mapping``), and mirror the tail kernels' index |
| 19 | math to show cross-request pollution before the fix and isolation after. |
| 20 | """ |
| 21 | |
| 22 | from types import SimpleNamespace |
| 23 | |
| 24 | import pytest |
| 25 | import torch |
| 26 | |
| 27 | from vllm.config import VllmConfig, set_current_vllm_config |
| 28 | from vllm.models.glm5next.common.attention import Glm5NextTailCache |
| 29 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 30 | from vllm.v1.attention.backends.mla.indexer import ( |
| 31 | KpoolTailBackend, |
| 32 | KpoolTailMetadataBuilder, |
| 33 | compute_kpool_tail_slot_mapping, |
| 34 | ) |
| 35 | from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides |
| 36 | from vllm.v1.kv_cache_layout import KVCacheLayout |
| 37 | from vllm.v1.worker.block_table import get_block_table_width |
| 38 | |
| 39 | KPOOL = 4 |
| 40 | |
| 41 | |
| 42 | def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): |
| 43 | (layout,) = KpoolTailBackend.supported_kv_cache_layouts() |
| 44 | spec = KpoolTailSpec( |
| 45 | block_size=KPOOL, |
| 46 | num_kv_heads=2, |
| 47 | head_size=128, |
| 48 | head_size_v=0, |
| 49 | dtype=torch.bfloat16, |
| 50 | sliding_window=KPOOL, |
| 51 | ) |
| 52 | strides = compute_layout_strides(spec, num_blocks=8, num_layers=3, layout=layout) |
| 53 | _, _, head_stride, state_stride, content_stride = strides |
| 54 | |
| 55 | assert layout is KVCacheLayout.LBHNC |
| 56 | assert head_stride == KPOOL * 128 * torch.bfloat16.itemsize |
| 57 | assert state_stride == 128 * torch.bfloat16.itemsize |
| 58 | assert content_stride == 1 |
| 59 | |
| 60 | |
| 61 | @pytest.mark.parametrize( |
| 62 | "num_speculative_tokens, ring", [(0, 4), (1, 8), (4, 8), (7, 16), (13, 32)] |
| 63 | ) |
| 64 | def test_tail_ring_divides_the_attention_block(num_speculative_tokens, ring): |
| 65 | """With 7 draft tokens the ring was 12 slots. 12 does not divide the |
| 66 | 640-token KDA attention block, so the scheduler block grew to their lcm |
| 67 | and prefix-cache hits were cut down to multiples of 1920.""" |
| 68 | with set_current_vllm_config(VllmConfig()): |
| 69 | cache = Glm5NextTailCache( |
| 70 | head_dim=128, |
| 71 | dtype=torch.bfloat16, |
| 72 | prefix="tail", |
| 73 | cache_config=SimpleNamespace(block_size=640), |
| 74 | index_kpool=KPOOL, |
| 75 | ) |
| 76 | spec = cache.get_kv_cache_spec( |
| 77 | SimpleNamespace(num_speculative_tokens=num_speculative_tokens) |
| 78 | ) |
| 79 | |
| 80 | assert spec.block_size == spec.sliding_window == ring |
| 81 | assert ring >= KPOOL + num_speculative_tokens |
| 82 | assert 640 % ring == 0 |
| 83 | |
| 84 | |
| 85 | def test_tail_spec_opts_out_of_generic_slot_mapping(): |
| 86 | """The tail row is one block wide (padded to the block-table alignment), so |
| 87 | the generic kernel's ``pos // kpool`` column index runs off the end of the |
| 88 | allocation for long prompts. The spec must opt out of it entirely.""" |
| 89 | spec = KpoolTailSpec( |
| 90 | block_size=KPOOL, |
| 91 | num_kv_heads=2, |
| 92 | head_size=128, |
| 93 | head_size_v=0, |
| 94 | dtype=torch.bfloat16, |
| 95 | sliding_window=KPOOL, |
| 96 | ) |
| 97 | max_len = 1 << 20 |
| 98 | width = get_block_table_width( |
| 99 | spec.max_num_blocks_per_req(None, max_len), |
| 100 | spec.block_size, |
| 101 | token_alignment=spec.block_table_token_alignment, |
| 102 | ) |
| 103 | |
| 104 | assert width * KPOOL < max_len |
| 105 | assert spec.uses_slot_mapping is False |
| 106 | |
| 107 | |
| 108 | def make_tail_block_table(own_blocks, width=64): |
| 109 | """Tail-group block table as BlockTables produces it: column 0 holds the |
| 110 | request's single KpoolTailManager block, the remaining columns are never |
| 111 | written and stay zero.""" |
| 112 | bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) |
| 113 | bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) |
| 114 | return bt |
| 115 | |
| 116 | |
| 117 | def legacy_generic_tail_slots(block_table, query_start_loc, positions): |
| 118 | """Reference of the generic ``_compute_slot_mappings_kernel`` arithmetic |
| 119 | (block_table.py:305-313) applied to the tail group's table.""" |
| 120 | slots = [] |
| 121 | for req in range(block_table.shape[0]): |
| 122 | for i in range(query_start_loc[req], query_start_loc[req + 1]): |
| 123 | pos = int(positions[i]) |
| 124 | block_number = int(block_table[req, pos // KPOOL]) |
| 125 | slots.append(block_number * KPOOL + pos % KPOOL) |
| 126 | return torch.tensor(slots, dtype=torch.int64) |
| 127 | |
| 128 | |
| 129 | def circular_tail_slots( |
| 130 | slot_mapping, block_table, query_start_loc, positions, num_actual, num_reqs |
| 131 | ): |
| 132 | return compute_kpool_tail_slot_mapping( |
| 133 | slot_mapping, |
| 134 | block_table, |
| 135 | query_start_loc, |
| 136 | positions, |
| 137 | num_actual, |
| 138 | num_reqs, |
| 139 | KPOOL, |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | def make_batch(per_req_positions, padded_len=None): |
| 144 | positions = torch.cat( |
| 145 | [torch.tensor(p, dtype=torch.int64) for p in per_req_positions] |
| 146 | ) |
| 147 | num_actual = positions.numel() |
| 148 | num_reqs = len(per_req_positions) |
| 149 | lens = [len(p) for p in per_req_positions] |
| 150 | qsl = torch.zeros(num_reqs + 1, dtype=torch.int64) |
| 151 | torch.cumsum(torch.tensor(lens, dtype=torch.int64), 0, out=qsl[1:]) |
| 152 | if padded_len is None: |
| 153 | padded_len = num_actual |
| 154 | slot_mapping = torch.full((padded_len,), -1, dtype=torch.int64) |
| 155 | return positions, qsl, slot_mapping, num_actual, num_reqs |
| 156 | |
| 157 | |
| 158 | def test_legacy_generic_mapping_collapses_onto_block_zero(): |
| 159 | """The bug: with the manager's 1-column block table, the generic kernel |
| 160 | maps every pos >= kpool onto tail block 0, and distinct requests collide.""" |
| 161 | own_blocks = [5, 9] |
| 162 | per_req = [list(range(10)), list(range(12))] # prompts of len 10 and 12 |
| 163 | positions, qsl, _, num_actual, num_reqs = make_batch(per_req) |
| 164 | bt = make_tail_block_table(own_blocks) |
| 165 | |
| 166 | legacy = legacy_generic_tail_slots(bt, qsl, positions) |
| 167 | |
| 168 | # Every token at pos >= kpool resolves to block 0, not the request's own. |
| 169 | off = 0 |
| 170 | for req, prompt in enumerate(per_req): |
| 171 | req_slots = legacy[off : off + len(prompt)] |
| 172 | for pos in range(len(prompt)): |
| 173 | slot = int(req_slots[pos]) |
| 174 | if pos >= KPOOL: |
| 175 | assert slot // KPOOL == 0, ( |
| 176 | f"expected collapse onto block 0 at pos {pos}" |
| 177 | ) |
| 178 | assert slot // KPOOL != own_blocks[req] |
| 179 | off += len(prompt) |
| 180 | |
| 181 | # The two requests share ring slots -> cross-request pollution. |
| 182 | a_slots = set(legacy[: len(per_req[0])].tolist()) |
| 183 | b_slots = set(legacy[len(per_req[0]) :].tolist()) |
| 184 | assert a_slots & b_slots, "legacy mapping must collide across requests" |
| 185 | |
| 186 | |
| 187 | def test_circular_mapping_isolates_requests(): |
| 188 | """The fix: every token lands in its own request's block at pos % kpool, |
| 189 | and no slot is ever shared by two different requests (slots do recur |
| 190 | within a request every kpool positions -- that is the circular design).""" |
| 191 | own_blocks = [5, 9] |
| 192 | per_req = [list(range(10)), list(range(12))] |
| 193 | positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) |
| 194 | bt = make_tail_block_table(own_blocks) |
| 195 | |
| 196 | out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) |
| 197 | |
| 198 | off = 0 |
| 199 | per_req_slots = [] |
| 200 | for req, prompt in enumerate(per_req): |
| 201 | req_slots = set() |
| 202 | for pos in range(len(prompt)): |
| 203 | slot = int(out[off + pos]) |
| 204 | assert slot // KPOOL == own_blocks[req], ( |
| 205 | f"req {req} pos {pos} left its tail block" |
| 206 | ) |
| 207 | assert slot % KPOOL == pos % KPOOL |
| 208 | req_slots.add(slot) |
| 209 | per_req_slots.append(req_slots) |
| 210 | off += len(prompt) |
| 211 | assert not per_req_slots[0] & per_req_slots[1] |
| 212 | |
| 213 | |
| 214 | @pytest.mark.parametrize("prompt_len", [1, 2, 3, 4]) |
| 215 | def test_circular_mapping_matches_generic_for_short_requests(prompt_len): |
| 216 | """For pos < kpool the generic kernel already picks the own block, so the |
| 217 | two mappings agree while every position fits the request's first block |
| 218 | (single-request behavior is unchanged).""" |
| 219 | own_blocks = [7] |
| 220 | per_req = [list(range(prompt_len))] |
| 221 | positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) |
| 222 | bt = make_tail_block_table(own_blocks) |
| 223 | |
| 224 | legacy = legacy_generic_tail_slots(bt, qsl, positions) |
| 225 | out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) |
| 226 | assert torch.equal(out, legacy) |
| 227 | |
| 228 | |
| 229 | def test_circular_mapping_preserves_padding_and_empty_batch(): |
| 230 | own_blocks = [5, 9] |
| 231 | per_req = [list(range(10)), list(range(12))] |
| 232 | padded_len = sum(len(p) for p in per_req) + 8 |
| 233 | positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( |
| 234 | per_req, padded_len=padded_len |
| 235 | ) |
| 236 | bt = make_tail_block_table(own_blocks) |
| 237 | |
| 238 | out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) |
| 239 | assert out.shape == slot_mapping.shape |
| 240 | assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1)) |
| 241 | |
| 242 | empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs) |
| 243 | assert torch.equal(empty, slot_mapping) |
| 244 | |
| 245 | |
| 246 | def make_common_metadata(per_req_positions, own_blocks, with_positions=True): |
| 247 | positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( |
| 248 | per_req_positions, padded_len=sum(len(p) for p in per_req_positions) + 4 |
| 249 | ) |
| 250 | bt = make_tail_block_table(own_blocks) |
| 251 | seq_lens = torch.tensor( |
| 252 | [max(p) + 1 if p else 1 for p in per_req_positions], dtype=torch.int64 |
| 253 | ) |
| 254 | return CommonAttentionMetadata( |
| 255 | query_start_loc=qsl, |
| 256 | query_start_loc_cpu=qsl.clone(), |
| 257 | seq_lens=seq_lens, |
| 258 | num_reqs=num_reqs, |
| 259 | num_actual_tokens=num_actual, |
| 260 | max_query_len=max((len(p) for p in per_req_positions), default=1), |
| 261 | max_seq_len=int(seq_lens.max()) if num_reqs else 1, |
| 262 | block_table_tensor=bt, |
| 263 | slot_mapping=slot_mapping, |
| 264 | positions=positions if with_positions else None, |
| 265 | ) |
| 266 | |
| 267 | |
| 268 | def make_tail_builder(block_size=KPOOL, max_num_batched_tokens=128): |
| 269 | builder = object.__new__(KpoolTailMetadataBuilder) |
| 270 | builder.kv_cache_spec = SimpleNamespace(block_size=block_size) |
| 271 | builder.slot_mapping_buffer = torch.empty(max_num_batched_tokens, dtype=torch.int64) |
| 272 | return builder |
| 273 | |
| 274 | |
| 275 | def test_builder_build_uses_circular_mapping(): |
| 276 | per_req = [list(range(10)), list(range(12))] |
| 277 | own_blocks = [5, 9] |
| 278 | cam = make_common_metadata(per_req, own_blocks) |
| 279 | meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam) |
| 280 | |
| 281 | out = meta.slot_mapping |
| 282 | off = 0 |
| 283 | for req, prompt in enumerate(per_req): |
| 284 | for pos in range(len(prompt)): |
| 285 | slot = int(out[off + pos]) |
| 286 | assert slot // KPOOL == own_blocks[req] |
| 287 | assert slot % KPOOL == pos % KPOOL |
| 288 | off += len(prompt) |
| 289 | # Padding tail of the buffer keeps the -1 sentinel. |
| 290 | assert torch.equal( |
| 291 | out[cam.num_actual_tokens :], torch.full_like(out[cam.num_actual_tokens :], -1) |
| 292 | ) |
| 293 | |
| 294 | |
| 295 | def test_builder_build_falls_back_without_positions(): |
| 296 | """Capture / dummy builds without positions keep the generic mapping.""" |
| 297 | per_req = [list(range(10))] |
| 298 | cam = make_common_metadata(per_req, [5], with_positions=False) |
| 299 | meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam) |
| 300 | assert meta.slot_mapping is cam.slot_mapping |
| 301 | |
| 302 | |
| 303 | def test_builder_reuses_slot_mapping_storage(): |
| 304 | builder = make_tail_builder() |
| 305 | first = make_common_metadata([list(range(10))], [5]) |
| 306 | first_meta = KpoolTailMetadataBuilder.build(builder, 0, first) |
| 307 | data_ptr = first_meta.slot_mapping.data_ptr() |
| 308 | |
| 309 | second = make_common_metadata([list(range(12))], [9]) |
| 310 | second_meta = KpoolTailMetadataBuilder.build(builder, 0, second) |
| 311 | |
| 312 | assert second_meta.slot_mapping.data_ptr() == data_ptr |
| 313 | assert second_meta.slot_mapping[:12].tolist() == [ |
| 314 | 9 * KPOOL + pos % KPOOL for pos in range(12) |
| 315 | ] |
| 316 | |
| 317 | |
| 318 | # --------------------------------------------------------------------------- |
| 319 | # Index-level mirror of the tail kernels: seed / stash / pool completion |
| 320 | # (addressing replicated from kpool_compress.py's Triton kernels). |
| 321 | # --------------------------------------------------------------------------- |
| 322 | |
| 323 | |
| 324 | class TailRingMirror: |
| 325 | """Mirror of the tail-ring addressing in _kpool_tail_seed_kernel and |
| 326 | _kpool_decode_update_batched_kernel.""" |
| 327 | |
| 328 | def __init__(self, num_blocks, kpool=KPOOL, ring=None): |
| 329 | self.kpool = kpool |
| 330 | self.ring = ring or kpool |
| 331 | self.k = torch.full((num_blocks, self.ring, 3), float("nan")) |
| 332 | self.s = torch.full((num_blocks, self.ring, 3), float("nan")) |
| 333 | |
| 334 | def stash(self, tail_slot, pos, k, s): |
| 335 | blk, off = tail_slot // self.ring, pos % self.ring |
| 336 | self.k[blk, off] = k |
| 337 | self.s[blk, off] = s |
| 338 | |
| 339 | seed = stash # the seed kernel writes with the same addressing |
| 340 | |
| 341 | def complete(self, tail_slot, pos, k, s): |
| 342 | blk = tail_slot // self.ring |
| 343 | start = pos - (self.kpool - 1) |
| 344 | kk = torch.stack( |
| 345 | [self.k[blk, (start + i) % self.ring] for i in range(self.kpool)] |
| 346 | ) |
| 347 | ss = torch.stack( |
| 348 | [self.s[blk, (start + i) % self.ring] for i in range(self.kpool)] |
| 349 | ) |
| 350 | kk[-1], ss[-1] = k, s # is_current for the completing token |
| 351 | w = torch.softmax(ss, dim=0) |
| 352 | return (kk * w).sum(0) |
| 353 | |
| 354 | |
| 355 | def token_kv(req, pos): |
| 356 | k = torch.tensor([pos + 100.0 * req, pos + 0.5, 2.0 * pos + 0.25]) |
| 357 | s = torch.tensor([0.1 * (pos + 1) + req, 0.2 * pos, 0.05 * pos]) |
| 358 | return k, s |
| 359 | |
| 360 | |
| 361 | def tail_slot_for(mapping, req, pos, own_block): |
| 362 | if mapping == "legacy": |
| 363 | bt_val = own_block if pos < KPOOL else 0 |
| 364 | return bt_val * KPOOL + pos % KPOOL |
| 365 | return own_block * KPOOL + pos % KPOOL |
| 366 | |
| 367 | |
| 368 | def run_scenario(mapping, interleave): |
| 369 | """Requests A (block 5, prompt len 9) and B (block 9, prompt len 11) |
| 370 | decode concurrently; returns A's boundary pool [8, 9, 10, 11].""" |
| 371 | ring = TailRingMirror(num_blocks=16) |
| 372 | blocks = {"A": 5, "B": 9} |
| 373 | prompts = {"A": 9, "B": 11} |
| 374 | |
| 375 | def slot(req, pos): |
| 376 | return tail_slot_for(mapping, 0 if req == "A" else 1, pos, blocks[req]) |
| 377 | |
| 378 | # Prefill: seed each request's trailing incomplete pool. |
| 379 | for req, L in prompts.items(): |
| 380 | for pos in range(L - (L % KPOOL or KPOOL), L): |
| 381 | if pos < 0: |
| 382 | continue |
| 383 | ring.stash(slot(req, pos), pos, *token_kv(0 if req == "A" else 1, pos)) |
| 384 | |
| 385 | # Decode: A emits pos 9, 10, 11; B emits 11, 12, 13. `interleave` |
| 386 | # processes B before A within a step, which is exactly what concurrent |
| 387 | # Triton programs do when they share tail block 0. |
| 388 | order = ["B", "A"] if interleave else ["A", "B"] |
| 389 | decode = {"A": [9, 10, 11], "B": [11, 12, 13]} |
| 390 | result = None |
| 391 | for step in range(3): |
| 392 | for req in order: |
| 393 | pos = decode[req][step] |
| 394 | k, s = token_kv(0 if req == "A" else 1, pos) |
| 395 | if pos % KPOOL == KPOOL - 1: |
| 396 | pool = ring.complete(slot(req, pos), pos, k, s) |
| 397 | if req == "A": |
| 398 | result = pool |
| 399 | ring.stash(slot(req, pos), pos, k, s) |
| 400 | return result |
| 401 | |
| 402 | |
| 403 | def test_interleaved_decode_pollution_legacy_vs_circular(): |
| 404 | """Ground truth: request A decoding alone (its ring used exclusively).""" |
| 405 | torch.manual_seed(0) |
| 406 | ground_truth = run_scenario("circular", interleave=False) |
| 407 | |
| 408 | legacy = run_scenario("legacy", interleave=True) |
| 409 | circular = run_scenario("circular", interleave=True) |
| 410 | |
| 411 | # The old mapping lets request B's tokens into A's ring: A's boundary |
| 412 | # pool is compressed from 2 of B's tokens -> wrong. |
| 413 | assert not torch.allclose(legacy, ground_truth), ( |
| 414 | f"legacy mapping unexpectedly clean: {legacy} vs {ground_truth}" |
| 415 | ) |
| 416 | |
| 417 | # The circular mapping keeps the rings isolated under interleaving. |
| 418 | torch.testing.assert_close(circular, ground_truth) |
| 419 | |
| 420 | |
| 421 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") |
| 422 | @pytest.mark.parametrize( |
| 423 | "per_req,num_actual,padded_len", |
| 424 | [ |
| 425 | ([list(range(10)), list(range(12))], 22, 22), |
| 426 | ([list(range(10)), list(range(12))], 22, 30), |
| 427 | ([[3, 4], [0], [7, 8, 9]], 6, 8), |
| 428 | ([[5]], 1, 1), |
| 429 | ], |
| 430 | ) |
| 431 | def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): |
| 432 | """The CUDA (Triton) path must match the CPU torch reference, including |
| 433 | tokens between the last request boundary and num_actual_tokens (mapped to |
| 434 | the last request) and untouched padding beyond num_actual.""" |
| 435 | positions, qsl, slot_mapping, _, num_reqs = make_batch( |
| 436 | per_req, padded_len=padded_len |
| 437 | ) |
| 438 | # Replace the all--1 placeholder slots with sentinel values to check the |
| 439 | # padding range is copied through untouched. |
| 440 | slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 |
| 441 | bt = make_tail_block_table(list(range(5, 5 + num_reqs))) |
| 442 | |
| 443 | ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) |
| 444 | got = circular_tail_slots( |
| 445 | slot_mapping.cuda(), |
| 446 | bt.cuda(), |
| 447 | qsl.cuda().to(torch.int32), |
| 448 | positions.cuda(), |
| 449 | num_actual, |
| 450 | num_reqs, |
| 451 | ) |
| 452 | torch.testing.assert_close(got.cpu(), ref) |
| 453 | |
| 454 | |
| 455 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") |
| 456 | def test_triton_mapping_reads_strided_block_table(): |
| 457 | """The kernel must address the block table through its real stride(0), |
| 458 | not a dense assumption (the #57477 class of bug): a tail block table that |
| 459 | is a column view of a wider allocation must still read column 0.""" |
| 460 | per_req = [list(range(10)), list(range(12))] |
| 461 | positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req) |
| 462 | own_blocks = [11, 22] |
| 463 | wide = torch.zeros(num_reqs, 8, dtype=torch.int32) |
| 464 | bt = wide[:, 1:4] # non-contiguous view, stride(0) == 8 |
| 465 | bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) |
| 466 | |
| 467 | ref = circular_tail_slots( |
| 468 | slot_mapping, bt.contiguous(), qsl, positions, num_actual, num_reqs |
| 469 | ) |
| 470 | got = circular_tail_slots( |
| 471 | slot_mapping.cuda(), |
| 472 | bt.cuda(), |
| 473 | qsl.cuda().to(torch.int32), |
| 474 | positions.cuda(), |
| 475 | num_actual, |
| 476 | num_reqs, |
| 477 | ) |
| 478 | torch.testing.assert_close(got.cpu(), ref) |
| 479 | for req, blk in enumerate(own_blocks): |
| 480 | start, end = int(qsl[req]), int(qsl[req + 1]) |
| 481 | assert (got[start:end] // KPOOL == blk).all() |
| 482 | |
| 483 | |
| 484 | @pytest.mark.parametrize("ring_pools", [1, 2]) |
| 485 | def test_rejected_completing_draft_needs_ring_slots(ring_pools): |
| 486 | """With a one-pool ring, the drafts behind a rejected pool-completing draft |
| 487 | overwrote the pool's earlier keys, so its redo compressed wrong keys.""" |
| 488 | ring_size = ring_pools * KPOOL |
| 489 | truth = TailRingMirror(num_blocks=2, ring=ring_size) |
| 490 | ring = TailRingMirror(num_blocks=2, ring=ring_size) |
| 491 | block = 1 |
| 492 | |
| 493 | def slot(pos): |
| 494 | return block * ring_size + pos % ring_size |
| 495 | |
| 496 | for pos in range(4, 7): |
| 497 | truth.stash(slot(pos), pos, *token_kv(0, pos)) |
| 498 | ring.stash(slot(pos), pos, *token_kv(0, pos)) |
| 499 | expected = truth.complete(slot(7), 7, *token_kv(0, 7)) |
| 500 | |
| 501 | # Draft 7 completes the pool and is rejected. With a one-pool ring, drafts |
| 502 | # 8..10 overwrite the slots of positions 4..6, which are read by the redo |
| 503 | # of 7. |
| 504 | for pos in range(7, 11): |
| 505 | k, s = token_kv(9, pos) # draft values |
| 506 | if pos % KPOOL == KPOOL - 1: |
| 507 | ring.complete(slot(pos), pos, k, s) |
| 508 | ring.stash(slot(pos), pos, k, s) |
| 509 | redo = ring.complete(slot(7), 7, *token_kv(0, 7)) |
| 510 | if ring_pools == 1: |
| 511 | assert not torch.allclose(redo, expected) |
| 512 | else: |
| 513 | torch.testing.assert_close(redo, expected) |