Summary of Changes

Hello @hlin99, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the flexibility and portability of the LMCache library by decoupling it from mandatory compiled CUDA extensions. It achieves this by providing robust pure-Python fallback implementations for all core operations and introducing a dynamic backend selection mechanism. This change allows LMCache to function seamlessly on systems without NVIDIA GPUs, broadening its compatibility and ease of use across diverse hardware environments.

Highlights

• Centralized Backend Selection: Implemented a predicate-based registry in lmcache/__init__.py to dynamically select the best available backend (CUDA > fallback) at import time, ensuring lmcache can be installed and imported without NVIDIA GPUs.
• Pure-Python Fallback Implementations: Introduced lmcache/non_cuda_equivalents.py with a complete set of pure-Python (+ NumPy/SciPy) fallback implementations for all operations previously exclusive to c_ops, including rotary embedding, CDF calculation, KV-cache reshape/transfer, encode/decode, and pinned/NUMA memory management.
• Numerical Equivalence Testing: Added tests/v1/test_non_cuda_equivalents.py to rigorously cross-compare results of each operation under CUDA c_ops, non-CUDA with GPU visible, and non-CUDA without GPU, ensuring numerical equivalence across backends.
• Simplified Import Logic: Removed CUDA-only import guards (if torch.cuda.is_available():) across the codebase, as c_ops import now always succeeds due to automatic fallback, simplifying module dependencies.
• Updated Test Skip Logic: Adapted test skip logic in tests/v1/test_mem_kernels.py to use torch.cuda.is_available() directly, reflecting the new backend selection mechanism.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

hi @hickeyma @sammshen . i see you're the creator of non cuda equivalents. can you take some time to review this PR and let me know your thoughts? many thanks for your time.

@hlin99 What about installing LMCache from source? AFIK, this requires CUDA?

hi @hickeyma thanks for your review. installing from source I did not encounter the issue, and I think that might be the reason that people are not complaining crash in import c ops. However, just viewing from logic perspective, cuda.is_available != c lib must be there to be safely imported. There might be other reasons to fail the import. So fallback path with a try & exception block is a safe way to properly handle this

resolve conflicts and rebase to latest.

@hlin99 can you see you unit test failures on AMD?

@pytest.fixture(scope="module")
--
def run_all_children():
"""Launch 3 child processes. Runs once for the entire module."""
if RESULTS_DIR.exists():
shutil.rmtree(RESULTS_DIR)
 
for mode, cuda_vis in [("CUDA_OPS", "0"), ("NON_CUDA", "0"), ("NON_CUDA", "")]:
r = run_scenario(mode, cuda_vis)
>           assert r.returncode == 0, (
f"Scenario {mode}/CUDA_VISIBLE_DEVICES='{cuda_vis}' failed:\n"
f"{r.stdout}\n{r.stderr}"
)
E           AssertionError: Scenario CUDA_OPS/CUDA_VISIBLE_DEVICES='0' failed:
E
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[transfer_direction_enum] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[multi_layer_kv_transfer] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[multi_layer_kv_transfer_unilateral] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[single_layer_kv_transfer] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[single_layer_kv_transfer_sgl] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[load_and_reshape_flash] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[reshape_and_cache_back_flash] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[lmcache_memcpy_async] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[encode_fast_new] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[decode_fast_new] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[decode_fast_prefsum] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[calculate_cdf] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[rotary_embedding_k_fused] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_pinned_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_pinned_numa_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_numa_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[get_gpu_pci_bus_id] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E
E             =================================== FAILURES ===================================
E             ____________________ test_scenario[multi_layer_kv_transfer] ____________________
E
E             name = 'multi_layer_kv_transfer'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_multi_layer_kv_transfer():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_layers = 2
E                     num_tokens = 4
E                     head_size = 16
E                     page_buffer_size = 10
E                     dtype = torch.float32
E
E                     slot_mapping = torch.tensor(
E                         [0, 2, 5, 9],
E                         dtype=torch.int64,
E                         device=device,
E                     )
E
E                     for direction in [True, False]:
E                         dir_tag = "paged2lmc" if direction else "lmc2paged"
E
E                         # 1. LMCache Tensor
E                         lmc_shape = (2, num_layers, num_tokens, head_size)
E                         key_value = torch.zeros(
E                             lmc_shape,
E                             dtype=dtype,
E                             device=device,
E                         )
E                         if not direction:  # LMC → Paged
E                             for ly in range(num_layers):
E                                 for t in range(num_tokens):
E                                     val = (
E                                         ly * 1000 + t * 10 + torch.arange(head_size, device=device)
E                                     ).to(dtype)
E                                     key_value[0, ly, t] = val
E                                     key_value[1, ly, t] = val + 500
E
E                         # 2. Paged Buffers
E                         page_buffers = []
E                         for ly in range(num_layers):
E                             pb = torch.zeros(
E                                 (2, page_buffer_size, head_size),
E                                 dtype=dtype,
E                                 device=device,
E                             )
E                             if direction:  # Paged → LMC
E                                 for s in range(page_buffer_size):
E                                     val = (
E                                         ly * 2000
E                                         + s * 10
E                                         + torch.arange(
E                                             head_size,
E                                             device=device,
E                                         )
E                                     ).to(dtype)
E                                     pb[0, s] = val
E                                     pb[1, s] = val + 700
E                             page_buffers.append(pb)
E
E                         # 3. Pointer Tensor
E                         key_value_ptrs = torch.tensor(
E                             [pb.data_ptr() for pb in page_buffers],
E                             dtype=torch.int64,
E                             device=device,
E                         )
E
E                         # 4. Execute
E             >           ops.multi_layer_kv_transfer(
E                             key_value,
E                             key_value_ptrs,
E                             slot_mapping,
E                             torch.device(device),
E                             page_buffer_size,
E                             direction,
E                             False,  # use_mla
E                         )
E             E           TypeError: multi_layer_kv_transfer(): incompatible function arguments. The following argument types are supported:
E             E               1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.device, arg4: typing.SupportsInt, arg5: lmcache.c_ops.TransferDirection, arg6: lmcache.c_ops.GPUKVFormat, arg7: typing.SupportsInt) -> None
E             E
E             E           Invoked with: tensor([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]],
E             E
E             E
E             E                   [[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]],
E             E                  device='cuda:0'), tensor([140133642798592, 140133642801152], device='cuda:0'), tensor([0, 2, 5, 9], device='cuda:0'), device(type='cuda', index=0), 10, True, False
E
E             tests/v1/test_non_cuda_equivalents.py:1093: TypeError
E             ______________ test_scenario[multi_layer_kv_transfer_unilateral] _______________
E
E             name = 'multi_layer_kv_transfer_unilateral'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_multi_layer_kv_transfer_unilateral():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_layers = 2
E                     num_tokens = 4
E                     head_size = 16
E                     page_buffer_size = 10
E                     dtype = torch.float32
E
E                     slot_mapping = torch.tensor(
E                         [1, 3, 4, 7],
E                         dtype=torch.int64,
E                         device=device,
E                     )
E
E                     for direction in [True, False]:
E                         dir_tag = "p2l" if direction else "l2p"
E
E                         # LMC Layout: [2, num_layers, num_tokens, head_size]
E                         lmc_shape = (2, num_layers, num_tokens, head_size)
E                         lmc_tensor = torch.zeros(
E                             lmc_shape,
E                             dtype=dtype,
E                             device=device,
E                         )
E
E                         if not direction:  # LMC → Paged
E                             for kv in range(2):
E                                 for ly in range(num_layers):
E                                     for t in range(num_tokens):
E                                         val = (
E                                             kv * 5000
E                                             + ly * 1000
E                                             + t * 10
E                                             + torch.arange(
E                                                 head_size,
E                                                 device=device,
E                                             )
E                                         ).to(dtype)
E                                         lmc_tensor[kv, ly, t] = val
E
E                         # 1. Paged Buffers
E                         buffers = {}
E                         for kv in range(2):
E                             for ly in range(num_layers):
E                                 pb = torch.zeros(
E                                     (page_buffer_size, head_size),
E                                     dtype=dtype,
E                                     device=device,
E                                 )
E                                 if direction:  # Paged → LMC
E                                     val = (
E                                         kv * 7000
E                                         + ly * 2000
E                                         + torch.arange(
E                                             head_size,
E                                             device=device,
E                                         )
E                                     ).to(dtype)
E                                     for s in range(page_buffer_size):
E                                         pb[s] = val + (s * 10)
E                                 buffers[(kv, ly)] = pb
E
E                         # 2. Grouped Pointer Tensor
E                         # C++: ptrs[layer_id] = Key,
E                         #      ptrs[layer_id + num_layers] = Value
E                         ptr_list = []
E                         for ly in range(num_layers):
E                             ptr_list.append(buffers[(0, ly)].data_ptr())
E                         for ly in range(num_layers):
E                             ptr_list.append(buffers[(1, ly)].data_ptr())
E
E                         key_value_ptrs = torch.tensor(
E                             ptr_list,
E                             dtype=torch.int64,
E                             device=device,
E                         ).contiguous()
E
E                         # 3. Execute
E             >           ops.multi_layer_kv_transfer_unilateral(
E                             lmc_tensor,
E                             key_value_ptrs,
E                             slot_mapping,
E                             torch.device(device),
E                             page_buffer_size,
E                             direction,
E                             False,  # use_mla
E                         )
E             E           TypeError: multi_layer_kv_transfer_unilateral(): incompatible function arguments. The following argument types are supported:
E             E               1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.device, arg4: typing.SupportsInt, arg5: lmcache.c_ops.TransferDirection, arg6: lmcache.c_ops.GPUKVFormat) -> None
E             E
E             E           Invoked with: tensor([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]],
E             E
E             E
E             E                   [[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]],
E             E                  device='cuda:0'), tensor([140133642804224, 140133642805760, 140133642807296, 140133642808320],
E             E                  device='cuda:0'), tensor([1, 3, 4, 7], device='cuda:0'), device(type='cuda', index=0), 10, True, False
E
E             tests/v1/test_non_cuda_equivalents.py:1213: TypeError
E             ___________________ test_scenario[single_layer_kv_transfer] ____________________
E
E             name = 'single_layer_kv_transfer'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_single_layer_kv_transfer():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_tokens = 64
E                     num_blocks = 256
E                     block_size = 16
E                     num_heads = 12
E                     head_size = 64
E                     hidden_size = num_heads * head_size
E
E                     slot_mapping = torch.arange(
E                         0,
E                         num_tokens * 2,
E                         2,
E                         device=device,
E                     ).to(torch.int64)
E
E                     # (use_mla, token_major, vllm_two_major, direction)
E                     # direction: False = LMC→vLLM, True = vLLM→LMC
E                     test_cases = [
E                         (False, True, True, False),
E                         (False, False, False, False),
E                         (False, True, True, True),
E                         (True, True, True, False),
E                         (True, True, True, True),
E                     ]
E
E                     for use_mla, token_major, vllm_two_major, direction in test_cases:
E                         dir_tag = "v2l" if direction else "l2v"
E                         case_desc = (
E                             f"MLA={use_mla}, TM={token_major}, 2Maj={vllm_two_major}, Dir={dir_tag}"
E                         )
E
E                         # 1. Setup Shapes
E                         if use_mla:
E                             lmc_shape = (num_tokens, hidden_size)
E                             vllm_shape = (num_blocks, block_size, hidden_size)
E                         else:
E                             lmc_shape = (
E                                 (num_tokens, 2, hidden_size)
E                                 if token_major
E                                 else (2, num_tokens, hidden_size)
E                             )
E                             if vllm_two_major:
E                                 vllm_shape = (
E                                     2,
E                                     num_blocks,
E                                     block_size,
E                                     num_heads,
E                                     head_size,
E                                 )
E                             else:
E                                 vllm_shape = (
E                                     num_blocks,
E                                     2,
E                                     block_size,
E                                     num_heads,
E                                     head_size,
E                                 )
E
E                         # 2. Deterministic Data
E                         lmc_size = 1
E                         for s in lmc_shape:
E                             lmc_size *= s
E                         vllm_size = 1
E                         for s in vllm_shape:
E                             vllm_size *= s
E
E                         lmc_tensor = (
E                             (torch.arange(lmc_size, device=device) % 1000)
E                             .to(torch.float16)
E                             .reshape(lmc_shape)
E                         )
E                         vllm_tensor = (
E                             (torch.arange(vllm_size, device=device) % 1000)
E                             .to(torch.float16)
E                             .reshape(vllm_shape)
E                         )
E
E                         # 3. Golden Reference
E                         lmc_ref = lmc_tensor.clone()
E                         vllm_ref = vllm_tensor.clone()
E                         block_indices = slot_mapping // block_size
E                         block_offsets = slot_mapping % block_size
E
E                         if not direction:  # LMC → vLLM
E                             if use_mla:
E                                 vllm_ref[block_indices, block_offsets, :] = lmc_ref
E                             else:
E                                 src = lmc_ref if token_major else lmc_ref.permute(1, 0, 2)
E                                 src = src.view(
E                                     num_tokens,
E                                     2,
E                                     num_heads,
E                                     head_size,
E                                 )
E                                 if vllm_two_major:
E                                     vllm_ref[0, block_indices, block_offsets] = src[:, 0, :, :]
E                                     vllm_ref[1, block_indices, block_offsets] = src[:, 1, :, :]
E                                 else:
E                                     vllm_ref[block_indices, 0, block_offsets] = src[:, 0, :, :]
E                                     vllm_ref[block_indices, 1, block_offsets] = src[:, 1, :, :]
E                         else:  # vLLM → LMC
E                             if use_mla:
E                                 lmc_ref = vllm_ref[block_indices, block_offsets, :]
E                             else:
E                                 if vllm_two_major:
E                                     k = vllm_ref[0, block_indices, block_offsets]
E                                     v = vllm_ref[1, block_indices, block_offsets]
E                                 else:
E                                     k = vllm_ref[block_indices, 0, block_offsets]
E                                     v = vllm_ref[block_indices, 1, block_offsets]
E                                 combined = torch.stack(
E                                     [k, v],
E                                     dim=1,
E                                 ).view(num_tokens, 2, hidden_size)
E                                 lmc_ref = combined if token_major else combined.permute(1, 0, 2)
E
E                         # 4. Execute
E             >           ops.single_layer_kv_transfer(
E                             lmc_tensor,
E                             vllm_tensor,
E                             slot_mapping,
E                             direction,
E                             token_major,
E                             vllm_two_major,
E                             use_mla,

Let me know if you need access to the Buildkite portal

@hlin99 can you see you unit test failures on AMD?

@pytest.fixture(scope="module")
--
def run_all_children():
"""Launch 3 child processes. Runs once for the entire module."""
if RESULTS_DIR.exists():
shutil.rmtree(RESULTS_DIR)
 
for mode, cuda_vis in [("CUDA_OPS", "0"), ("NON_CUDA", "0"), ("NON_CUDA", "")]:
r = run_scenario(mode, cuda_vis)
>           assert r.returncode == 0, (
f"Scenario {mode}/CUDA_VISIBLE_DEVICES='{cuda_vis}' failed:\n"
f"{r.stdout}\n{r.stderr}"
)
E           AssertionError: Scenario CUDA_OPS/CUDA_VISIBLE_DEVICES='0' failed:
E
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[transfer_direction_enum] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[multi_layer_kv_transfer] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[multi_layer_kv_transfer_unilateral] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[single_layer_kv_transfer] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[single_layer_kv_transfer_sgl] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             FAILED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[load_and_reshape_flash] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[reshape_and_cache_back_flash] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[lmcache_memcpy_async] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[encode_fast_new] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[decode_fast_new] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[decode_fast_prefsum] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[calculate_cdf] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[rotary_embedding_k_fused] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_pinned_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_pinned_numa_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[alloc_free_numa_ptr] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E             tests/v1/test_non_cuda_equivalents.py::test_scenario[get_gpu_pci_bus_id] >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             >>> Importing lmcache.c_ops as ops (Mode: CUDA_OPS)
E             PASSED
E
E             =================================== FAILURES ===================================
E             ____________________ test_scenario[multi_layer_kv_transfer] ____________________
E
E             name = 'multi_layer_kv_transfer'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_multi_layer_kv_transfer():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_layers = 2
E                     num_tokens = 4
E                     head_size = 16
E                     page_buffer_size = 10
E                     dtype = torch.float32
E
E                     slot_mapping = torch.tensor(
E                         [0, 2, 5, 9],
E                         dtype=torch.int64,
E                         device=device,
E                     )
E
E                     for direction in [True, False]:
E                         dir_tag = "paged2lmc" if direction else "lmc2paged"
E
E                         # 1. LMCache Tensor
E                         lmc_shape = (2, num_layers, num_tokens, head_size)
E                         key_value = torch.zeros(
E                             lmc_shape,
E                             dtype=dtype,
E                             device=device,
E                         )
E                         if not direction:  # LMC → Paged
E                             for ly in range(num_layers):
E                                 for t in range(num_tokens):
E                                     val = (
E                                         ly * 1000 + t * 10 + torch.arange(head_size, device=device)
E                                     ).to(dtype)
E                                     key_value[0, ly, t] = val
E                                     key_value[1, ly, t] = val + 500
E
E                         # 2. Paged Buffers
E                         page_buffers = []
E                         for ly in range(num_layers):
E                             pb = torch.zeros(
E                                 (2, page_buffer_size, head_size),
E                                 dtype=dtype,
E                                 device=device,
E                             )
E                             if direction:  # Paged → LMC
E                                 for s in range(page_buffer_size):
E                                     val = (
E                                         ly * 2000
E                                         + s * 10
E                                         + torch.arange(
E                                             head_size,
E                                             device=device,
E                                         )
E                                     ).to(dtype)
E                                     pb[0, s] = val
E                                     pb[1, s] = val + 700
E                             page_buffers.append(pb)
E
E                         # 3. Pointer Tensor
E                         key_value_ptrs = torch.tensor(
E                             [pb.data_ptr() for pb in page_buffers],
E                             dtype=torch.int64,
E                             device=device,
E                         )
E
E                         # 4. Execute
E             >           ops.multi_layer_kv_transfer(
E                             key_value,
E                             key_value_ptrs,
E                             slot_mapping,
E                             torch.device(device),
E                             page_buffer_size,
E                             direction,
E                             False,  # use_mla
E                         )
E             E           TypeError: multi_layer_kv_transfer(): incompatible function arguments. The following argument types are supported:
E             E               1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.device, arg4: typing.SupportsInt, arg5: lmcache.c_ops.TransferDirection, arg6: lmcache.c_ops.GPUKVFormat, arg7: typing.SupportsInt) -> None
E             E
E             E           Invoked with: tensor([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]],
E             E
E             E
E             E                   [[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]],
E             E                  device='cuda:0'), tensor([140133642798592, 140133642801152], device='cuda:0'), tensor([0, 2, 5, 9], device='cuda:0'), device(type='cuda', index=0), 10, True, False
E
E             tests/v1/test_non_cuda_equivalents.py:1093: TypeError
E             ______________ test_scenario[multi_layer_kv_transfer_unilateral] _______________
E
E             name = 'multi_layer_kv_transfer_unilateral'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_multi_layer_kv_transfer_unilateral():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_layers = 2
E                     num_tokens = 4
E                     head_size = 16
E                     page_buffer_size = 10
E                     dtype = torch.float32
E
E                     slot_mapping = torch.tensor(
E                         [1, 3, 4, 7],
E                         dtype=torch.int64,
E                         device=device,
E                     )
E
E                     for direction in [True, False]:
E                         dir_tag = "p2l" if direction else "l2p"
E
E                         # LMC Layout: [2, num_layers, num_tokens, head_size]
E                         lmc_shape = (2, num_layers, num_tokens, head_size)
E                         lmc_tensor = torch.zeros(
E                             lmc_shape,
E                             dtype=dtype,
E                             device=device,
E                         )
E
E                         if not direction:  # LMC → Paged
E                             for kv in range(2):
E                                 for ly in range(num_layers):
E                                     for t in range(num_tokens):
E                                         val = (
E                                             kv * 5000
E                                             + ly * 1000
E                                             + t * 10
E                                             + torch.arange(
E                                                 head_size,
E                                                 device=device,
E                                             )
E                                         ).to(dtype)
E                                         lmc_tensor[kv, ly, t] = val
E
E                         # 1. Paged Buffers
E                         buffers = {}
E                         for kv in range(2):
E                             for ly in range(num_layers):
E                                 pb = torch.zeros(
E                                     (page_buffer_size, head_size),
E                                     dtype=dtype,
E                                     device=device,
E                                 )
E                                 if direction:  # Paged → LMC
E                                     val = (
E                                         kv * 7000
E                                         + ly * 2000
E                                         + torch.arange(
E                                             head_size,
E                                             device=device,
E                                         )
E                                     ).to(dtype)
E                                     for s in range(page_buffer_size):
E                                         pb[s] = val + (s * 10)
E                                 buffers[(kv, ly)] = pb
E
E                         # 2. Grouped Pointer Tensor
E                         # C++: ptrs[layer_id] = Key,
E                         #      ptrs[layer_id + num_layers] = Value
E                         ptr_list = []
E                         for ly in range(num_layers):
E                             ptr_list.append(buffers[(0, ly)].data_ptr())
E                         for ly in range(num_layers):
E                             ptr_list.append(buffers[(1, ly)].data_ptr())
E
E                         key_value_ptrs = torch.tensor(
E                             ptr_list,
E                             dtype=torch.int64,
E                             device=device,
E                         ).contiguous()
E
E                         # 3. Execute
E             >           ops.multi_layer_kv_transfer_unilateral(
E                             lmc_tensor,
E                             key_value_ptrs,
E                             slot_mapping,
E                             torch.device(device),
E                             page_buffer_size,
E                             direction,
E                             False,  # use_mla
E                         )
E             E           TypeError: multi_layer_kv_transfer_unilateral(): incompatible function arguments. The following argument types are supported:
E             E               1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.device, arg4: typing.SupportsInt, arg5: lmcache.c_ops.TransferDirection, arg6: lmcache.c_ops.GPUKVFormat) -> None
E             E
E             E           Invoked with: tensor([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]],
E             E
E             E
E             E                   [[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
E             E
E             E                    [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
E             E                     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]],
E             E                  device='cuda:0'), tensor([140133642804224, 140133642805760, 140133642807296, 140133642808320],
E             E                  device='cuda:0'), tensor([1, 3, 4, 7], device='cuda:0'), device(type='cuda', index=0), 10, True, False
E
E             tests/v1/test_non_cuda_equivalents.py:1213: TypeError
E             ___________________ test_scenario[single_layer_kv_transfer] ____________________
E
E             name = 'single_layer_kv_transfer'
E
E                 @pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
E                 def test_scenario(name):
E             >       SCENARIO_REGISTRY[name]()
E
E             tests/v1/test_non_cuda_equivalents.py:1400:
E             _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
E
E                 def scenario_single_layer_kv_transfer():
E                     ops, scene_info = get_test_context()
E                     is_cuda_backend = scene_info.startswith("cuda_ops")
E
E                     torch.manual_seed(42)
E                     if torch.cuda.is_available():
E                         torch.cuda.manual_seed_all(42)
E
E                     device = f"cuda:{torch.cuda.current_device()}" if is_cuda_backend else "cpu"
E
E                     num_tokens = 64
E                     num_blocks = 256
E                     block_size = 16
E                     num_heads = 12
E                     head_size = 64
E                     hidden_size = num_heads * head_size
E
E                     slot_mapping = torch.arange(
E                         0,
E                         num_tokens * 2,
E                         2,
E                         device=device,
E                     ).to(torch.int64)
E
E                     # (use_mla, token_major, vllm_two_major, direction)
E                     # direction: False = LMC→vLLM, True = vLLM→LMC
E                     test_cases = [
E                         (False, True, True, False),
E                         (False, False, False, False),
E                         (False, True, True, True),
E                         (True, True, True, False),
E                         (True, True, True, True),
E                     ]
E
E                     for use_mla, token_major, vllm_two_major, direction in test_cases:
E                         dir_tag = "v2l" if direction else "l2v"
E                         case_desc = (
E                             f"MLA={use_mla}, TM={token_major}, 2Maj={vllm_two_major}, Dir={dir_tag}"
E                         )
E
E                         # 1. Setup Shapes
E                         if use_mla:
E                             lmc_shape = (num_tokens, hidden_size)
E                             vllm_shape = (num_blocks, block_size, hidden_size)
E                         else:
E                             lmc_shape = (
E                                 (num_tokens, 2, hidden_size)
E                                 if token_major
E                                 else (2, num_tokens, hidden_size)
E                             )
E                             if vllm_two_major:
E                                 vllm_shape = (
E                                     2,
E                                     num_blocks,
E                                     block_size,
E                                     num_heads,
E                                     head_size,
E                                 )
E                             else:
E                                 vllm_shape = (
E                                     num_blocks,
E                                     2,
E                                     block_size,
E                                     num_heads,
E                                     head_size,
E                                 )
E
E                         # 2. Deterministic Data
E                         lmc_size = 1
E                         for s in lmc_shape:
E                             lmc_size *= s
E                         vllm_size = 1
E                         for s in vllm_shape:
E                             vllm_size *= s
E
E                         lmc_tensor = (
E                             (torch.arange(lmc_size, device=device) % 1000)
E                             .to(torch.float16)
E                             .reshape(lmc_shape)
E                         )
E                         vllm_tensor = (
E                             (torch.arange(vllm_size, device=device) % 1000)
E                             .to(torch.float16)
E                             .reshape(vllm_shape)
E                         )
E
E                         # 3. Golden Reference
E                         lmc_ref = lmc_tensor.clone()
E                         vllm_ref = vllm_tensor.clone()
E                         block_indices = slot_mapping // block_size
E                         block_offsets = slot_mapping % block_size
E
E                         if not direction:  # LMC → vLLM
E                             if use_mla:
E                                 vllm_ref[block_indices, block_offsets, :] = lmc_ref
E                             else:
E                                 src = lmc_ref if token_major else lmc_ref.permute(1, 0, 2)
E                                 src = src.view(
E                                     num_tokens,
E                                     2,
E                                     num_heads,
E                                     head_size,
E                                 )
E                                 if vllm_two_major:
E                                     vllm_ref[0, block_indices, block_offsets] = src[:, 0, :, :]
E                                     vllm_ref[1, block_indices, block_offsets] = src[:, 1, :, :]
E                                 else:
E                                     vllm_ref[block_indices, 0, block_offsets] = src[:, 0, :, :]
E                                     vllm_ref[block_indices, 1, block_offsets] = src[:, 1, :, :]
E                         else:  # vLLM → LMC
E                             if use_mla:
E                                 lmc_ref = vllm_ref[block_indices, block_offsets, :]
E                             else:
E                                 if vllm_two_major:
E                                     k = vllm_ref[0, block_indices, block_offsets]
E                                     v = vllm_ref[1, block_indices, block_offsets]
E                                 else:
E                                     k = vllm_ref[block_indices, 0, block_offsets]
E                                     v = vllm_ref[block_indices, 1, block_offsets]
E                                 combined = torch.stack(
E                                     [k, v],
E                                     dim=1,
E                                 ).view(num_tokens, 2, hidden_size)
E                                 lmc_ref = combined if token_major else combined.permute(1, 0, 2)
E
E                         # 4. Execute
E             >           ops.single_layer_kv_transfer(
E                             lmc_tensor,
E                             vllm_tensor,
E                             slot_mapping,
E                             direction,
E                             token_major,
E                             vllm_two_major,
E                             use_mla,

Let me know if you need access to the Buildkite portal

hi @sammshen . thanks for pointing this out! not aware recently there was an API update in these kernels. i'm working on an upgrade now.

thanks @chunxiaozheng for the detailed review

thanks @chunxiaozheng for the detailed review

thank you @sammshen & @chunxiaozheng

the CODEOWNERS is a little too fine-grained, let me change that

this PR should be able to merge after: #3016

@maobaolong @chunxiaozheng PTAL!

@hlin99 Maybe you forget to update the import lmcache.c_ops as lmc_ops in lmcache/v1/multiprocess/gpu_context.py?

@DongDongJu Would you like to take a look at this PR?

@hlin99 Maybe you forget to update the import lmcache.c_ops as lmc_ops in lmcache/v1/multiprocess/gpu_context.py?

Hi @maobaolong . you're right! i was aware of it. since this PR was opened, there're some changes on c_ops, i only updated necessary part(signature, enum def) in this PR, because i really don't want this PR to go bigger and bigger.

I will raise separate PR to fill the known gaps bettween c ops and non cuda equivalents soon. or anybody wants to contribute is also welcome