[Bugfix] Fix DeepseekV32 AssertionError: num_kv_heads == 1

[chaunceyjiang]

Purpose

Fix #33074

Introduced #30207

vllm/vllm/distributed/kv_transfer/kv_connector/utils.py

Lines 322 to 326 in 64e3d67

	# Figure out whether the first dimension of the cache is K/V
	# or num_blocks. This is used to register the memory regions correctly.
	kv_cache_shape = self.attn_backend.get_kv_cache_shape(
	num_blocks=1, block_size=16, num_kv_heads=4, head_size=1
	)

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses a bug by removing a restrictive assertion related to num_kv_heads in the DeepseekV32IndexerBackend. While this resolves the immediate AssertionError, it's crucial to ensure that the backend's get_kv_cache_shape and underlying attention kernels correctly handle scenarios where num_kv_heads is greater than 1. The PR description could benefit from more details regarding the specific DeepseekV3.2 configurations that triggered the error and how the backend is expected to behave with varying num_kv_heads values. Additionally, filling out the 'Purpose', 'Test Plan', and 'Test Result' sections in the PR body would greatly enhance clarity and reviewability.

[chaunceyjiang]

vllm/vllm/distributed/kv_transfer/kv_connector/utils.py

Lines 322 to 326 in 64e3d67

	# Figure out whether the first dimension of the cache is K/V
	# or num_blocks. This is used to register the memory regions correctly.
	kv_cache_shape = self.attn_backend.get_kv_cache_shape(
	num_blocks=1, block_size=16, num_kv_heads=4, head_size=1
	)

/cc @NickLucche

I’m not sure whether this fix is correct. PTAL.

[NickLucche]

@chaunceyjiang I think the issue is in how we're using the function, let me look into it

[NickLucche]

#33090

[chaunceyjiang]

Hi @NickLucche, for a change like #33090, I’m inclined to remove this assert. This would make the use of get_kv_cache_shape more flexible.

[chaunceyjiang]

same as: FlashMLASparseBackend

[chaunceyjiang]

@NickLucche WDYT?