Fp8 deepseek

[blueswhen]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @blueswhen, 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 introduces significant improvements to the KV cache memory management and FP8 quantization system. The core change centralizes the logic for managing request-to-token indexes within the memory allocation process, leading to a cleaner and more efficient design. Additionally, it enhances the flexibility of FP8 quantization scale handling to better support diverse model architectures and updates relevant calibration data.

Highlights

• Centralized KV Cache Index Management: I've refactored the KV cache memory management by moving the logic for initializing and copying request-to-token indexes (init_req_to_token_indexes and copy_kv_index_to_req) directly into the MemoryManager.alloc method. This centralizes the memory allocation and index management, simplifying the codebase in basemodel.py and other inference paths.
• Flexible FP8 Quantization Scales: The OfflineFp8QuantMemManager has been enhanced to support more flexible handling of FP8 quantization scales. It now correctly adjusts scales for attention head configurations where the model's head count is a multiple or divisor of the total head number, particularly when using enable_fa3.
• Updated FP8 Calibration Data: I've updated the FP8 calibration JSON files for Qwen3 models (test_kv_cache_calib_per_head_qwen3_235b.json and test_kv_cache_calib_per_tensor_qwen3_235b.json) to reflect changes in num_head and scales_shape, ensuring compatibility with new model configurations or improved quantization parameters.
• Streamlined Memory Allocation Calls: Calls to mem_manager.alloc across various model inference, CUDA graph warmup, and benchmarking scripts have been updated to pass necessary request and sequence length information directly. This aligns with the centralized memory management logic and ensures proper index handling during prefill and decode operations.

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 configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for FP8 quantization for DeepSeek models and includes a significant refactoring of the memory allocation logic. The refactoring centralizes token index mapping within the MemoryManager, which improves code structure. The changes to support flexible FP8 quantization configurations are also a good addition. I've found a critical bug in the MTP benchmark test where incorrect tensors are being used for memory allocation, and a high-severity issue in the MemoryManager where a missing assertion could lead to a runtime error. Please address these points to ensure correctness and robustness.

[gemini-code-assist]

There seems to be a bug here. You are passing b_req_idx and b_seq_len to alloc, but these tensors have a size of batch_size and are from the prefill stage. The allocation is for batch_size * (len(draft_models) + 1) tokens for the decode stage. You should be using nopad_b_seq_idx and nopad_b_seq_len which are correctly sized for this allocation. Using the wrong tensors will lead to incorrect indexing in the copy_kv_index_to_req kernel.

Suggested change

	mem_indexes = main_model.req_manager.mem_manager.alloc(
	batch_size * (len(draft_models) + 1), b_req_idx, b_seq_len, None, False
	).cuda()
	mem_indexes = main_model.req_manager.mem_manager.alloc(
	batch_size * (len(draft_models) + 1), nopad_b_seq_idx, nopad_b_seq_len, None, False
	).cuda()

[gemini-code-assist]

When is_prefill is True, b_ready_cache_len is passed to init_req_to_token_indexes. However, the alloc function signature allows b_ready_cache_len to be None, which would cause a runtime error inside init_req_to_token_indexes as it calls .cpu().numpy() on it.

Please add an assertion to ensure b_ready_cache_len is not None when is_prefill is True. For example:

if is_prefill:
    assert b_ready_cache_len is not None, "b_ready_cache_len must be provided for prefill"
    init_req_to_token_indexes(
        ...
    )