[MLA][Graph] Improve assertion on Graph mode with MLA

[MengqingCao]

What this PR does / why we need it?

Improve assertion on Graph mode with MLA.

When running deepseek with graph mode, the fused MLA op only support numHeads / numKvHeads ∈ {32, 64, 128}, thus we improve the assertion info here to avoid users confused with this.

Does this PR introduce any user-facing change?

Adjusting tp size is required when running deepseek-v3/r1 with graph mode. deepseek-v2-lite is not supported in graph mode.

How was this patch tested?

Test locally as the CI machine could not run V3 due to the HBM limits.

[linfeng-yuan]

LGTM. Note that MLA kernel may support numHeads / numKvHeads < 16 at future.

[MengqingCao]

thanks for the info！

…

---- Replied Message ---- | From | ***@***.***> | | Date | 05/22/2025 20:58 | | To | ***@***.***> | | Cc | Mengqing ***@***.***>***@***.***> | | Subject | Re: [vllm-project/vllm-ascend] [MLA][Graph] Improve assertion on Graph mode with MLA (PR #933) | linfeng-yuan left a comment (vllm-project/vllm-ascend#933) LGTM. Note that MLA kernel may support numHeads / numKvHeads < 16 at future. — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.Message ID: ***@***.***>

[Yikun]

LGTM. Note that MLA kernel may support numHeads / numKvHeads < 16 at future.

Maybe we should add note or TODO on code and doc current limit supported in FAQ or somewhere.

[MengqingCao]

Maybe we should add note or TODO on code and doc current limit supported in FAQ or somewhere.

Done, PTAL, thanks!

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[guanyuzhu]

assert self.num_queries_per_kv in _ALLOWED_NUM_QUERIES_PER_KV, \
            ("The allowed number of queries per kv when enabling both MLA and Graph mode"
            " only support {32, 64, 128}, Thus this is not supported for DeepSeek-V2-Lite,"
            " as it only has 16 attention heads. And if you're using DeepSeek-V3 or DeepSeek-R1,"
            " please make sure after the tensor parallel split, num_heads / num_kv_heads in "
            "{32, 64, 128}.")

The constraint here is affected by what? Why does it have to be 32/64/128? Can other values like 28 be used (or why is it not recommended)? Because there's a user-defined model with num_heads=28 and num_kv_heads=1, which fails to run here.

[MengqingCao]

assert self.num_queries_per_kv in _ALLOWED_NUM_QUERIES_PER_KV, \
            ("The allowed number of queries per kv when enabling both MLA and Graph mode"
            " only support {32, 64, 128}, Thus this is not supported for DeepSeek-V2-Lite,"
            " as it only has 16 attention heads. And if you're using DeepSeek-V3 or DeepSeek-R1,"
            " please make sure after the tensor parallel split, num_heads / num_kv_heads in "
            "{32, 64, 128}.")

The constraint here is affected by what? Why does it have to be 32/64/128? Can other values like 28 be used (or why is it not recommended)? Because there's a user-defined model with num_heads=28 and num_kv_heads=1, which fails to run here.

This constraint comes from cann op, and it will be moved after #1653 and #1508

[guanyuzhu]

1.Although the restriction on multiples of 32 has been lifted, is it still recommended to set parameters as multiples of 32? I've seen analyses suggesting that using multiples of 32 allows for hardware alignment and full utilization of hardware resources.
2.Just to confirm, after lifting this restriction, configurations with 28 heads should also work, right? (Given: head_size=576, num_heads=32, num_kv_heads=1)

[guanyuzhu]

1.Although the restriction on multiples of 32 has been lifted, is it still recommended to set parameters as multiples of 32? I've seen analyses suggesting that using multiples of 32 allows for hardware alignment and full utilization of hardware resources.
2.Just to confirm, after lifting this restriction, configurations with 28 heads should also work, right? (Given: head_size=576, num_heads=32, num_kv_heads=1)

assert self.num_queries_per_kv in _ALLOWED_NUM_QUERIES_PER_KV, \
            ("The allowed number of queries per kv when enabling both MLA and Graph mode"
            " only support {32, 64, 128}, Thus this is not supported for DeepSeek-V2-Lite,"
            " as it only has 16 attention heads. And if you're using DeepSeek-V3 or DeepSeek-R1,"
            " please make sure after the tensor parallel split, num_heads / num_kv_heads in "
            "{32, 64, 128}.")

The constraint here is affected by what? Why does it have to be 32/64/128? Can other values like 28 be used (or why is it not recommended)? Because there's a user-defined model with num_heads=28 and num_kv_heads=1, which fails to run here.

This constraint comes from cann op, and it will be moved after #1653 and #1508

1.Although the restriction on multiples of 32 has been lifted, is it still recommended to set parameters as multiples of 32? I've seen analyses suggesting that using multiples of 32 allows for hardware alignment and full utilization of hardware resources.
2.Just to confirm, after lifting this restriction, configurations with 28 heads should also work, right? (Given: head_size=576, num_heads=32, num_kv_heads=1)