Add XPU support to LMCache for CPU offloading

[zhenwei-intel]

TODO:

• add test
• mla check
• rebase dev
• disk offload
• optimize from_gpu and to_gpu

offload to cpu

offload to local disk

[zhenwei-intel]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Intel XPUs, which is a significant enhancement for platform compatibility. The changes primarily involve adding a Platform utility class for hardware detection and making various parts of the codebase platform-agnostic by replacing CUDA-specific calls with conditional logic. A new VLLMPagedMemXPUConnectorV2 is also added.

The overall approach is solid, but I've found a few critical issues where the optional import of lmcache.c_ops can lead to runtime errors on non-CUDA platforms, as the code doesn't handle the case where the import fails. I've also included some suggestions to improve code maintainability by reducing duplication and cleaning up comments.

[gemini-code-assist]

The docstring for to_gpu mentions "CUDA kernel" on line 121. Since this connector is for XPU, this is misleading. Please update the comment to be platform-agnostic. For example, you could replace "CUDA kernel" with "computation kernel" or "underlying kernel". The same issue exists in the docstring for from_gpu on line 177.

[zhenwei-intel]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Intel XPUs by adding a Platform utility to abstract device-specific logic. It includes a new VLLMPagedMemXPUConnectorV2 for XPU integration and updates various parts of the codebase to be device-agnostic. The changes are well-structured. However, there are critical issues where the optional import of CUDA-specific C operations (lmc_ops) is not handled correctly, which will lead to crashes on non-CUDA platforms like XPU when certain features (e.g., cachegen serde) are used. There are also some minor code cleanup opportunities in the new XPU connector.

[zhenwei-intel]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds XPU support to LMCache, primarily by introducing a Platform utility for device detection and implementing a new VLLMPagedMemXPUConnectorV2 for XPU devices. The changes are well-structured to abstract away device-specific logic. However, I've found a critical issue with stream initialization for XPU that could lead to a runtime error. Additionally, there are opportunities to improve maintainability by reducing code duplication and to enhance clarity by correcting docstrings and error messages in the new XPU connector. Addressing these points will strengthen the XPU integration.

[zhenwei-intel]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Intel XPUs to LMCache, which is a valuable addition for enabling CPU offloading on different hardware. The changes primarily involve abstracting device-specific code (CUDA vs. XPU) and adding a new VLLMPagedMemXPUConnectorV2 for XPU devices. The overall approach is solid, and the device-agnostic changes like replacing .cuda() with .to(device) are well-executed. My review focuses on improving maintainability by addressing code duplication, fixing minor documentation errors, and highlighting potential performance and correctness issues in the new XPU connector.

[gemini-code-assist]

Using zip with strict=False can hide potential bugs if the input iterables (memory_objs, starts, ends) have mismatched lengths, as it will silently truncate to the shortest one. It is safer to use strict=True (available in Python 3.10+) to raise a ValueError in such cases. If compatibility with older Python versions is required, adding an explicit length check before the loop would be a robust alternative.

Suggested change

	for memory_obj, start, end in zip(memory_objs, starts, ends, strict=False):
	for memory_obj, start, end in zip(memory_objs, starts, ends, strict=True):

[gemini-code-assist]

The logic to detect the available device (cuda or xpu) is duplicated in this function, both here and at lines 536-541. This could be refactored into a helper function to improve maintainability and reduce code duplication. For example, a function could return the device module and name, which can then be used in both places.

[gemini-code-assist]

VLLMPagedMemXPUConnectorV2 inherits from VLLMPagedMemGPUConnectorV2, but it seems to override most of the core logic (__init__, to_gpu, from_gpu) and uses a different implementation strategy (pure PyTorch vs. custom CUDA kernels). The parent class also contains CUDA-specific code that makes direct inheritance for other devices problematic. This can lead to maintenance challenges.

A cleaner design might be for VLLMPagedMemXPUConnectorV2 to inherit directly from GPUConnectorInterface. Any logic that is truly common between the GPU and XPU connectors could be extracted into a shared base class or utility functions.

[gemini-code-assist]

The docstring mentions a CUDA kernel, which appears to be a copy-paste artifact from the GPU connector. To maintain accuracy and clarity, this should be updated to be device-agnostic, for example, by simply referring to it as the underlying kernel.

Suggested change

	underlying CUDA kernel will never see -1 in slot_mapping)
	should NEVER overlap with the prefix caching (which means the
	underlying kernel will never see -1 in slot_mapping)

[gemini-code-assist]

The explicit Python loop over layers to perform index_copy_ can introduce performance overhead, especially for models with a large number of layers. This is in contrast to the GPU connector which uses a single fused CUDA kernel for this operation. While this pure-PyTorch implementation is a great way to support XPU without custom kernels, it's worth noting this potential performance bottleneck. If performance becomes an issue, creating a batched version of this operation or a custom XPU kernel could be considered for optimization.

[gemini-code-assist]

Similar to a previous comment, the docstring mentions a CUDA kernel. This should be updated to be device-agnostic to avoid confusion.

Suggested change

	underlying CUDA kernel will never see -1 in slot_mapping)
	should NEVER overlap with the prefix caching (which means the
	underlying kernel will never see -1 in slot_mapping)

[gemini-code-assist]

Similar to the to_gpu method, the use of list comprehensions here to iterate over layers ([... for kvcache in kvcaches]) effectively creates a Python loop. This can be a performance bottleneck compared to a fused kernel, especially for models with many layers. It's a good idea to be aware of this potential performance impact.