fix image cache overhead

[shihaobai]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @shihaobai, 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 aims to reduce the overhead associated with caching and retrieving image embeddings, likely improving performance or memory usage. The core change involves optimizing the serialization and deserialization process for these embeddings by moving away from PyTorch's default methods towards more direct buffer manipulation with NumPy, while also adding data type awareness.

Highlights

• Image Embedding Serialization: Refactored the tensor2bytes and bytes2tensor utility functions used for caching image embeddings to switch from torch.save/torch.load to using NumPy buffer operations (tobytes, frombuffer).
• Data Type Handling: Introduced explicit handling for different tensor data types (float32 vs. others, potentially bfloat16 via uint16 conversion) during serialization and deserialization to potentially reduce memory overhead.
• Model Integration: Updated the image embedding loading logic in the Gemma3 and Qwen-VL pre-layer inference modules to pass the required dtype to the modified bytes2tensor function.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configureGemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The tensor serialization method in lightllm/server/embed_cache/utils.py for float16 and bfloat16 types will lead to data corruption due to incorrect value casting. Use bit reinterpretation instead. The changes in the pre_layer_infer.py files are appropriate but depend on the correctness of these utility functions.

[gemini-code-assist]

For non-float32 types, the tensor2bytes function uses t.cpu().to(torch.uint16), which performs a value cast instead of a bit-wise reinterpretation, leading to data loss. Use t.cpu().contiguous().view(torch.uint16) to preserve the bit representation. Explicitly check for torch.float16 and torch.bfloat16.

Suggested change

	else:
	t = t.cpu().to(torch.uint16).numpy().tobytes()
	elif t.dtype == torch.float16 or t.dtype == torch.bfloat16:
	t_view = t.cpu().contiguous().view(torch.uint16)
	t = t_view.numpy().tobytes()
	else:
	raise TypeError(f"Unsupported dtype for tensor2bytes: {t.dtype}. Only float32, float16, bfloat16 are explicitly supported.")

[gemini-code-assist]

The bytes2tensor function incorrectly casts uint16 values back to float, instead of reinterpreting the bits. Use .view(torch_dtype) on the tensor created from arr_loaded to reinterpret the bits correctly. Use .copy() on the NumPy array returned by np.frombuffer for safety. Explicitly handle float32, float16/bfloat16.

    if torch_dtype == torch.float32:
        arr_loaded = np.frombuffer(b, dtype=np.float32)
        return torch.from_numpy(arr_loaded)
    elif torch_dtype == torch.float16 or torch_dtype == torch.bfloat16:
        arr_loaded_uint16 = np.frombuffer(b, dtype=np.uint16)
        return torch.from_numpy(arr_loaded_uint16.copy()).view(torch_dtype)
    else:
        raise TypeError(f"Unsupported torch_dtype for bytes2tensor: {torch_dtype}. This function is optimized for float32, float16, bfloat16.")