[Bugfix] Ensure calculated KV scales are applied in attention.

[adabeyta]

Purpose:

Resolves bug #27102 .

Test Plan:

Throughput

vllm bench throughput   --model Qwen/Qwen3-8B   --quantization fp8   --kv-cache-dtype fp8_e4m3   --tensor-parallel-size 2   --dataset-name random --input-len 1024 --output-len 256

E2E Correctness
Server with KV scale calculation ON (remove --calculate-kv-scales flag for OFF case):

vllm serve Qwen/Qwen3-8B \
    --tensor-parallel-size 2 \
    --quantization fp8 \
    --kv-cache-dtype fp8_e4m3 \
    --calculate-kv-scales \
    --port 8000

lm_eval

lm_eval --model local-completions --model_args pretrained <model>,base_url=http://0.0.0.0:8000/v1/completions,num_concurrent=50,max_retries=3 --tasks gsm8k --num_fewshot 5 --batch_size auto --limit 100 

Throughput results

Main:calculate_kv_scales=False | enforce_eager=False


Throughput: 1.73 requests/s, 2210.02 total tokens/s, 442.00 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000



Main:calculate_kv_scales=True | enforce_eager=False

Throughput: 1.73 requests/s, 2211.36 total tokens/s, 442.27 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000

Main: calulate_kv_scales=True | enforce_eager=True

Throughput: 1.72 requests/s, 2202.42 total tokens/s, 440.48 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000

PR: calculate_kv_scales=False | enforce_eager=False


Throughput: 1.73 requests/s, 2211.37 total tokens/s, 442.27 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000

PR: calculate_kv_scales=True | enforce_eager=False


Throughput: 1.73 requests/s, 2210.78 total tokens/s, 442.16 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000

PR: calulate_kv_scales=True | enforce_eager=True

Throughput: 1.72 requests/s, 2200.57 total tokens/s, 440.11 output tokens/s
Total num prompt tokens:  1024000
Total num output tokens:  256000

E2E GSM8K Accuracy Results

Main: calulate_kv_scales=False | enforce_eager=False


|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.84|±  |0.0368|
|     |       |strict-match    |     5|exact_match|↑  | 0.83|±  |0.0378|```

Main: calulate_kv_scales=True | enforce_eager=False

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.85|±  |0.0359|
|     |       |strict-match    |     5|exact_match|↑  | 0.86|±  |0.0349|

Main: calulate_kv_scales=True | enforce_eager=True

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.88|±  |0.0327|
|     |       |strict-match    |     5|exact_match|↑  | 0.89|±  |0.0314|

PR: calculate_kv_scales=False | enforce_eager=False


|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.86|±  |0.0349|
|     |       |strict-match    |     5|exact_match|↑  | 0.86|±  |0.0349|

PR: calculate_kv_scales=True | enforce_eager=False


|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.86|±  |0.0349|
|     |       |strict-match    |     5|exact_match|↑  | 0.86|±  |0.0349|

PR: calulate_kv_scales=True | enforce_eager=True

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.89|±  |0.0314|
|     |       |strict-match    |     5|exact_match|↑  | 0.90|±  |0.0302|

[gemini-code-assist]

Code Review

This pull request correctly addresses a bug where calculated KV scales were not being applied during attention. The fix introduces a mechanism to calculate these scales on the first forward pass and then disables subsequent calculations by adding a state flag kv_scales_calculated to GPUModelRunner. The logic is sound and the changes are well-targeted. I've included one suggestion to refactor the implementation slightly, which will improve maintainability by removing a small piece of redundant logic.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[ProExpertProg]

Thx for the fix, I got a suggestion to further simplify the logic

[ProExpertProg]

Could you just sanity check with a deepseek model that it works with MLA too?

[ProExpertProg]

Seems like there's a few more paths left over

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @adabeyta.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ProExpertProg]

Just a CI file merge note