fix(driver/cuda): size attention float workspace from prefill split-KV bound

[ingim]

Summary

Fixes #414 — the intermittent batch_prefill_tmp_s ... only 0 bytes available in AlignedAllocator. Increase the workspace buffer size. warning during prefill on larger GPUs / certain head configs.

flashinfer's PrefillPlan carves the split-KV partial outputs out of the float workspace as:

tmp_v = num_qo_heads * padded_batch_size * cta_tile_q * head_dim_vo * 4
tmp_s = num_qo_heads * padded_batch_size * cta_tile_q              * 4

padded_batch_size is bounded by max_batch_size_if_split, and the plan hardcodes num_blocks_per_sm = 2, giving per rank max_batch_size_if_split = ceil(2 * num_sm / num_kv_heads_local); cta_tile_q tops out at 128 (64 for head_dim >= 256). The requirement therefore scales with SM count and head config, neither of which a fixed buffer can track.

Root cause in our code

attention_float_workspace_bytes() already received a cudaDeviceProp but ignored it, sizing instead from a decode-shaped heuristic (cta_tile_q = 16, padded_batch keyed off max_requests). That only over-provisioned by accident when max_requests was large, and it early-returned the base for GQA ratios outside the decode fast-path, sliding-window layouts, and TP>1 — all of which still hit the prefill kernel (force_prefill_path is set precisely when GQA isn't decode-supported).

Change

Compute the actual prefill split-KV bound from prop.multiProcessorCount and the per-rank head config, sized unconditionally for every arch that reaches the prefill kernel. The bound is tp-invariant (qo_heads shrinks as num_kv_heads shrinks); base (80 MiB, 128 MiB for Qwen hybrids) remains a floor.

Resulting sizes:

Config	New workspace	Old behavior
Qwen3-8B, L40S (142 SM)	96 MiB	64 MiB → overflow (the report)
Qwen3-8B, L40S, TP=2	96 MiB (tp-invariant)	80 MiB base
32q / 2kv, L40S	304 MiB	80 MiB base → overflow
64q / 8kv, B200 (148 SM)	176 MiB	80 MiB base

Validation

• The formula matches the flashinfer prefill plan in v0.6.9 (the pinned version), v0.6.12, and current main — that sizing block is byte-identical across all three, so this is not coupled to a transient implementation detail.
• flashinfer exposes no workspace-size query API for attention (only its CUTLASS GEMM paths do); its own recommended default is a flat 128 MiB, which is itself insufficient for low-KV-head / high-SM configs — hence the device-aware computation here.
• Sizing arithmetic checked standalone; not yet built in-tree or run on GPU. Needs a CUDA build + tests/inferlets/test_async_lm.py --driver cuda_native on an affected GPU (e.g. L40S) to confirm the warning is gone.

Notes / follow-ups

• parity_harness.cpp still calls AttentionWorkspace::allocate() with the 80 MiB default (it doesn't go through the planner); left as-is since it's a dev-only path.
• A runtime guard at the PrefillPlan call sites (catch the bump-allocator failure / compare computed-vs-available bytes) would be the only thing fully robust to a future flashinfer change to the sizing constants; out of scope here.

Summary by CodeRabbit

• Refactor
  • Optimized GPU memory allocation calculations for attention computations, reducing conditional constraints and improving buffer sizing logic based on GPU architecture properties.

[coderabbitai]

Caution

Review failed

Pull request was closed or merged during review

Walkthrough

attention_float_workspace_bytes in attention_workspace.cpp is refactored to compute the float workspace size from GPU SM count (prop.multiProcessorCount) and per-rank head geometry instead of a fixed default. Early-return guards tied to tp_size, max_requests, GQA decode support, and attention layout are removed. Buffer sizing now derives padded_batch_size from SM count and KV heads, and cta_tile_q from head_dim.

Changes

Dynamic float workspace sizing

Layer / File(s)	Summary
Simplified validity guards driver/cuda/src/attention_workspace.cpp	Removes early-return conditions based on tp_size == 1, max_requests > 0, GQA-in-decode support, sliding_window/non-full_attention layer checks, and head-count divisibility. Replaces them with a minimal validity check on head counts, supported head_dim_kernel, and positive multiProcessorCount. The Qwen hybrid base calculation is unchanged.
SM- and head-geometry-derived buffer sizing driver/cuda/src/attention_workspace.cpp	Replaces cta_tile_q = 16 and request-aligned batch padding with padded_batch_size derived from prop.multiProcessorCount and per-rank KV heads, and cta_tile_q derived from head_dim. Recomputes tmp_v and tmp_s using these values; max_requests is no longer used. Returns max(base, align_up(tmp_v + tmp_s + slack)).

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix(driver/cuda): size attention float workspace from prefill split-KV bound' directly and specifically describes the main change: fixing workspace sizing by computing it from prefill split-KV requirements instead of using decode heuristics.
Linked Issues check	✅ Passed	The pull request fully addresses the coding requirements of issue #414: it computes float_workspace_bytes dynamically using cudaDeviceProp.multiProcessorCount and per-rank head configuration, eliminating the 64 MiB fixed allocation that caused intermittent prefill workspace overflow errors.
Out of Scope Changes check	✅ Passed	All changes in attention_workspace.cpp are directly related to the stated objective of fixing workspace sizing; the refactoring removes decode-path gating and replaces it with SM-based prefill split-KV bound computation, which is entirely in scope.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/prefill-split-kv-workspace-414

---

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

_{Comment @coderabbitai help to get the list of available commands and usage tips.}