Remove prefill cpu sync op

[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 significantly optimizes the model's prefill performance by strategically decoupling CPU and GPU tensor operations. By introducing explicit CPU-side tensor management within the ModelInput structure and deferring CUDA transfers until necessary, it effectively reduces costly CPU-GPU synchronization overhead. The changes also include broader refactoring to centralize ModelInput usage, extend support for new DeepSeekV3 model types, and integrate new Triton kernel autotuning settings for NVIDIA H200 GPUs. These modifications are complemented by refined benchmarking practices to accurately assess the performance gains.

Highlights

• CPU-GPU Data Separation: Implemented explicit separation of CPU-bound and CUDA-bound tensors within the ModelInput object, using _cpu suffixes for CPU tensors to manage data transfer more efficiently.
• Optimized Prefill Synchronization: Refactored the to_cuda mechanism to perform tensor transfers from CPU to GPU only when explicitly required, eliminating implicit CPU synchronization during the prefill phase.
• Unified Model Input Handling: Standardized the init_some_extra_state and gen_prefill_params functions to accept the complete ModelInput object, streamlining data access and improving code clarity across various model implementations.
• Expanded DeepSeekV3 Support: Updated model configuration checks to properly recognize and support new DeepSeekV3 variants, including "deepseekv3_vanilla" and "deepseekv3_eagle".
• Improved Benchmarking Accuracy: Enhanced benchmark scripts by utilizing pinned CPU memory for faster data transfers and repositioning torch.cuda.synchronize() for more precise measurement of GPU computation times.
• New Triton Kernel Autotuning: Introduced a new autotuning configuration file for a Triton kernel, specifically targeting MoE alignment on NVIDIA H200 GPUs, to further boost performance.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 refactors the prefill and decode paths to remove CPU-GPU synchronization points, which should improve performance. The main strategy is to prepare input metadata on CPU tensors and perform non-blocking copies to the GPU within the forward pass. The changes are extensive and consistently applied across many model-specific files. I've found a critical issue where a missing .item() call will lead to a TypeError at runtime, and another issue in a test file where a hardcoded file path breaks reproducibility. After addressing these points, this PR will be a solid performance improvement.

[gemini-code-assist]

The .max() method on a PyTorch tensor returns a 0-dimensional tensor, not a Python scalar. Downstream code that uses max_q_seq_len and max_kv_seq_len (e.g., for creating new tensors with torch.empty) expects an integer and will fail with a TypeError. You should add .item() to convert the 0-dim tensors to scalars. Since these operations are on CPU tensors, calling .item() will not cause a device synchronization.

Suggested change

	max_q_seq_len = (b_seq_len_cpu - b_ready_cache_len_cpu).max()
	max_kv_seq_len = b_seq_len_cpu.max()
	max_q_seq_len = (b_seq_len_cpu - b_ready_cache_len_cpu).max().item()
	max_kv_seq_len = b_seq_len_cpu.max().item()

[gemini-code-assist]

The benchmark test data generation has been changed from creating random data to loading a hardcoded file test.npy. This will cause the benchmark to fail for anyone who does not have this specific file. This change should be reverted to ensure the benchmark is self-contained and reproducible.

Suggested change

	# test_data = np.vstack([np.random.randint(0, 50256, input_len) for _ in range(batch_size)])
	test_data = np.load("test.npy")
	test_data = np.vstack([np.random.randint(0, 50256, input_len) for _ in range(batch_size)])