| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm import _custom_ops as ops |
| 9 | |
| 10 | # DeepSeek V3 MLA dimensions |
| 11 | NOPE_DIM = 512 # NoPE latent dimension (FP8 quantized in cache) |
| 12 | ROPE_DIM = 64 # RoPE dimension (stored as BF16 in cache) |
| 13 | NUM_TILES = 4 # NOPE_DIM / GROUP_SIZE = 512 / 128 |
| 14 | GROUP_SIZE = 128 # FP8 quantization group size (one scale per group) |
| 15 | ENTRY_BYTES = 656 # 512 (FP8) + 16 (4×float32 scales) + 128 (64×BF16 RoPE) |
| 16 | |
| 17 | |
| 18 | def _build_test_case(seq_lens, block_size, seed=42): |
| 19 | """Build a synthetic FP8 cache and compute the expected BF16 output. |
| 20 | |
| 21 | This simulates what concat_and_cache_ds_mla_kernel writes into the |
| 22 | KV cache, then computes what cp_gather_and_upconvert should produce. |
| 23 | |
| 24 | Args: |
| 25 | seq_lens: List of sequence lengths, one per request. |
| 26 | block_size: Number of tokens per physical cache block. |
| 27 | seed: Random seed for reproducibility. |
| 28 | |
| 29 | Returns: |
| 30 | Tuple of (cache, block_table, workspace_starts_t, num_reqs, |
| 31 | total_tokens, expected_output). |
| 32 | |
| 33 | """ |
| 34 | torch.manual_seed(seed) |
| 35 | |
| 36 | num_reqs = len(seq_lens) |
| 37 | total_tokens = sum(seq_lens) |
| 38 | |
| 39 | # workspace_starts[r] = sum of seq_lens[0..r-1] |
| 40 | # This tells the kernel where in the output buffer each request's |
| 41 | # gathered tokens should be written. |
| 42 | workspace_starts = [] |
| 43 | s = 0 |
| 44 | for sl in seq_lens: |
| 45 | workspace_starts.append(s) |
| 46 | s += sl |
| 47 | |
| 48 | # How many physical cache blocks each request needs |
| 49 | blocks_per_req = [math.ceil(s / block_size) for s in seq_lens] |
| 50 | total_blocks = sum(blocks_per_req) |
| 51 | max_blocks = max(blocks_per_req) |
| 52 | |
| 53 | # Block table maps (request, logical_block_idx) -> physical_block_id. |
| 54 | # Here we assign blocks contiguously: request 0 gets blocks [0, 1, ...], |
| 55 | # request 1 gets the next set, etc. |
| 56 | block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda") |
| 57 | block_idx = 0 |
| 58 | for r in range(num_reqs): |
| 59 | for b in range(blocks_per_req[r]): |
| 60 | block_table[r, b] = block_idx |
| 61 | block_idx += 1 |
| 62 | |
| 63 | # The raw paged cache: [num_blocks, block_size, 656] as uint8 |
| 64 | cache = torch.zeros( |
| 65 | total_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda" |
| 66 | ) |
| 67 | # Expected kernel output: [total_tokens, 576] as BF16 |
| 68 | expected = torch.zeros( |
| 69 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 70 | ) |
| 71 | |
| 72 | # Fill each token's cache entry and compute expected output |
| 73 | for r in range(num_reqs): |
| 74 | for t in range(seq_lens[r]): |
| 75 | out_idx = workspace_starts[r] + t |
| 76 | # Map token position -> (physical_block, offset_within_block) |
| 77 | phys = block_table[r, t // block_size].item() |
| 78 | off = t % block_size |
| 79 | |
| 80 | # --- NoPE section: 4 tiles of 128 FP8 values, each with a scale --- |
| 81 | for tile in range(NUM_TILES): |
| 82 | start = tile * GROUP_SIZE |
| 83 | |
| 84 | # Generate random data and quantize to FP8 e4m3 |
| 85 | fp8_vals = torch.randn(GROUP_SIZE, device="cuda").to( |
| 86 | torch.float8_e4m3fn |
| 87 | ) |
| 88 | # Pack FP8 bytes into cache at bytes [start : start+128] |
| 89 | cache[phys, off, start : start + GROUP_SIZE] = fp8_vals.view( |
| 90 | torch.uint8 |
| 91 | ) |
| 92 | |
| 93 | # Random positive scale in [0.1, 2.1] |
| 94 | scale = (torch.rand(1, device="cuda") * 2.0 + 0.1).item() |
| 95 | scale_t = torch.tensor([scale], dtype=torch.float32, device="cuda") |
| 96 | # Pack scale as 4 raw bytes at bytes [512 + tile*4 : ...] |
| 97 | cache[phys, off, NOPE_DIM + tile * 4 : NOPE_DIM + (tile + 1) * 4] = ( |
| 98 | scale_t.view(torch.uint8) |
| 99 | ) |
| 100 | |
| 101 | # Reference dequant: fp8 -> float32, multiply scale, -> bf16. |
| 102 | # This matches the CUDA path: fp8 -> half -> float * scale -> bf16. |
| 103 | # (fp8 -> half is exact, half -> float is exact, so fp8 -> float |
| 104 | # gives the same result regardless of intermediate type.) |
| 105 | expected[out_idx, start : start + GROUP_SIZE] = ( |
| 106 | fp8_vals.float() * scale |
| 107 | ).bfloat16() |
| 108 | |
| 109 | # --- RoPE section: 64 BF16 values, direct copy (no dequant) --- |
| 110 | rope = torch.randn(ROPE_DIM, dtype=torch.bfloat16, device="cuda") |
| 111 | # Pack RoPE bytes into cache at bytes [528 : 656] |
| 112 | cache[phys, off, NOPE_DIM + 16 :] = rope.view(torch.uint8) |
| 113 | # Expected output: exact copy |
| 114 | expected[out_idx, NOPE_DIM:] = rope |
| 115 | |
| 116 | workspace_starts_t = torch.tensor( |
| 117 | workspace_starts, dtype=torch.int32, device="cuda" |
| 118 | ) |
| 119 | |
| 120 | return ( |
| 121 | cache, |
| 122 | block_table, |
| 123 | workspace_starts_t, |
| 124 | num_reqs, |
| 125 | total_tokens, |
| 126 | expected, |
| 127 | ) |
| 128 | |
| 129 | |
| 130 | def _build_test_case_fast(seq_lens, block_size, seed=42): |
| 131 | """Vectorized test-case builder for large sequence lengths. |
| 132 | |
| 133 | Same logic as _build_test_case but uses tensor operations instead of |
| 134 | per-token Python loops, making it practical for seq_lens up to 128K+. |
| 135 | """ |
| 136 | torch.manual_seed(seed) |
| 137 | |
| 138 | num_reqs = len(seq_lens) |
| 139 | total_tokens = sum(seq_lens) |
| 140 | |
| 141 | workspace_starts = [] |
| 142 | s = 0 |
| 143 | for sl in seq_lens: |
| 144 | workspace_starts.append(s) |
| 145 | s += sl |
| 146 | |
| 147 | blocks_per_req = [math.ceil(sl / block_size) for sl in seq_lens] |
| 148 | total_blocks = sum(blocks_per_req) |
| 149 | max_blocks = max(blocks_per_req) |
| 150 | |
| 151 | # Contiguous block allocation |
| 152 | block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda") |
| 153 | block_idx = 0 |
| 154 | for r in range(num_reqs): |
| 155 | for b in range(blocks_per_req[r]): |
| 156 | block_table[r, b] = block_idx |
| 157 | block_idx += 1 |
| 158 | |
| 159 | cache = torch.zeros( |
| 160 | total_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda" |
| 161 | ) |
| 162 | |
| 163 | # Generate all data vectorized |
| 164 | nope_fp8 = torch.randn(total_tokens, NOPE_DIM, device="cuda").to( |
| 165 | torch.float8_e4m3fn |
| 166 | ) |
| 167 | scales = (torch.rand(total_tokens, NUM_TILES, device="cuda") * 2.0 + 0.1).float() |
| 168 | rope = torch.randn(total_tokens, ROPE_DIM, dtype=torch.bfloat16, device="cuda") |
| 169 | |
| 170 | # Compute expected output vectorized (same dequant logic as kernel) |
| 171 | expected = torch.zeros( |
| 172 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 173 | ) |
| 174 | for tile in range(NUM_TILES): |
| 175 | start = tile * GROUP_SIZE |
| 176 | expected[:, start : start + GROUP_SIZE] = ( |
| 177 | nope_fp8[:, start : start + GROUP_SIZE].float() * scales[:, tile : tile + 1] |
| 178 | ).bfloat16() |
| 179 | expected[:, NOPE_DIM:] = rope |
| 180 | |
| 181 | # Build per-token cache entries as [total_tokens, 656] uint8 |
| 182 | token_data = torch.zeros( |
| 183 | total_tokens, ENTRY_BYTES, dtype=torch.uint8, device="cuda" |
| 184 | ) |
| 185 | token_data[:, :NOPE_DIM] = nope_fp8.view(torch.uint8) |
| 186 | token_data[:, NOPE_DIM : NOPE_DIM + 16] = scales.view(torch.uint8) |
| 187 | token_data[:, NOPE_DIM + 16 :] = rope.view(torch.uint8) |
| 188 | |
| 189 | # Scatter into paged cache (loop over requests, not tokens) |
| 190 | block_start = 0 |
| 191 | for r in range(num_reqs): |
| 192 | sl = seq_lens[r] |
| 193 | nb = blocks_per_req[r] |
| 194 | ws = workspace_starts[r] |
| 195 | flat_cache = cache[block_start : block_start + nb].reshape(-1, ENTRY_BYTES) |
| 196 | flat_cache[:sl] = token_data[ws : ws + sl] |
| 197 | block_start += nb |
| 198 | |
| 199 | workspace_starts_t = torch.tensor( |
| 200 | workspace_starts, dtype=torch.int32, device="cuda" |
| 201 | ) |
| 202 | |
| 203 | return ( |
| 204 | cache, |
| 205 | block_table, |
| 206 | workspace_starts_t, |
| 207 | num_reqs, |
| 208 | total_tokens, |
| 209 | expected, |
| 210 | ) |
| 211 | |
| 212 | |
| 213 | @pytest.mark.parametrize( |
| 214 | "seq_lens,block_size", |
| 215 | [ |
| 216 | # Production block_size=64 (only supported value for FlashMLA sparse). |
| 217 | # Realistic prefill scenarios with varying request counts. |
| 218 | ([1], 64), # single token edge case |
| 219 | ([64], 64), # 1 req, exactly one block |
| 220 | ([128], 64), # 1 req, crosses block boundary |
| 221 | ([512], 64), # 1 req, longer prefill |
| 222 | ([256, 128, 384], 64), # 3 reqs, varying lengths |
| 223 | ([128] * 4, 64), # 4 reqs, equal lengths |
| 224 | ([64] * 16, 64), # 16 reqs, shorter prefills |
| 225 | ], |
| 226 | ) |
| 227 | def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size): |
| 228 | """Core correctness test: build cache, run kernel, compare output.""" |
| 229 | ( |
| 230 | cache, |
| 231 | block_table, |
| 232 | workspace_starts_t, |
| 233 | num_reqs, |
| 234 | total_tokens, |
| 235 | expected, |
| 236 | ) = _build_test_case(seq_lens, block_size) |
| 237 | |
| 238 | dst = torch.zeros( |
| 239 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 240 | ) |
| 241 | |
| 242 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 243 | cache, dst, block_table, workspace_starts_t, num_reqs |
| 244 | ) |
| 245 | |
| 246 | # NoPE: fp8 dequant has rounding error, so we allow small tolerance. |
| 247 | # The fp8 -> float -> bf16 path can differ by up to ~1 ULP of bf16. |
| 248 | torch.testing.assert_close( |
| 249 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 250 | ) |
| 251 | |
| 252 | # RoPE: pure bf16 copy, must be bit-exact |
| 253 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |
| 254 | |
| 255 | |
| 256 | def test_cp_gather_fp8_shuffled_blocks(): |
| 257 | """Test that the kernel correctly follows the block table when |
| 258 | physical blocks are non-contiguous and out of order. |
| 259 | |
| 260 | Here we allocate 4 physical blocks but map the request's 2 logical |
| 261 | blocks to physical blocks [3, 1] (reversed, with gaps). |
| 262 | """ |
| 263 | torch.manual_seed(123) |
| 264 | block_size = 4 |
| 265 | seq_lens = [8] # needs 2 blocks (tokens 0-3 in block 0, 4-7 in block 1) |
| 266 | total_tokens = 8 |
| 267 | |
| 268 | # 4 physical blocks, but only blocks 3 and 1 are used (in that order). |
| 269 | # Tokens 0-3 -> physical block 3, tokens 4-7 -> physical block 1. |
| 270 | num_phys_blocks = 4 |
| 271 | cache = torch.zeros( |
| 272 | num_phys_blocks, block_size, ENTRY_BYTES, dtype=torch.uint8, device="cuda" |
| 273 | ) |
| 274 | block_table = torch.tensor([[3, 1]], dtype=torch.int32, device="cuda") |
| 275 | workspace_starts = torch.tensor([0], dtype=torch.int32, device="cuda") |
| 276 | expected = torch.zeros( |
| 277 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 278 | ) |
| 279 | |
| 280 | # Fill cache at the shuffled physical locations |
| 281 | for t in range(total_tokens): |
| 282 | # Follow the same block_table lookup the kernel will use |
| 283 | phys = block_table[0, t // block_size].item() |
| 284 | off = t % block_size |
| 285 | |
| 286 | for tile in range(NUM_TILES): |
| 287 | start = tile * GROUP_SIZE |
| 288 | fp8_vals = torch.randn(GROUP_SIZE, device="cuda").to(torch.float8_e4m3fn) |
| 289 | cache[phys, off, start : start + GROUP_SIZE] = fp8_vals.view(torch.uint8) |
| 290 | |
| 291 | # Use a fixed scale to keep this test simple |
| 292 | scale = 1.5 |
| 293 | scale_t = torch.tensor([scale], dtype=torch.float32, device="cuda") |
| 294 | cache[phys, off, NOPE_DIM + tile * 4 : NOPE_DIM + (tile + 1) * 4] = ( |
| 295 | scale_t.view(torch.uint8) |
| 296 | ) |
| 297 | |
| 298 | expected[t, start : start + GROUP_SIZE] = ( |
| 299 | fp8_vals.float() * scale |
| 300 | ).bfloat16() |
| 301 | |
| 302 | rope = torch.randn(ROPE_DIM, dtype=torch.bfloat16, device="cuda") |
| 303 | cache[phys, off, NOPE_DIM + 16 :] = rope.view(torch.uint8) |
| 304 | expected[t, NOPE_DIM:] = rope |
| 305 | |
| 306 | dst = torch.zeros( |
| 307 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 308 | ) |
| 309 | |
| 310 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 311 | cache, dst, block_table, workspace_starts, len(seq_lens) |
| 312 | ) |
| 313 | |
| 314 | torch.testing.assert_close( |
| 315 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 316 | ) |
| 317 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |
| 318 | |
| 319 | |
| 320 | def test_cp_gather_fp8_mixes_resident_and_pinned_host_rows(): |
| 321 | seq_lens = [7, 5] |
| 322 | block_size = 4 |
| 323 | ( |
| 324 | cache, |
| 325 | block_table, |
| 326 | workspace_starts, |
| 327 | num_reqs, |
| 328 | total_tokens, |
| 329 | expected, |
| 330 | ) = _build_test_case(seq_lens, block_size) |
| 331 | |
| 332 | device_cache = cache.clone() |
| 333 | host_cache = cache.cpu().pin_memory().view(-1, ENTRY_BYTES) |
| 334 | device_rows = torch.arange(host_cache.shape[0], dtype=torch.int32, device="cuda") |
| 335 | use_device = device_rows % 2 == 0 |
| 336 | host_rows = device_rows.clone() |
| 337 | device_rows.masked_fill_(~use_device, -1) |
| 338 | |
| 339 | host_cache[use_device.cpu()] = 0 |
| 340 | device_cache.view(-1, ENTRY_BYTES)[~use_device] = 0 |
| 341 | dst = torch.empty( |
| 342 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 343 | ) |
| 344 | |
| 345 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 346 | device_cache, |
| 347 | dst, |
| 348 | block_table, |
| 349 | workspace_starts, |
| 350 | num_reqs, |
| 351 | host_cache=host_cache, |
| 352 | host_row_ids=host_rows, |
| 353 | device_row_ids=device_rows, |
| 354 | ) |
| 355 | |
| 356 | torch.testing.assert_close( |
| 357 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 358 | ) |
| 359 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |
| 360 | |
| 361 | |
| 362 | @pytest.mark.parametrize( |
| 363 | "gather_seq_lens,seq_starts", |
| 364 | [ |
| 365 | ([6, 5, 3], [3, 4, 5]), |
| 366 | ([0, 5, 3], [12, 4, 5]), |
| 367 | ], |
| 368 | ) |
| 369 | def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts): |
| 370 | """Gather request slices beginning at arbitrary cache positions.""" |
| 371 | full_seq_lens = [12, 11, 9] |
| 372 | ( |
| 373 | cache, |
| 374 | block_table, |
| 375 | _workspace_starts_t, |
| 376 | num_reqs, |
| 377 | _total_tokens, |
| 378 | full_expected, |
| 379 | ) = _build_test_case(full_seq_lens, block_size=4) |
| 380 | |
| 381 | workspace_starts = torch.tensor( |
| 382 | [0, gather_seq_lens[0], sum(gather_seq_lens[:2])], |
| 383 | dtype=torch.int32, |
| 384 | device="cuda", |
| 385 | ) |
| 386 | seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda") |
| 387 | dst = torch.empty( |
| 388 | sum(gather_seq_lens), |
| 389 | NOPE_DIM + ROPE_DIM, |
| 390 | dtype=torch.bfloat16, |
| 391 | device="cuda", |
| 392 | ) |
| 393 | |
| 394 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 395 | cache, |
| 396 | dst, |
| 397 | block_table, |
| 398 | workspace_starts, |
| 399 | num_reqs, |
| 400 | seq_starts_t, |
| 401 | ) |
| 402 | |
| 403 | full_workspace_starts = [0, full_seq_lens[0], sum(full_seq_lens[:2])] |
| 404 | expected = torch.cat( |
| 405 | [ |
| 406 | full_expected[ |
| 407 | full_workspace_starts[i] + seq_starts[i] : full_workspace_starts[i] |
| 408 | + seq_starts[i] |
| 409 | + gather_seq_lens[i] |
| 410 | ] |
| 411 | for i in range(num_reqs) |
| 412 | ] |
| 413 | ) |
| 414 | torch.testing.assert_close( |
| 415 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 416 | ) |
| 417 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |
| 418 | |
| 419 | |
| 420 | @pytest.mark.parametrize( |
| 421 | "seq_lens,block_size", |
| 422 | [ |
| 423 | # Large sequence lengths matching end-to-end benchmark scenarios. |
| 424 | # Uses vectorized builder since per-token Python loops would be too slow. |
| 425 | ([8000], 64), |
| 426 | ([16000], 64), |
| 427 | ([32000], 64), |
| 428 | ([64000], 64), |
| 429 | ([96000], 64), |
| 430 | ([128000], 64), |
| 431 | ], |
| 432 | ) |
| 433 | def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): |
| 434 | """Correctness test with large sequence lengths matching benchmark |
| 435 | scenarios (8K-128K prefill).""" |
| 436 | ( |
| 437 | cache, |
| 438 | block_table, |
| 439 | workspace_starts_t, |
| 440 | num_reqs, |
| 441 | total_tokens, |
| 442 | expected, |
| 443 | ) = _build_test_case_fast(seq_lens, block_size) |
| 444 | |
| 445 | dst = torch.zeros( |
| 446 | total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" |
| 447 | ) |
| 448 | |
| 449 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 450 | cache, dst, block_table, workspace_starts_t, num_reqs |
| 451 | ) |
| 452 | |
| 453 | torch.testing.assert_close( |
| 454 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 455 | ) |
| 456 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |
| 457 | |
| 458 | |
| 459 | def test_cp_gather_fp8_large_uneven_sequences_with_starts(): |
| 460 | """Gather uneven long request slices through the page-oriented path.""" |
| 461 | gather_seq_lens = [17, 32_768, 71] |
| 462 | seq_starts = [3, 17, 5] |
| 463 | full_seq_lens = [ |
| 464 | start + length for start, length in zip(seq_starts, gather_seq_lens) |
| 465 | ] |
| 466 | ( |
| 467 | cache, |
| 468 | block_table, |
| 469 | full_workspace_starts, |
| 470 | num_reqs, |
| 471 | _total_tokens, |
| 472 | full_expected, |
| 473 | ) = _build_test_case_fast(full_seq_lens, block_size=64) |
| 474 | |
| 475 | workspace_starts = torch.zeros(num_reqs, dtype=torch.int32, device="cuda") |
| 476 | workspace_starts[1:] = torch.tensor( |
| 477 | gather_seq_lens[:-1], dtype=torch.int32, device="cuda" |
| 478 | ).cumsum(dim=0) |
| 479 | seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda") |
| 480 | dst = torch.empty( |
| 481 | sum(gather_seq_lens), |
| 482 | NOPE_DIM + ROPE_DIM, |
| 483 | dtype=torch.bfloat16, |
| 484 | device="cuda", |
| 485 | ) |
| 486 | |
| 487 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 488 | cache, |
| 489 | dst, |
| 490 | block_table, |
| 491 | workspace_starts, |
| 492 | num_reqs, |
| 493 | seq_starts_t, |
| 494 | ) |
| 495 | |
| 496 | expected = torch.cat( |
| 497 | [ |
| 498 | full_expected[ |
| 499 | full_workspace_starts[req_id] |
| 500 | + seq_starts[req_id] : full_workspace_starts[req_id] |
| 501 | + seq_starts[req_id] |
| 502 | + gather_seq_lens[req_id] |
| 503 | ] |
| 504 | for req_id in range(num_reqs) |
| 505 | ] |
| 506 | ) |
| 507 | torch.testing.assert_close( |
| 508 | dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 |
| 509 | ) |
| 510 | assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) |