[askliar]

📌 Description

The BF16 GDN MTP decoder recompiles for changing batch sizes when serving paged state pools. Pool resizing also invalidates warmup specializations.

Strided descriptors now make only the outer batch and pool dimensions symbolic. Exact strides, inner geometry, dtype, device, and launch tactics remain specialized. TVM-FFI binds actual tensors at launch. SUPPORTS_DYNAMIC_STRIDED_SHAPES=True lets consumers enable compatible warmup.

🔍 Related Issues

Consumer: vLLM PR #65657.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

Changed-file pre-commit hooks passed. The all-files run was not performed.

🧪 Tests

• Tests have been added or updated as needed.

• All tests are passing (unittest, etc.).

• This PR requires full-scope CI, not targeted tests (expensive; justify below).
  Reason: This change affects one decoder and its existing regression suite. Targeted CI is sufficient.

  Suggested CI test targets (optional; paths under tests/, whitespace-separated — a starting
  point for @flashinfer-bot run <targets> / /bot run <targets>, reviewers may adjust): tests/gdn/test_decode_ucache.py

Equivalent focused command on NVIDIA GB300, from a checkout with the candidate package and dependencies installed:

uv run --no-project --python .venv/bin/python .venv/bin/python -m pytest --noconftest --import-mode=importlib -q \
  tests/gdn/test_decode_ucache.py::test_strided_batch_and_pool_reuse \
  tests/gdn/test_decode_ucache.py::test_strided_ab_matches_compact

Provider	Focused result	Test duration
Released 0.7.3, test change only	1 failed, 1 passed	9.95 s
Released 0.7.3 with this provider change	2 passed	10.60 s

The baseline failed the intended cache-reuse assertion after batch and pool sizes changed. Output and pool comparisons passed before that assertion.

The regression covers T=8, changing batch/pool sizes, singleton views, page strides, flushes, wrapped rings, and graph replay with negative indices. It compares output, state, and rings exactly against the compact path.

The provider-only model comparison used NVIDIA GB300 and nvidia/Qwen3.5-122B-A10B-NVFP4, with BF16 recurrent state and natural MTP7. Both runs used vLLM b95a484d72ba, batch size 256, prompt lengths 36–146, 256 output tokens per request, and disabled prefix caching.

Phase	Baseline compile/load calls	Candidate compile/load calls
Workload warmup	11	0
Measured generation	1	0
Factual output checks	0	1

Baseline workload warmup spent 10.235 s in compile/load calls. Its measured-generation miss took 0.949 s. The candidate's remaining factual-request miss took 1.470 s.

The remaining miss selects a distinct input stride and launch tactic. The companion vLLM warmup covers that combination and requires the capability marker.

A final run used vLLM c350280dacca with the same workload and this provider patch. All 12 variants compiled during startup warmup.
Those variants cover MBP 2–7 for QKV strides 12,288 and 20,480. Workload warmup, measured generation, and factual prompts recorded zero decoder compile/load calls.
All four factual token arrays matched the baseline and stopped normally. The final run had zero preemptions.

Compile logging affects host timing. These results do not establish an uninstrumented throughput speedup. All four factual answers matched and stopped normally. Timed continuations were not all identical, so this is not full-model parity evidence.

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

The compact path and recurrence arithmetic are unchanged. The prefix materializer's separate pool-capacity specialization remains outside this decoder fix. BF16 STP is also outside this change.

The model comparison used an isolated patched 0.7.3 provider. It does not show this behavior in an unpatched release. Full-suite tests were not run.

AI assistance: OpenAI Codex.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn