[WIP][Kernel] Enable GDN batch-invariant prefix caching

[quanliu1991]

Purpose

Enable VLLM_BATCH_INVARIANT=1 together with prefix caching on Qwen GDN hybrid models.

Today this fails at scheduler init (NotImplementedError: ... does not yet support Mamba cache alignment) because prefix cache puts GDN into mamba align mode while BIC needs a 64-token FLA grid.

This PR:

• Uses one shared grid lcm(FLA_CHUNK_SIZE=64, mamba_state_block_size) for prefill chunks and prefix-cache hits (hits are floored onto the grid; a full-prompt hit is not).
• Rounds hybrid block_size up to a multiple of 64 (1056 → 1088) so the grid stays usable for short prefixes.
• Temporarily includes GDN BIC from [Model] Enable batch-invariant mixed decode and prefill for Qwen GDN #49827 / [Feature] Add batch invariance support to GDN_ATTN backend #45819 so CI can run while those PRs are still open. Drop that commit when they land.

Related: #49827, #45819, #48613, #42960.

Test Plan

Set VLLM_GDN_PC_TEST_MODEL to a GDN checkpoint (e.g. Qwen3.5-35B-A3B). Tests skip without it.

export VLLM_BATCH_INVARIANT=1
export VLLM_GDN_PC_TEST_MODEL=/path/to/Qwen3.5-35B-A3B
export VLLM_TP_SIZE=2
export VLLM_ENABLE_V1_MULTIPROCESSING=0
pytest tests/v1/determinism/test_gdn_prefix_cache_batch_invariant.py -s -v

TTFT (start vllm serve with BI=1, then PC on / off):

BASE_URL=http://127.0.0.1:8000 TAG=on \
  bash tests/v1/determinism/run_gdn_prefix_cache_bench.sh

Hardware for the numbers below: 2× L20, TP=2, Qwen3.5-35B-A3B, vllm-openai:v0.28.0 overlay with the same changes.

Test Result

Before: engine does not start (NotImplementedError above).

After (bitwise): engine starts, block_size 1056 → 1088, alignment=1088. Same-path hit BS=1 vs BS=N is bitwise identical (prefix-1100 real hit; ragged-200 floors to 0). RESULT: PASS.

After (TTFT, shared 2000-token prefix, BI=1):

prefix cache	c	mean TTFT (ms)	median TTFT (ms)	mean TPOT (ms)
off	1	261	262	18.7
on	1	210	203	18.5
off	16	2154	2020	145.8
on	16	1670	1469	153.3

c=1 TTFT −19%, c=16 median TTFT −27%. TPOT unchanged. 2000 tokens only skip the first 1088-token page, so this is not decode-only TTFT.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[quanliu1991]

@claude review