[arpera]

Purpose

Eliminate a GPU→CPU sync caused by .tolist() inside prepare_chunk_indices during GDN prefill.

prepare_chunk_indices (in vllm/model_executor/layers/fla/ops/index.py) calls .tolist() on a GPU tensor, which triggers a blocking GPU→CPU synchronization. The function is decorated with tensor_cache (identity-based LRU cache), so only the first call per step actually syncs — but that one sync is enough to stall the pipeline.

Fix: Pre-compute chunk_indices and chunk_offsets on CPU in GDNAttentionMetadataBuilder.build() (where cu_seqlens_cpu is already available, so no sync needed), async-copy them to GPU, and thread them as explicit optional parameters through the entire FLA ops chain. When provided, downstream ops skip both the computation and the tensor_cache lookup. The @tensor_cache fallback remains unchanged for callers that do not pass pre-computed values (KDA, OLMo Hybrid).

Also extracts hardcoded chunk_size=64 into FLA_CHUNK_SIZE constant.

Changes:

• vllm/model_executor/layers/fla/ops/utils.py: add FLA_CHUNK_SIZE constant
• vllm/v1/attention/backends/gdn_attn.py: pre-compute chunk_indices/chunk_offsets on CPU, store in GDNAttentionMetadata
• vllm/model_executor/layers/mamba/gdn_linear_attn.py: pass chunk_indices/chunk_offsets from metadata to FLA ops
• vllm/model_executor/layers/fla/ops/chunk.py, cumsum.py, chunk_o.py, chunk_delta_h.py, chunk_scaled_dot_kkt.py, solve_tril.py, wy_fast.py: accept optional chunk_indices (and chunk_offsets for chunk_delta_h), skip tensor_cache when provided
• vllm/model_executor/layers/fla/ops/kda.py: use FLA_CHUNK_SIZE constant

Test Plan

• Ran vllm serve nvidia/Qwen3.5-397B-A17B-NVFP4 with GDN attention and verified no GPU→CPU sync from prepare_chunk_indices in Nsight Systems profile.

Test Result

Before: Nsight Systems shows a GPU→CPU sync (.tolist()) on the first prepare_chunk_indices call per step during GDN prefill.
After: No GPU→CPU sync, no tensor_cache overhead on the GDN prefill path.

Perf & Eval Testing

Server (same for all tests):

vllm serve nvidia/Qwen3.5-397B-A17B-NVFP4 \
    --port 8000 -tp 1 -pp 1 -dp 8 \
    --enable-expert-parallel --language-model-only \
    --reasoning-parser qwen3 --stream-interval 100

Prefill Benchmark

vllm bench serve --backend vllm --model nvidia/Qwen3.5-397B-A17B-NVFP4 \
    --port 8000 --endpoint /v1/completions --dataset-name random \
    --random-input 32768 --random-output 1 --max-concurrency 128 \
    --num-prompt 128 --ignore-eos --temperature 0.0

	Total Throughput (tok/s)	Mean TTFT (ms)	Median TTFT (ms)	P99 TTFT (ms)
Before (main)	154,344	14,020	13,999	26,151
After (this PR)	153,575	14,274	14,103	26,204

GSM8K Eval

python3 tests/evals/gsm8k/gsm8k_eval.py

Config	Accuracy	Invalid Rate	Tokens/sec
Before (main)	0.879	0.024	1,474
After (this PR)	0.886	0.017	1,547

No perf or accuracy regression — results are within run-to-run noise.

Nsight Systems Profile (prefill, input 8192 tokens)

After this change, the first GDN layer in a step no longer stands out from subsequent GDN layers in CPU time — all GDN blocks are now uniform. Additionally, there are zero DtoH (GPU→CPU) memory copies during execution.

Before: DtoH memcpy is present (<0.1%), the first GDN block stalls on the GPU→CPU sync from .tolist(), forward ≈ 410 ms.

After: No DtoH memcpy at all, all GDN blocks have uniform CPU time, forward ≈ 385 ms.

---

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.

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

[gemini-code-assist]

Code Review

This pull request introduces a register method to the FLA operations cache utility, enabling manual insertion of entries. This functionality is used in the GDN attention backend to pre-compute chunk indices on the CPU and cache them against GPU keys, preventing performance-degrading GPU-to-CPU synchronization. Feedback includes a recommendation to refactor duplicated cache management logic into a helper function and to replace a hardcoded chunk size with a named constant.

[mergify]

Hi @arpera, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @arpera, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

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

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

[arpera]

cc: @vadiklyutiy

[vadiklyutiy]

/gemini review

[gemini-code-assist]

Code Review

This pull request centralizes the chunk size configuration by introducing a global FLA_CHUNK_SIZE constant across various FLA Triton kernels. It also enhances the tensor_cache utility with a new register method, which allows for manual insertion of entries into the cache. This functionality is utilized in the GDN attention backend to pre-compute chunk indices on the CPU and register them using GPU tensor keys, thereby avoiding performance-degrading GPU-to-CPU synchronizations. I have no feedback to provide as there were no review comments.

[vadiklyutiy]

@claude review

[vadiklyutiy]

I investigated a little bit.

chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) is always the same for fixed step: cu_seqlens and chunk_size the same within step.

So, actually we even don't need @tensor_cache - may just calculate once prepare_chunk_indices in MetaDataBuilder and use it.
Cons: implementation of above approach add new arg chunk_size to around 10 funcs.

But for prefill_size=1000, GDN CPU part takes 2x+ longer then GPU part, and default cudagraph_max_capture_size is 512, so we don't hide CPU in this case.

So, for this case make sense to speed up CPU part because it is a bottleneck. Searching in tensor_cache takes around 4% from GDN CPU times (purple on screenshot above).

What do you think?

[arpera]

I agree that storing chunk_indices in metadata would be a more clear solution in this case. Initially I tried that approach but the patch wasn't neat enough. If you think that storing chunk_indices in metadata would be a better approach, then I'll implement it.

[vadiklyutiy]

wasn't neat enough

what exact things looked not so good?

[mergify]

Hi @arpera, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[arpera]

Initially the patch itself was huge for such a small issue

[vadiklyutiy]

@claude review

[arpera]

@claude review

[ZJY0516]

@claude review

[ZJY0516]

@arpera Thanks for your contribution, do we have some perf and eval tests?