| 344 | |
| 345 | |
| 346 | def test_cache(): |
| 347 | @I.ir_module |
| 348 | class Module: |
| 349 | @R.function |
| 350 | def main( |
| 351 | key: R.Tensor(("num_tokens", 1, 8), dtype="float16"), |
| 352 | value: R.Tensor(("num_tokens", 1, 8), dtype="float16"), |
| 353 | key_cache: R.Tensor(("num_blocks", 1, 1, 16, 8), dtype="float16"), |
| 354 | value_cache: R.Tensor(("num_blocks", 1, 8, 16), dtype="float16"), |
| 355 | slot_mapping: R.Tensor(("num_tokens",), dtype="int32"), |
| 356 | ) -> R.Tuple( |
| 357 | [ |
| 358 | R.Tensor(("num_blocks", 1, 8, 16, 8), dtype="float16"), |
| 359 | R.Tensor(("num_blocks", 1, 8, 16), dtype="float16"), |
| 360 | ] |
| 361 | ): |
| 362 | with R.dataflow(): |
| 363 | kv = R.call_pure_packed( |
| 364 | "tvm.contrib.vllm.reshape_and_cache", |
| 365 | key, |
| 366 | value, |
| 367 | key_cache, |
| 368 | value_cache, |
| 369 | slot_mapping, |
| 370 | sinfo_args=[key_cache.struct_info, value_cache.struct_info], |
| 371 | ) |
| 372 | out = (kv[0], kv[1]) |
| 373 | R.output(out) |
| 374 | return out |
| 375 | |
| 376 | np.random.seed(0) |
| 377 | num_heads = 1 |
| 378 | head_dim = 8 |
| 379 | vec_size = 8 |
| 380 | block_size = 16 |
| 381 | num_tokens = 8 |
| 382 | num_blocks = 1 |
| 383 | key = np.random.randn(num_tokens, num_heads, head_dim).astype("float16") |
| 384 | value = np.random.randn(num_tokens, num_heads, head_dim).astype("float16") |
| 385 | key_cache_before = np.random.randn( |
| 386 | num_blocks, num_heads, head_dim // vec_size, block_size, vec_size |
| 387 | ).astype("float16") |
| 388 | value_cache_before = np.random.randn(num_blocks, num_heads, head_dim, block_size).astype( |
| 389 | "float16" |
| 390 | ) |
| 391 | slot_mapping = np.arange(num_tokens).astype("int32") |
| 392 | |
| 393 | key_cache = key_cache_before.copy() |
| 394 | value_cache = value_cache_before.copy() |
| 395 | |
| 396 | out_key_cache, out_value_cache = build_and_run( |
| 397 | Module, |
| 398 | [key, value, key_cache, value_cache, slot_mapping], |
| 399 | "cuda", |
| 400 | ) |
| 401 | |
| 402 | ref_key_cache = np.array( |
| 403 | [ |