[BugFix] register quant scale tensors as buffer

[BoyuanFeng]

Prior to this pr, pytest -sv "tests/compile/test_fusion_attn.py::test_attention_quant_pattern[AttentionBackendEnum.TRITON_ATTN-nvidia/Llama-4-Scout-17B-16E-Instruct-FP8-TestAttentionFp8StaticQuantPatternModel-+quant_fp8-dtype1-256-128-40-8]" leads to illegal memory access error. I found that the IMA comes from torch.ops._C.static_scaled_fp8_quant as described in #31377.

@lengrongfu pointed out the issue is that scale tensor is not on GPU so IMA happened. In tests/compile/test_fusion_attn.py we actually have model_unfused.to(device). So why it is not on gpu?

vllm/tests/compile/test_fusion_attn.py

Line 365 in 87f1b8c

model_unfused = model_unfused.to(device)

It turns out that _k_scale, _v_scale, _q_scale, _prob_scale are not registered as buffers. model_unfused.to(device) only moves parameters and buffers. The fix is to register these 4 tensors as buffers.

This fixes the unit test failure.

Close: #31377

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request correctly addresses an illegal memory access error by registering the quantization scale tensors (_k_scale, _v_scale, _q_scale, _prob_scale) as buffers. This is the proper solution to ensure these tensors are moved to the correct device when model.to(device) is called. While this fix is correct, I've identified a similar critical issue that has been overlooked in the same file. The q_range, k_range, and v_range tensors in the Attention and MLAAttention classes are also initialized as raw tensors and not registered as buffers. When calculate_kv_scales is enabled, this will lead to a device mismatch RuntimeError during calculations, which is the same class of bug. To prevent this, I strongly recommend also registering q_range, k_range, and v_range as buffers.

[ProExpertProg]

This doesn't actually solve the issue for me, I had to run in a Docker container to repro:

$ docker run -it --runtime nvidia --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    -v ~/git/vllm:/vllm \ # path for the local vllm repo
    --ipc=host \
    --entrypoint /bin/bash \
    vllm/vllm-openai:latest
# in container
$ pip install tblib
$ pytest /vllm/tests/compile/test_fusion_attn.py -s -v

________________ test_attention_quant_pattern[AttentionBackendEnum.TRITON_ATTN-nvidia/Llama-4-Scout-17B-16E-Instruct-FP8-TestAttentionFp8StaticQuantPatternModel-+quant_fp8-dtype0-256-128-40-8] ________________

num_qo_heads = 40, num_kv_heads = 8, head_size = 128, batch_size = 256, dtype = torch.bfloat16, custom_ops = '+quant_fp8', model_name = 'nvidia/Llama-4-Scout-17B-16E-Instruct-FP8'
model_class = <class 'tests.compile.test_fusion_attn.TestAttentionFp8StaticQuantPatternModel'>, backend = <AttentionBackendEnum.TRITON_ATTN: 'vllm.v1.attention.backends.triton_attn.TritonAttentionBackend'>
dist_init = None

>   ???
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 1310660 / 1310720 (100.0%)
E   Greatest absolute difference: 6752.0 at index (156, 3552) (up to 0.01 allowed)
E   Greatest relative difference: inf at index (32, 0) (up to 0.01 allowed)

[BoyuanFeng]

@ProExpertProg there are 3 separate issues for attn_fusion. One is the numeric mismatch issue you have seen. @angelayi is fixing the issue.

Another one is an IMA error which is fixed by this PR.

There is also an implicit compilation config cache issue, fixed by #31376.

[ProExpertProg]

Ok, please post the PRs in #pr-merge-requests once CI is done so we can get the CI fixed ASAP

[robertgshaw2-redhat]

nice job @BoyuanFeng

[robertgshaw2-redhat]

note that the test failure is related

[BoyuanFeng]

@robertgshaw2-redhat the unit test failure comes from dummy weight load. This PR adds the q/k/v/prob scales to state dict. So dummy weight load randomly initializes its value. When quant is not used, model checkpoint does not contain q/k/v/prob scales so their values are still random values from dummy load, instead of expected default value (i.e., value 1).

The issue is fixed by updating process_weights_after_loading to handle cases that we should not load quant weights.