Skip to content

[KV Cache] Zero compressed attention storage pages - #373

Merged
yangzhuxinyzx merged 1 commit into
1CatAI:mainfrom
Leonccaa:fix/qwen38-flash-next-compressed-qsa-zeroer
Aug 28, 2026
Merged

yangzhuxinyzx merged 1 commit into
1CatAI:mainfrom
Leonccaa:fix/qwen38-flash-next-compressed-qsa-zeroer

Conversation

@Leonccaa

Copy link
Copy Markdown
Contributor

Summary

  • derive the zeroer's physical attention page layout from storage_block_size
  • keep the existing virtual-block split for ordinary full attention pages
  • cover ordinary and compressed layouts with real cache specs and a CUDA page-zeroing test

Problem

Compressed attention exposes a logical scheduler block that is larger than its
physical cache page. For example, Qwen3.8 Flash Next uses a logical block size
of 784 and a compression ratio of 8, so its compressed QSA cache stores 98 rows
per physical page.

KVBlockZeroer previously formed its virtual-page ratio from the logical block
size and the group's kernel block size. For a compressed cache this multiplied
the physical span by the compression ratio again, so a reused block could clear
the wrong physical range instead of exactly one compressed storage page.

Fix

Resolve the physical kernel block size and virtual-page ratio from the cache
spec before deriving segment spans:

  • ordinary attention: keep the group kernel block size and its existing ratio
  • compressed attention: use storage_block_size as one physical page with ratio 1
  • reject layouts whose storage block is not divisible by the selected kernel block

Validation

  • 3 passed, 1 skipped for the focused controller test selection (CUDA test skipped there)
  • Ruff check and format check pass for both changed files
  • CUDA smoke on a Tesla V100-PCIE-32GB:
    • ordinary layout: (16, 49)
    • compressed layout: (98, 1)
    • zeroing block 1 clears only physical page 1
  • Qwen3.8 Flash Next 125B NVFP4, TP4 on 4x V100 completed model load,
    4 GiB/GPU KV-cache initialization, PIECEWISE+FULL CUDA Graph capture, and
    1K/4K generation with the equivalent patch before rebasing onto current main

Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
@yangzhuxinyzx
yangzhuxinyzx marked this pull request as ready for review August 28, 2026 03:44
@yangzhuxinyzx

Copy link
Copy Markdown
Contributor

审计通过:压缩注意力零页逻辑按物理 storage_block_size 解析,普通页行为保持不变;锁定提交在最新 main 的隔离 worktree 上通过 13 项定向测试(含 V100 CUDA 实页清零),完整 pre-commit 与 DCO 通过。按锁定 SHA 合并。

@yangzhuxinyzx
yangzhuxinyzx merged commit 2b631b9 into 1CatAI:main Aug 28, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants