Skip to content

[BugFix] Fix attention buffer sizing for per-layer KV heads #8043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chang-wenbin merged 1 commit into from
Jun 14, 2026
+2 −1
Merged

[BugFix] Fix attention buffer sizing for per-layer KV heads #8043

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion fastdeploy/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1749,7 +1749,8 @@ def _initialize_attn_backend(self) -> None:
decoder_block_shape_q=decoder_block_shape_q,
decoder_step_token_num=self.speculative_config.num_speculative_tokens + 1,
num_heads=num_heads,
kv_num_heads=max(kv_num_heads_per_layer),
# This requires the largest possible group size, corresponding to the smallest kv-num-heads.
kv_num_heads=min(kv_num_heads_per_layer),

@PaddlePaddle-bot PaddlePaddle-bot Jun 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Bug min(kv_num_heads_per_layer) 只覆盖了按 group_size 递增的 buffer，但 decode-unified 的 decode_block_indices 容量还会乘以 kv_num_heads，这里会在部分分层 KV-head 配置下低估容量。

allocate_decode_unified_related_buffer() 和 allocate_launch_related_buffer() 中的容量是 max_batch_size * kv_num_heads * max_num_chunk * q_tile_num，而 config_for_attention() 实际按 q_tile_num * kv_chunk_num * kv_num_heads 写入。比如 num_heads=64、decoder_step_token_num=1、per-layer KV heads 为 [8, 1] 时，按 min=1 只预留 1 * ceil(64 / 1 / 16) = 4 个 q/kv-head tile；layer 0 若使用 8 个 KV heads，会需要 8 * ceil(64 / 8 / 16) = 8 个，存在越界写或错误配置风险。

建议修复方式：
保留最小 KV heads 用于 decoder_batch_ids/encoder_batch_ids 等只按 group_size 放大的 buffer，但对 decode_block_indices 单独按所有 layer 的最坏值分配，例如计算 max(k * ceil(decoder_step_token_num * (num_heads // k) / 16) for k in kv_num_heads_per_layer) 后作为容量上界；或者扩展 allocation API，分别传入最大 group_size 和 decode_block_indices 容量。

block_size=self.fd_config.cache_config.block_size,
head_dim=head_dim,
dtype=self.model_config.dtype,
Expand Down
Loading