Skip to content

[Cute,Fwd,Sm90] Ceil div in paged kv manager to prevent size 0#2446

Open
imbr92 wants to merge 2 commits intoDao-AILab:mainfrom
adaptive-ml:ceil_div_hdim256
Open

[Cute,Fwd,Sm90] Ceil div in paged kv manager to prevent size 0#2446
imbr92 wants to merge 2 commits intoDao-AILab:mainfrom
adaptive-ml:ceil_div_hdim256

Conversation

@imbr92
Copy link
Copy Markdown
Contributor

@imbr92 imbr92 commented Apr 8, 2026

When running the SM90 cute-dsl fwd with kv cache with gemma 3 4B, I got ValueError: Expected size in shape to be strictly positive, but got 0 since n_block_size is 80 (or 64 for the local layers) but num_threads = 128. Added a ceil div to prevent this and added a correctness test to the existing cute test file. Also enabled test_flash_attn_paged_deepseek for SM90 since paged kv is now supported on SM90.

Tests:

Ran on H200.

test_flash_attn_paged_deepseek

image #### test_flash_attn_paged_hdim256 image

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.

1 participant