[Hybrid] Pass kernel block size to builders

[tdoublep]

Purpose

Solves #27264 and #26936

This PR makes two changes:

• The GPU model runner will now pass the kernel block size to the metadata builders, fixing a pretty bad bug that exists on main.
• I also restrict the kernel block sizes for the FlashAttention backend based on the discussion in [Bug]: Cache malformation in hybrid models with SSM cache dtype float32 and block allocation wrap around #27264. This is required since for block_size >= 128, FA will read partial block data that may contain NaN's if mamba state is in fp32, which leads to NaN's appearing in output.

Test Plan

I'm using the following test script:

from vllm import LLM, SamplingParams
import os

os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASH_ATTN'
#os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASHINFER'

prompts = ["Solve the following math problem step by step. The last line of your response should be of the form Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\nConsider the paths of length $16$ that follow the lines from the lower left corner to the upper right corner on an $8\\times 8$ grid. Find the number of such paths that change direction exactly four times, as in the examples shown below.\n\nRemember to put your answer on its own line after \"Answer:\"."]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="nvidia/NVIDIA-Nemotron-Nano-9B-v2",
    trust_remote_code=True,
    num_gpu_blocks_override=10,
    compilation_config={'cudagraph_capture_sizes': [1]},
    mamba_ssm_cache_dtype="float32",
    max_num_seqs=1,
    enable_prefix_caching=True,
)

outputs = []
for _ in range(10):
    outputs.append(llm.generate(prompts, sampling_params)[0])

# Print the outputs.
print(f"Prompt: {prompts[0]!r}")
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Generated text: {generated_text!r}")
    print(f"Generated token IDs: {output.outputs[0].token_ids!r}")

Which on main using FLASH_ATTN produces:

Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

and on main with FLASHINFER produces:

Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Generated text: ''
Generated token IDs: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Test Result

With this PR, FLASH_ATTN produces:

Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n\nExample 1: \n\nExample 2:\n\n\n<think>\n'
Generated token IDs: [1032, 1267, 20396, 1032, 1049, 1058, 1032, 1267, 20396, 1032, 1050, 2100, 1010, 49250, 2077, 1561]
Generated text: '</think>\nThe number of paths of length 16 from the lower left'
Generated token IDs: [1885, 74045, 1561, 1784, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: '</think>\nThe problem involves finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 19263, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' output: To solve the problem of finding the number of paths of length 1'
Generated token IDs: [4848, 1058, 3870, 15047, 1278, 4127, 1307, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049]
Generated text: ' \n</think>\nThe problem requires finding the number of paths of length '
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032]
Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n</think>\nThe paths of length 16 from the lower left'
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: " \n\n<think>\nOkay, let's see. I need to find the"
Generated token IDs: [1032, 1267, 49250, 2077, 1561, 44053, 1044, 2878, 1681, 3219, 1046, 1362, 2534, 1317, 3081, 1278]
Generated text: '<think>\nOkay, so I need to find the number of paths on an'
Generated token IDs: [1060, 74045, 1561, 44053, 1044, 1878, 1362, 2534, 1317, 3081, 1278, 2782, 1307, 22344, 1408, 1420]

and FLASHINFER produces:

Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n\nExample paths:\n- Path 1: Right, Right, Down,'
Generated token IDs: [1032, 1267, 20396, 22344, 1877, 1045, 17669, 1032, 1049, 1058, 21285, 1044, 21285, 1044, 16999, 1044]
Generated text: '</think>\nThe number of paths of length 16 from the lower left'
Generated token IDs: [1885, 74045, 1561, 1784, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: '</think>\nThe problem involves finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 19263, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' output: To solve the problem of finding the number of paths of length 1'
Generated token IDs: [4848, 1058, 3870, 15047, 1278, 4127, 1307, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049]
Generated text: ' \n</think>\nThe problem requires finding the number of paths of length '
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032]
Generated text: '</think>\nThe problem requires finding the number of paths of length 16'
Generated token IDs: [1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032, 1049, 1054]
Generated text: ' \n</think>\nThe paths of length 16 from the lower left'
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 22344, 1307, 5592, 1032, 1049, 1054, 1562, 1278, 4953, 3979]
Generated text: " \n\n<think>\nOkay, let's see. I need to find the"
Generated token IDs: [1032, 1267, 49250, 2077, 1561, 44053, 1044, 2878, 1681, 3219, 1046, 1362, 2534, 1317, 3081, 1278]
Generated text: ' \n</think>\nThe problem requires finding the number of paths of length '
Generated token IDs: [1032, 1010, 1885, 74045, 1561, 1784, 4127, 10867, 13170, 1278, 2782, 1307, 22344, 1307, 5592, 1032]

---

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.

[mergify]

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

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

[tdoublep]

@heheda12345 This is ready for review now.

[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".

[heheda12345]

@tdoublep my only concern is codex's comment. Other parts looks great.

[tdoublep]

@heheda12345 The v1-test-attention (B200) job is failing but it looks unrelated (the container doesn't even start-up).

[heheda12345]

LGTM!