Skip to content

Enable causal blocking#908

Open
kdulla wants to merge 32 commits intoquic:mainfrom
kdulla:enable_causal_blocking
Open

Enable causal blocking#908
kdulla wants to merge 32 commits intoquic:mainfrom
kdulla:enable_causal_blocking

Conversation

@kdulla
Copy link
Copy Markdown
Contributor

@kdulla kdulla commented Apr 6, 2026 •
edited
Loading

Summary

  • This PR introduces QEfficient attention blocking for causal language models, adding Head, Batch, KV and Q blocking strategies that preserve numerical outputs and scalability at long sequence lengths.
  • New PR made to avoid rebase issues

Key Features

  • KV Blocking: Blocked compute over KV cache along sequence dim.
  • Q Blocking: Blocked compute over query sequence.
  • Head Blocking: Blocked compute over num heads, based on TS
  • Batch Blocking: Blocked compute over batch, still needs compiler changes to work correctly so not auto computed by default, must be manually passed.
  • Currently, Configurable via qaic_config (explicit or auto). Blocking is disabled by default but can be enabled by passing "enable_blocking": True via qaic_config
  • Created a generic interface for attention that can be called across different models for easier extension of attention methods in the future.
  • Supported Model Architectures for Blocking:
    • Llama3
    • Gpt-oss
    • Gemma
    • Gemma2
    • Granite
    • GraniteMOE
    • Mistral
    • Mixtral
    • MPT
    • Qwen2
    • Qwen3 (text only)
    • Qwen3MOE (text only)
    • Qwen3.5VL
    • Qwen3.5VL MOE
    • Starcoder

WIP

  • Implement the Auto Blocking
  • Unit tests

@vbaddi
Copy link
Copy Markdown
Contributor

vbaddi commented Apr 6, 2026

@kdulla pls fix the DCO, thanks

@vbaddi vbaddi mentioned this pull request Apr 6, 2026
Open
3 tasks
@kdulla kdulla force-pushed the enable_causal_blocking branch from e6fdc1f to 033550e Compare April 6, 2026 11:36
vbaddi and others added 27 commits April 8, 2026 08:46
@vbaddi @kdulla
feat: Add KV blocking attention optimization for Causal LMs
92fbee3 
- Add strategy registry and AttentionBlockingConfig for extensible blocking
- Implement BlockedKVAttentionTransform for supported attention modules
- Add auto-blocking policy with device and model specific params.
- Integrate KV blocking into Llama-like models using QEffDynamicCache

Key components:
* attention_blocking.py: Strategy registry and config
* attention_blocking_policy.py: Auto-derive policy
* blocked_attention_utils.py: KV blocked attention kernels
* pytorch_transforms.py: Module-level blocking application

Usage:
  qaic_config = {"num_kv_blocks": 2} # explicit
  qaic_config = {"attn_blocking_auto": True}  # automatic

Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@vbaddi @kdulla
nit: fix reducesum for kv blocking w/sub and hash params to handle qa…
25450eb 
…ic_config

Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@vbaddi @kdulla
nit: rename BlockedKVAttentionTransform to KVBlockingAttentionTransform
9bb7c4d 
Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@vbaddi @kdulla
nit: add kv+q blocking generalize support for qwen3/qwen3_moe/mllama
56cfcf4 
Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@kdulla
Initial changes for moving transforms out of pretrained
406706f 
Signed-off-by: Kushal Dulla <quic_kdulla@quicinc.com>
@kdulla
fixed CB issue with skip futures and added additional support for dif…
dbb82c0 
…ferent kinds of blocking

Signed-off-by: Kushal Dulla <quic_kdulla@quicinc.com>
@kdulla
added additional support for gpt-oss and qwen2.5vl
64796d3 
Signed-off-by: Kushal Dulla <quic_kdulla@quicinc.com>
@kdulla
Added support for num kv blocks that do not exactly divide CL
b5b93b7 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Fixed blocking configurator
1538d48 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Fixing formatting
c345b68 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Further formatting corrections
6c142ab 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Test formatting error
ab6069f 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
restructured some file paths and added optional qaic_config override …
d911e25 
…in compile

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
minor formatting fix
17ce9f6 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Restructured blocking changes to have generic blocked attention inter…
0823a5b 
…face and calls to blocking directory

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Removed print statement
ccedb77 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
added in progress qwen2.5 vl support and mistral support
3fba5b4 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed minor formatting issues
a665a20 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixing generic blocked attention based on code reviews
66cba12 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
removed incomplete blocking code from untested models, added generic …
ca7dae9 
…blocked interface to further models.

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed formatting and a few incorrect merges
63a635a 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
additional merge errors
d14dd52 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
added batch blocking
e409baa 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed past key value merge error
c9fa160 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed blocking test config
bcd263a 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
added a fix for batch blocking not exporting correctly
6ff0bbe 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed non-dynamic kv blocking
795806e 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
Added Qwen3VL blocking support
cdd9823 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla kdulla force-pushed the enable_causal_blocking branch from d5dfdb8 to cdd9823 Compare April 8, 2026 10:53
@kdulla
Fixes to qwen3 automatic blocking
98ecfd1 
Signed-off-by: Kushal Dulla <quic_kdulla@quicinc.com>
@kdulla kdulla force-pushed the enable_causal_blocking branch from 1de9644 to f8e61d5 Compare April 9, 2026 09:25
@kdulla
added missing init file
aa328f9 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla kdulla force-pushed the enable_causal_blocking branch from f8e61d5 to aa328f9 Compare April 9, 2026 09:30
@vishwasdivakar
Copy link
Copy Markdown

vishwasdivakar commented Apr 10, 2026

@vbaddi, PR908, with the latest commit fails to export Qwen2.5-VL-32B (same export script used as PR774)

commit- aa328f9

kdulla added 2 commits April 10, 2026 06:33
@kdulla
fix for qwen2.5vl
0f673cb 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
@kdulla
fixed mixtral modeling and removed unnecessary transform if we alread…
96275b3 
…y have onnx

Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@kdulla @vbaddi @vishwasdivakar