[Bugfix] record cos and sin cache in AscendRotaryEmbedding

[Debonex]

What this PR does / why we need it?

In scenarios where models like Moonlight (using MLA but without rope_scaling in config.json) invoke AscendRotaryEmbedding. _cos_cache and _sin_cache are not recorded correctly.

Does this PR introduce any user-facing change?

No

How was this patch tested?

• vLLM version: v0.13.0
• vLLM main: vllm-project/vllm@45c1ca1

[gemini-code-assist]

Code Review

This pull request addresses a bug where _cos_cache and _sin_cache for rotary embeddings were not being correctly recorded in certain scenarios. The fix involves introducing a new function, _record_cos_and_sin_cache_interleaved, to derive these caches from the combined cos_sin_cache and calling it during the initialization of AscendRotaryEmbedding. My review focuses on improving the implementation of this new function for better readability and maintainability, as the current version is overly complex for its task.

[gemini-code-assist]

The implementation of this function is overly complex and hard to follow. It can be significantly simplified for better readability and maintainability. The current logic is equivalent to splitting the cos_sin_cache and repeating each part, which can be expressed more directly using chunk and repeat.

Additionally, the function name _record_cos_and_sin_cache_interleaved is misleading because the cos_sin_cache tensor is concatenated ([cos, sin]), not interleaved ([c0, s0, c1, s1, ...]). Consider renaming it to something like _record_cos_and_sin_cache_from_combined to more accurately reflect its purpose.

Suggested change

	def _record_cos_and_sin_cache_interleaved(cos_sin_cache):
	global _cos_cache
	global _sin_cache
	if _cos_cache is not None or _sin_cache is not None:
	return
	hidden_dim = cos_sin_cache.shape[-1] // 2
	cos_cache, sin_cache = cos_sin_cache.view(-1, 2, hidden_dim).repeat(
	1, 1, 2).chunk(2, dim=1)
	_cos_cache = cos_cache.squeeze(1)
	_sin_cache = sin_cache.squeeze(1)
	def _record_cos_and_sin_cache_interleaved(cos_sin_cache):
	global _cos_cache
	global _sin_cache
	if _cos_cache is not None or _sin_cache is not None:
	return
	# cos_sin_cache is concatenated from cos and sin, each of size rotary_dim/2.
	cos_part, sin_part = cos_sin_cache.chunk(2, dim=-1)
	# For neox style, cos and sin are duplicated to match rotary_dim.
	_cos_cache = cos_part.repeat(1, 2)
	_sin_cache = sin_part.repeat(1, 2)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.