[HMA]Move hybrid blksize to update_block_size_for_backend to fix attn supported block size is not 16 issue

[xuechendi]

Purpose

--

Issue description:

Testing nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-bf16 with default block_size = 64.

if FA block_size = 16 w/ TP4 => PASS: 540672 / 8192 = 66

[kv_cache_utils.py:934] layer_spec.block_size=16, layer_spec.page_size_bytes=8192
[kv_cache_utils.py:934] layer_spec.block_size=262144, layer_spec.page_size_bytes=540672

if FA block_size = 64 w/ TP4 => ERROR!!! SSM_page_size will not be evenly divided by FA_page_size: 540672/32768 = 16.5

[kv_cache_utils.py:934] layer_spec.block_size=64, layer_spec.page_size_bytes=32768
[kv_cache_utils.py:934] layer_spec.block_size=262144, layer_spec.page_size_bytes=540672

--

Root Cause:

Default hybrid model block_size alignment is happened before platform.check_and_update_config
sequence is
platform init =>cache_config init=> model_config init(hybrid model block_size alignment) => __post_init => platform.check_and_update_config

because of the sequence, current alignment for hybrid model is only calculated based on block_size=16. While XPU block_size will be updated to 64 in platform.check_and_update_config, which cause the unalignment issue.

--

Solution:

After discussion with Reviewers in this PR, suggested to update block_size in platform.update_block_size_for_backend for XPU.

Including:

• redo hybrid model block_size alignment.
• update block_size per attn_backend.

--

Test Plan - re-tested on 0323 with latest fix

validation after fix

case 0 - Hybrid model

lm_eval   \
--model vllm   \
--model_args pretrained=/mnt/data/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-bf16,tensor_parallel_size=4,trust_remote_code=True,enable_expert_parallel=True,attention_backend=TRITON_ATTN   \
--tasks gsm8k   \
--num_fewshot 5   \
--batch_size auto

(Worker_TP0_EP0 pid=250806) INFO 03-23 20:58:48 [default_loader.py:384] Loading weights took 5.72 seconds
(Worker_TP0_EP0 pid=250806) INFO 03-23 20:58:49 [gpu_model_runner.py:4596] Model loading took 14.76 GiB memory and 7.550555 seconds
(Worker_TP0_EP0 pid=250806) INFO 03-23 20:58:49 [xpu.py:287] Update hybrid model block size to 1088
(Worker_TP2_EP2 pid=250808) INFO 03-23 20:58:49 [xpu.py:287] Update hybrid model block size to 1088
(Worker_TP1_EP1 pid=250807) INFO 03-23 20:58:50 [xpu.py:287] Update hybrid model block size to 1088
(Worker_TP3_EP3 pid=250809) INFO 03-23 20:58:50 [xpu.py:287] Update hybrid model block size to 1088

case 1 - full attn model

 lm_eval   \
--model vllm   \
--model_args pretrained=Qwen/Qwen3-30B-A3B-Thinking-2507,tensor_parallel_size=4,trust_remote_code=True,enable_expert_parallel=True   \
--tasks gsm8k   \
--num_fewshot 5  \
--batch_size auto

(Worker_TP0_EP0 pid=251994) INFO 03-23 22:12:58 [default_loader.py:384] Loading weights took 16.17 seconds
(Worker_TP1_EP1 pid=251995) INFO 03-23 22:12:59 [xpu.py:306] Update FLASH_ATTN block size to 64
(Worker_TP0_EP0 pid=251994) INFO 03-23 22:12:59 [gpu_model_runner.py:4596] Model loading took 14.3 GiB memory and 133.598931 seconds
(Worker_TP0_EP0 pid=251994) INFO 03-23 22:12:59 [xpu.py:306] Update FLASH_ATTN block size to 64
(Worker_TP2_EP2 pid=251996) INFO 03-23 22:12:59 [xpu.py:306] Update FLASH_ATTN block size to 64
(Worker_TP3_EP3 pid=251997) INFO 03-23 22:12:59 [xpu.py:306] Update FLASH_ATTN block size to 64

[gemini-code-assist]

Code Review

This pull request updates the KV cache page size unification logic to handle models with non-integral page size ratios between layers. It replaces the use of max() with math.lcm() to determine the unified page size, which is a more robust approach for this scenario. The changes also correctly propagate this scaling to page_size_padded if it is present in the cache specification. While this change is correct, I've identified a potential issue where the calculated LCM could become excessively large, leading to high memory consumption. I've added a high-severity comment with a suggestion to add a safeguard. An unrelated change for XPU support is also included in vllm/model_executor/layers/fused_moe/layer.py.

[mergify]

Hi @xuechendi, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @xuechendi, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @xuechendi, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[yma11]

There is another fix at #37429 and I think that one is considering more?

[NickLucche]

2 things off the top of my head

• lm_eval scores are lower than they should be for this model (iirc we're at 80+ for nemotronH @ZhanqiuHu )
• what's the setup that is leading to this assert? I don't recall seeing this issue before.
  In particular

has two page_size for linear_attn and full_attn

this is only true when kernel_block_size is required, and then again in that case the page_size is divided by the logical/physical ratio (duplicating num_blocks accordingly).

Let's look deeper into the cause of this.

[xuechendi]

@NickLucche

this is only true when kernel_block_size is required, and then again in that case the page_size is divided by the logical/physical ratio (duplicating num_blocks accordingly).

Oh, I am testing with Intel GPU which uses block_size = 64

Just tested with block_size = 16, and we won't need this FIX. I'll update the title and purpose

---

for accuracy, I am following the config here - https://github.com/vllm-project/vllm/blob/main/.buildkite/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml

If using full set (no limit) and without multiturn, strict-match acc gets to 0.837. Wondering if you're testing with lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml or using a different gsm8k setting?

[ZhanqiuHu]

If using full set (no limit) and without multiturn, strict-match acc gets to 0.837. Wondering if you're testing with lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml or using a different gsm8k setting?

Yes, without apply_chat_template and fewshot_as_multiturn, the expected acc is ~0.84.

[ZhanqiuHu]

@xuechendi btw was this function called: verify_and_update_config?

[xuechendi]

we'd better rename title before merge.

Any suggestion?

[jikunshang]

something like Refactor block size selection/update and fix hybrid model alignment

[yma11]

@xuechendi I found on platform XPU, the final block size is 64 dividable using FLASH_ATTN but not using TRITON_ATTN, which will cause gdn attn incorrect output. How can we assure the final block size is 64 dividable no matter what the attention backend is.

INFO 03-30 07:05:40 [config.py:228] Setting attention block size to 528 tokens to ensure that attention page size is >= mamba page size.
INFO 03-30 07:05:40 [config.py:259] Padding mamba page size by 0.76% to ensure that mamba page size and attention page size are exactly equal.

[xuechendi]

@xuechendi I found on platform XPU, the final block size is 64 dividable using FLASH_ATTN but not using TRITON_ATTN, which will cause gdn attn incorrect output. How can we assure the final block size is 64 dividable no matter what the attention backend is.

INFO 03-30 07:05:40 [config.py:228] Setting attention block size to 528 tokens to ensure that attention page size is >= mamba page size.
INFO 03-30 07:05:40 [config.py:259] Padding mamba page size by 0.76% to ensure that mamba page size and attention page size are exactly equal.

I did same test for Triton, and since Triton should not limit block_size to 64, the logic here is to set as 16, not sure why you're seeing misalign, did you set triton_attn block_size=64 somewhere in your codes?

[yma11]

@xuechendi I found on platform XPU, the final block size is 64 dividable using FLASH_ATTN but not using TRITON_ATTN, which will cause gdn attn incorrect output. How can we assure the final block size is 64 dividable no matter what the attention backend is.

INFO 03-30 07:05:40 [config.py:228] Setting attention block size to 528 tokens to ensure that attention page size is >= mamba page size.
INFO 03-30 07:05:40 [config.py:259] Padding mamba page size by 0.76% to ensure that mamba page size and attention page size are exactly equal.

I did same test for Triton, and since Triton should not limit block_size to 64, the logic here is to set as 16, not sure why you're seeing misalign, did you set triton_attn block_size=64 somewhere in your codes?

TRITON_ATTN has no limitation on block size to be dividable by 64, but in gdn attention, we do computation leveraging on kv cache so it's still need block-size dividable by 64.

[yma11]

@xuechendi I found on platform XPU, the final block size is 64 dividable using FLASH_ATTN but not using TRITON_ATTN, which will cause gdn attn incorrect output. How can we assure the final block size is 64 dividable no matter what the attention backend is.

INFO 03-30 07:05:40 [config.py:228] Setting attention block size to 528 tokens to ensure that attention page size is >= mamba page size.
INFO 03-30 07:05:40 [config.py:259] Padding mamba page size by 0.76% to ensure that mamba page size and attention page size are exactly equal.

I did same test for Triton, and since Triton should not limit block_size to 64, the logic here is to set as 16, not sure why you're seeing misalign, did you set triton_attn block_size=64 somewhere in your codes?

TRITON_ATTN has no limitation on block size to be dividable by 64, but in gdn attention, we do computation leveraging on kv cache so it still need block-size dividable by 64. This is XPU gdn attention limitation for the time being.

[xuechendi]

@xuechendi I found on platform XPU, the final block size is 64 dividable using FLASH_ATTN but not using TRITON_ATTN, which will cause gdn attn incorrect output. How can we assure the final block size is 64 dividable no matter what the attention backend is.

INFO 03-30 07:05:40 [config.py:228] Setting attention block size to 528 tokens to ensure that attention page size is >= mamba page size.
INFO 03-30 07:05:40 [config.py:259] Padding mamba page size by 0.76% to ensure that mamba page size and attention page size are exactly equal.

I did same test for Triton, and since Triton should not limit block_size to 64, the logic here is to set as 16, not sure why you're seeing misalign, did you set triton_attn block_size=64 somewhere in your codes?

TRITON_ATTN has no limitation on block size to be dividable by 64, but in gdn attention, we do computation leveraging on kv cache so it's still need block-size dividable by 64.

Share more details? Is that specific to current xpu_kv_cache impl? If that is only for GDN, maybe we can fix in the Qwen3.5 PR since I am not able to test it right now?

[MatthewBonanni]

Since #33657 hasn't landed yet, I think this PR is fine as is for now. If GDN does truly require block size 64 for all backends on XPU, that change can be made as part of 33657