Adding PagedAttention support for CausalLM models by vaibverm · Pull Request #1209 · quic/efficient-transformers

vaibverm · 2026-07-23T01:59:07Z

This PR is a clone of #982. This PR adds the PagedAttention (https://arxiv.org/pdf/2309.06180) support for all CausalLM models in QEfficient.
The major change is that KV cache is not treated as a contiguous memory under this implementation but rather a collection of blocks which can reside in a non-contiguous fashion inside the memory. This forces cache scatter and gather operations to happen per KV block.

Summary of changes compared to BlockedKV:

The cache shape changes from [BS, num_kv_heads, CL, dh] to [total_num_kv_blocks, num_kv_heads, kv_block_size, dh].
num_kv_blocks = -(-ctx_len // kv_block_size) = physical blocks required for 1 batch element in K cache.
Total_num_kv_blocks = BS (<kv_batch_size>) * num_kv_blocks = total physical blocks available for K cache.
2 new inputs block_table [BS, num_kv_blocks] and slot_id [BS] are passed as inputs to the ONNX.
4) a) block_id is each entry in the block_table and points to the physical K/V block that needs to be read/written corresponding to (position_id // kv_block_size)th entry in block_table. ‘-1’ signifies invalid/unallocated block.
4) b) slot_id denotes how many entries are already filled in currently active block => read up to / write after (slot_id – 1)
Limitation - Cache writes to only 1 block at a time per batch element => CPL is less than or equal to kv_block_size. Hence, cache writes should not cross the block boundary.
vLLM provides KV Cache Manager implementation which maintains the KV cache block_table with logical to physical block mapping and slot_id for location mapping within the active block.

anujgupt-github · 2026-07-24T06:46:03Z

"Limitation - Cache writes to only 1 block at a time per batch element => CPL is less than or equal to kv_block_size. Hence, cache writes should not cross the block boundary."

Where is this enforced in the code?

anujgupt-github · 2026-07-24T06:31:06Z

-        fbs: int = constants.ONNX_EXPORT_EXAMPLE_FBS
+            num_kv_blocks = self.hash_params["blocking_kwargs"].num_kv_blocks
+            self.supports_paged_attention = "paged" in self.hash_params["blocking_kwargs"].mode
+            seq_len = kv_block_size = -(-seq_len // num_kv_blocks) if self.supports_paged_attention else seq_len


if paged attention is false, we don't want block_size to be seq_len

kv_block_Size is not used anywhere if self.supports_paged_attention is False. Hence the fallback to default of kv_block_size = seq would not effect any non-pagedAttention path here. That fallback to default is to ensure seq_len = seq_len if self.supports_paged_attention is False.

Please let me know if you see any additional issues with this.

Needs confirmation for no remaining changes

anujgupt-github · 2026-07-24T06:37:13Z

+        return mode
    if num_q_blocks > 1:
-        return "q"
+        return mode


_normalize_attention_mode can return kv_paged.
if num_q_blocks > 1 and num_kv_blocks == 1, this still returns "kv_paged"
Is that ok?

Same as above. Corrected.

Marked resolved

anujgupt-github · 2026-07-24T06:42:50Z

+                "block_table": block_table,
+                "slot_id": slot_id,
+            }
+            if sliding_window is not None:


Why separate handling? are you assuming that block size will always be > SWA size?

In the current implementation, PagedAttention hasn't been evaluated with sliding window attention yet. The reason is that in sliding window attention, different layers will require different max num_kv_blocks and how VLLM handles this heterogeneity hasn't been studied yet. This will be done in a future implementation.

Better to assert the path which is not supported instead of letting it fall through

Please add a NotImplemented exception in this case.

anujgupt-github · 2026-07-24T06:44:09Z

+        if hasattr(self.model, "model"):
+            self.model.model.qaic_config = qaic_config
+            if hasattr(self.model.model, "model"):
+                self.model.model.model.qaic_config = qaic_config


Prefer passing qaic_config explicitly through constructors, or attach it once on the top wrapper and always read via a helper that walks up.

this code is failing in the unit test. please check

kdulla

The approach of adding new blocking modes adds a lot of repeated unnecessary code, changing that should help streamline many parts of this

kdulla · 2026-07-24T04:49:10Z

 class BlockingMode(str, Enum):
    NONE = ""
    KV = "kv"
+    KV_PAGED = "kv_paged"


Would be better to not changed BlockingMode and instead add a paged_attention_flag to AttentionBlockingConfig

I was trying to stay consistent with the current QEfficient blocking infrastructure where various blocking modes each have a separate forward and a separate mode. But I agree, we can probably introduce another qaic_config flag like "supports_pagedAttention" to switch between the 2 KV blocking modes. That would be an added consolidation of the existing code but wouldn't change functionality.

better to make the paged as a bool on the blocking config at the KV read/write step. that way we can remove a lot of duplicate code.

kdulla · 2026-07-24T04:56:00Z

+        ):
            mode_from_config = "kv" + mode_from_config
            blocking_config.num_kv_blocks = _get_valid_num_blocks(qaic_config, "num_kv_blocks")
+        if (


Only kv_paged attention is possible from qaic_config from this code

I am not able to fully understand the reason for this observation. We have both "paged" in blocking_mode and "paged" not in blocking_mode if statements. Why do you think only kv_paged is possible?

@kdulla Can you clarify?

kdulla · 2026-07-24T05:37:14Z

+                bs, num_kv_blocks
+            )
+            example_inputs["slot_id"] = torch.zeros(bs, dtype=torch.int64)
+            dynamic_axes["block_table"] = {0: "batch_size", 1: "num_kv_blocks"}


Num kv blocks can't be a proper dynamic axis since it is fixed at export time based on the number of loops the attention forward goes through

The generic CausalLM path was fixed, but Qwen2.5-VL, Qwen3-VL, and Qwen3-VL-MoE still declare num_kv_blocks dynamic. Dual QPC VLM prefill keeps slot_id at zero for all chunks.

Good catch — 41beb26 only fixed the generic CausalLM export path; qwen2_5_vl.py, qwen3_vl.py, and qwen3_vl_moe.py still declared {0: "batch_size", 1: "num_kv_blocks"} on block_table. Fixed in 3496aa1 — mirrored the same one-line change (drop the 1: "num_kv_blocks" entry, since the block count is fixed at export/trace time, not a real ONNX dynamic axis) into all three Qwen files. Pushed to PR_branch.

On the VLM slot_id bug (never built in _execute_chunked_prefill): slot_id was never populated at all for any of the 3 VLM prefill entry points in vlm_generation.py, same root mechanism as the CausalLM bug you flagged on text_generation_inference.py:820. Fixed in 4f4ce2e: slot_id is now recomputed per prefill chunk instead of staying frozen/unset before the chunk loop.

While fixing that, I noticed VisionLanguageGeneration had no way to actually populate block_table/num_kv_blocks at all — the dual-QPC constructor never wired it through, so paged attention on the language QPC was unreachable at runtime even though the export-side plumbing already supported it for qwen2_5_vl/qwen3_vl/qwen3_vl_moe. Added that wiring in 425c8e8: num_kv_blocks is now an optional constructor param (mirrors the non-VLM TextGeneration path), and self.block_table stays None/inert when it's not supplied — no behavior change for existing non-paged VLM callers.

anujgupt-github · 2026-07-24T08:06:43Z

    return ""


 def _resolve_effective_blocking_mode(attention_cfg: Dict[str, Any], requested_mode: str) -> str:


When paged attention is not requested, why should blocking mode resolution differ from mainline?

Resolved, blocking mode to match mainline

anujgupt-github · 2026-07-24T09:06:07Z

        v_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out)
        return k_out, v_out

+    def read_only_pagedAttention(self, block_index, updated, cache_kwargs):


Do we want mixed snake_case/camelCase cache APIs like read_only_pagedAttention?

I will leave the final decision to the QEff team regarding the case to use. My reasoning was that PagedAttention is a single word in the original paper: https://arxiv.org/pdf/2309.06180, so probably shouldn't be broken into 2 words like paged_attention. But please let me know if we need a clean snake_case for names involving PagedAttention and I can do that.

We will stick with the snake_case. Please update.

anujgupt-github · 2026-07-24T09:06:09Z

+            self.keys = CtxScatterFuncPagedAttention.apply(self.keys, block_id, ctx_indices, key_states)
+            self.values = CtxScatterFuncPagedAttention.apply(self.values, block_id, ctx_indices, value_states)
+
+    def get_seq_lengthPagedAttention(self, cache_position=None) -> int:


Should this return allocated paged-cache capacity or the active sequence length?

Today it returns allocated capacity: self.keys.shape[-2] * self.keys.shape[0] (block size × number of allocated blocks) — fixed at allocation time, doesn't shrink/grow with how many tokens have actually been written for a given request. Right now, this method has zero callers anywhere outside its own definition. Do you see a use case for returning active sequence length? Would be glad to receive further direction on this.

Removed this function as it was unused.

anujgupt-github · 2026-07-24T09:06:12Z

+        num_kv_blocks = self._get_num_kv_blocks()
+        kv_block_size = self._get_kv_block_size()
+
+        if self._is_paged_attention and num_kv_blocks and kv_block_size:


Is this intended to check the _is_paged_attention method object instead of calling it?

it should be self._is_paged_attention()

anujgupt-github · 2026-07-24T09:06:13Z

@@ -112,6 +112,8 @@ def forward(
        hidden_states: torch.Tensor,


Do we want every model wrapper signature to carry block_table and slot_id when most paths won't use paged attention?

Only the causalLM and ImageTextToText models which support KV Blocking should carry block_table and slot_id. Non causalLM models which do not use KV cache should not need to carry these inputs.

Need further clarification on how to proceed/ to resolve this comment.

anujgupt-github · 2026-07-24T09:06:14Z

        seq_len: int = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN

-        # increase seq_len if using a larger number of blocks
+        self.supports_paged_attention = False


Should export input construction mutate self.supports_paged_attention as model state?

Change self.paged_attention to a local variable? going through Blocking_kwargs as the single source of truth?

anujgupt-github · 2026-07-24T09:06:16Z

+        # increase seq_len if using a larger number of blocks and set PagedAttention params if required
        if self.hash_params.get("blocking_kwargs", None):
            max_blocks = -1
            for num_blocks in self.hash_params.get("blocking_kwargs").__dict__.values():


Do we want max_blocks derived from every int in blocking_kwargs.dict, including unrelated fields?

vaibverm · 2026-07-24T19:39:44Z

"Limitation - Cache writes to only 1 block at a time per batch element => CPL is less than or equal to kv_block_size. Hence, cache writes should not cross the block boundary."

Where is this enforced in the code?

This is not enforced in code right now since this is not a hard limitation. This limitation was introduced for current implementation based on the decision that CPL and kv_block_size should have same granularity. This is a "To do" for future implementation and probably the only one.

If needed, a warning/status message can be added to alert the user of this limitation when blocking_mode="*kv_paged" and CPL > kv_block_size.

anujgupt-github · 2026-07-27T08:45:44Z

Can you add tests for the cases below?

get_compilation_dims() should still support existing callers that unpack 3 values by default.
num_kv_blocks should be covered through an explicit opt-in path, not by changing the default return shape.
Non-paged generation should behave unchanged when num_kv_blocks is absent.
CPL > kv_block_size should not export into a wrong-output path.
A paged cache write where slot_id + seq_len > kv_block_size should be covered.
CB should be covered with two active slots using different physical KV blocks in block_table.
CB prefix sharing should be covered where two slots share some physical blocks and then diverge.
Non-paged blocking mode resolution should be covered for kv, qkv, hq, hkv, and hqkv.
Paged attention disabled + blocking enabled should be covered to prove old blocked attention behavior is unchanged.
Unsupported paged configs should be covered so we fail clearly instead of exporting a graph that can generate wrong output.

quic-hemagnih

Add unit test cases for read_only_pagedAttention / write_only_pagedAttention on QEffLayerCache

Suggested Optimizations to reduce unnecessary loops-

read_only_pagedAttention() is being called redundantly inside the head/q/batch loops, even though its inputs only depend on KV block index j.

For each j:

block_index = block_table[:, j]
updated = (position_ids.max(...) // kv_block_size) == j
cache_kwargs is unchanged

These do not depend on head_block_idx, q_block_idx, or b_block_idx. As a result, the same expensive CtxGatherFuncPagedAttention / GatherND over the full KV cache is repeated multiple times, and only cheap head/batch slicing differs afterward.

For example:

blocked_hqkv: calls = num_head_blocks × num_q_blocks × num_kv_blocks, but only num_kv_blocks unique gathers are needed.
blocked_bhqkv: calls = num_head_blocks × num_q_blocks × num_batch_blocks × num_kv_blocks, but again only num_kv_blocks unique gathers are needed.

Suggest hoisting the gather outside the outer loops:

Gather k/v once per KV block j.
Cache the gathered k/v states.
Reuse them inside head/q/batch loops and apply only the required slicing there.

This should reduce redundant CtxGatherFuncPagedAttention calls significantly, e.g. from 64 to 8 in hqkv and from 128 to 8 in bhqkv for the given configs, with the tradeoff of holding gathered KV blocks in memory temporarily.

Same applies to blocked_bhqkv_paged_attention_forward also.

quic-hemagnih · 2026-07-27T08:22:41Z

+        num_kv_blocks = self._get_num_kv_blocks()
+        kv_block_size = self._get_kv_block_size()
+
+        if self._is_paged_attention and num_kv_blocks and kv_block_size:


Shouldn't it be self._is_paged_attention()? What's the intention here?

Marked resolved.

quic-hemagnih · 2026-07-27T08:25:41Z

+                "block_table": block_table,
+                "slot_id": slot_id,
+            }
+            if sliding_window is not None:


This peice of code will never execute, as at line#159, sliding _window should be none to enter in the code leg.
Now at this line you are checking that sliding window is not NONE.

Marked resolved.

quic-hemagnih · 2026-07-27T08:32:37Z

+        v_out = CtxGatherFuncPagedAttention.apply(v_out, block_indices, ctx_indices)
+
+        v_out = torch.where((invalid_mask.unsqueeze(1)).unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out)
+        return k_out, v_out


Any reasons why k_out is not masked here? If we don't mask k_out then it might set a wrong current_max in the online softmax accumulator, which corrupts the numerical normalization of all subsequent valid blocks, producing a near-zero output instead of the correct attention-weighted value?
What do you think?

Based on my analysis, k_out doesn't need masking because the code always masks the scores computed from it before those scores are used for anything. Any garbage slot in k_out is, by definition, a slot the sequence hasn't written to yet. The causal mask always blocks future positions, so that garbage score gets overwritten before it can affect current_max or the final output.
We can resolve this if you agree, or I can look into it if you have any suggestions.

Need further clarification on how to proceed/ to resolve this comment.

quic-hemagnih · 2026-07-27T08:34:18Z

+        v_out = torch.where((invalid_mask.unsqueeze(1)).unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out)
+        return k_out, v_out
+
+    def write_only_pagedAttention(self, key_states, value_states, cache_kwargs):


Align function naming conventions with existing code of QEFF

Resolved to match snake case

quic-hemagnih · 2026-07-27T08:39:04Z

+        use_causal_mask = True
+    position_ids = cache_kwargs.get("position_ids")
+    block_table = cache_kwargs.get("block_table")  # [BS, num_kv_blocks] -> each entry is block_id value
+    kv_block_size = past_key_value.get_seq_length() if past_key_value is not None else 0


Why don't we use get_seq_lengthPagedAttention() in place of get_seq_length()?

For paged attention, keys has shape [total_num_kv_blocks, num_kv_heads, kv_block_size, dh], so keys.shape[-2] is kv_block_size — which happens to be
correct. But this is fragile as based on the paged layout. The newly added get_seq_lengthPagedAttention() method exists precisely for this purpose.

This comment applies for other instances too

Checked this — get_seq_length() and get_seq_length_paged_attention() return different quantities. get_seq_length() returns self.keys.shape[-2], which for the paged layout is kv_block_size (per-block size). get_seq_length_paged_attention() returns self.keys.shape[-2] * self.keys.shape[0] — total capacity across all blocks.
get_seq_lengt() is used to get kv_block_size to size each block read — swapping in the paged-specific method would change that to total capacity and break block indexing. So current usage looks intentional rather than fragile.

Please let me know your thoughts/ further direction on this.

Removed get_seq_length_paged_attention function as it was unused.

quic-rishinr · 2026-07-30T09:37:40Z

                lang_inputs["input_ids"], dtype=lang_inputs["mm_token_type_ids"].dtype
            )
+        if num_kv_blocks:
+            lang_inputs["slot_id"] = (np.max(lang_inputs["position_ids"]) % kv_block_size).reshape(batch_size)


np.max(lang_inputs["position_ids"]) would fail if the bs > 1. this will return a scalar. You would need to fetch the max value for all batch ids.

quic-rishinr · 2026-07-30T09:42:01Z

+        num_kv_blocks = self._get_num_kv_blocks()
+        kv_block_size = self._get_kv_block_size()
+
+        if self._is_paged_attention and num_kv_blocks and kv_block_size:


it should be self._is_paged_attention()

quic-rishinr · 2026-07-30T09:47:54Z

 class BlockingMode(str, Enum):
    NONE = ""
    KV = "kv"
+    KV_PAGED = "kv_paged"


better to make the paged as a bool on the blocking config at the KV read/write step. that way we can remove a lot of duplicate code.

quic-rishinr · 2026-07-30T09:57:13Z

@vaibverm can you please rebase the PR against main?

vaibverm · 2026-07-30T20:19:19Z

@vaibverm can you please rebase the PR against main?

Have done the rebase. But the changes in cache_utils.py for paged gather and scatter functions need to be done in line with mainline changes for the dynamo path coming from QEfficient/customop/utils.py. I need to understand in more detail about dynamo related changes.

quic-rishinr · 2026-08-03T08:28:51Z

@vaibverm can you resolve the comments which you have already addressed? this will make it easier for reviewers while revieing the PR. Also please fix the lint issues.

quic-rishinr · 2026-08-10T04:48:12Z

ci_ready

Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

…hqkv Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

…ute 'qaic_config' Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

CtxScatterFunc/CtxScatterFuncPagedAttention cloned the cache tensor before scattering, purely to avoid mutating the caller's live buffer in eager mode. The clone is never part of the emitted ONNX: the ScatterND-based FunctionProto comes entirely from symbolic(), not from forward()'s eager body (verified by dumping the CtxScatter FunctionProto and diffing it byte-for-byte against the onnxscript definition -- no Clone/Identity node present either way). The clone's only effect during export is on the tracer's tensor-identity/aliasing bookkeeping feeding the CSE pass, which was causing layer 0 and layer 1's decoder calls to end up with mismatched signatures and emit as two separate FunctionProtos (QEffLlamaDecoderLayer / .1) instead of deduping into one shared function under use_onnx_subfunctions=True. Tested: - Standalone repro (tmp_debug/diff_subfunctions.py): 2-layer llama export under use_onnx_subfunctions=True now emits a single QEffLlamaDecoderLayer function instead of QEffLlamaDecoderLayer/.1. - Full suite: tests/transformers/subfunction/test_single_subfunction.py, 15/15 passed, covering gpt2, falcon, gptj, llama, mistral, mpt, phi3, qwen2, qwen3, granite, olmo2, qwen3_moe, gemma, gemma2, and the tinyllama single-decoder export smoke test. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

build_transformer_blocking_config's ValueError guard hardcoded an allow-list of ["kv", "qkv", "hqkv", "bhqkv"] for which blocking modes may combine with paged attention, omitting "hkv". This contradicts the rule an earlier commit (04ee17b) already documented in a comment: paged attention is valid for any mode containing KV-blocking, since paged attention reads the cache block-wise and that only makes sense when the cache itself is KV-blocked. The hardcoded list was an unintended side effect of a later mode-resolution refactor (b64aca6), not a deliberate restriction -- there's no PR comment, call note, or commit message anywhere calling for hkv_paged to be rejected. Replaced the allow-list with the general "kv" in effective_mode check that matches the pre-existing documented intent, restoring hkv_paged as a supported combination while still rejecting paged modes with no KV-blocking at all (h_paged, hq_paged, q_paged), which never had a working code path. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

write_only_paged_attention assumed a single write never spans more than one physical KV block, backed only by a comment ("Assuming only 1 block is written at max") with no actual check. If seq_len exceeded block_size, or slot_id's offset pushed a write past the block boundary, the resulting ctx_indices would run past block_size for a fixed-size physical block -- CtxScatterFuncPagedAttention would then scatter using those out-of-range indices, silently wrapping/clobbering data with no error. This raises NotImplementedError in both cases instead: seq_len > block_size (the write alone is too long for any block), and slot_id + seq_len > block_size (the write fits alone but the slot offset pushes it over). Tested: - tests/unit_test/models/test_cache_correctness.py:: TestPagedAttentionSingleBlockWriteLimit (added, 4 tests): write within block_size succeeds; seq_len > block_size raises; slot_id + seq_len > block_size raises; slot_id + seq_len == block_size (exact boundary) succeeds. All pass. - Full test_cache_correctness.py: 30 passed, no regressions. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

construct such a block_table by hand (kv_block_size=4, block_table=[[3,7],[3,12]], position_ids=[[4],[4]], slot_id=[0,0]) to verify the per-row lookup (block_table[rows, block_index]) handles it correctly rather than degenerating to row 0 for every row. Four tests: - test_shared_prefix_written_to_common_physical_block: both rows' writes land in the shared physical block their block_table row points to. - test_divergent_decode_writes_go_to_separate_physical_blocks: each row's decode-step write lands in its own physical block, without bleeding into the other row's block or the shared block; untouched blocks stay at their initial value. - test_read_back_gathers_shared_prefix_and_divergent_tail_per_row: read-back correctly returns the shared block for both rows and each row's own divergent block. - test_replacement_prefill_uses_request_specific_block_table_row: a replacement prefill (a new request reusing a decode slot after the previous occupant finished) keys off each request's own block_table row, not row 0's for every row. Tested: full tests/unit_test/models/test_cache_correctness.py, 34 passed. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

…blocking ValueError, non-paged blocking numeric parity, and CPL > kv_block_size guard Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

…ked attention forwards The gather's inputs (block_index, updated, cache_kwargs) depend only on the KV-block index, not on the outer head/q/batch loop indices, so gather once per KV block and reuse the cached k/v states for slicing. Adds regression tests asserting the gather call count, numeric correctness, and skip_kv early-break behavior for qkv/hqkv/bhqkv. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

…attention The existing skip_kv test only asserted the KV gather's call count after the early break, not that the surviving blocks still produce correct output. Add a companion test comparing against plain causal attention. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

blocked_bhqkv_attention_forward's non-paged numeric parity against the unblocked model was untested anywhere in the repo (predates this PR -- introduced in 410d21b). test_cpu_logits_allclose_original_vs_blocked already proves this for kv/qkv/hqkv; extend its parametrize list to close the same gap for bhqkv. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Mirrors the earlier fix for the generic CausalLM export path: block_table's KV-block dimension is fixed at cache-allocation time, not a genuinely varying ONNX export axis, so it must not be declared dynamic for Qwen2.5-VL, Qwen3-VL, or Qwen3-VL-MoE either Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

… the earlier CausalLM fix (668e237): slot_id must be recomputed per prefill chunk (chunk_start_position_id % kv_block_size), not left frozen from before the chunk loop, or later chunks silently corrupt earlier chunks' KV rows. Threads block_table through _execute_chunked_prefill and its 3 call sites via getattr(self, \"block_table\", None), a safe no-op today since VisionLanguageGeneration's constructor never sets that attribute yet. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Adds an optional num_kv_blocks constructor param to VisionLanguageGeneration, mirroring the non-VLM TextGeneration path. When not supplied, self.block_table stays None and paged attention is disabled, matching the existing opt-in convention. Also forwards num_kv_blocks_comp read from compilation dims in modeling_auto.py's generate() instead of discarding it. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

BlockingMode.resolve() does a strict enum lookup and has no _paged members (paged_attention is a separate boolean), so passing a raw mode string like "kv_paged" through to build_transformer_blocking_config_for_transform raised ValueError. Strip the paged suffix before resolving the base mode, while still passing the original string to build_transformer_blocking_config so its own paged detection still works. Also add the now-required ctx_len argument to direct blocked_*_attention_forward calls and AttentionBlockingConfig fixtures in the KV-block gather and blocking-transform parity tests. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

quic-rishinr · 2026-09-09T05:12:41Z

CI-Ready

quic-rishinr · 2026-09-09T05:17:19Z

+    layer_idx: int,
+    kv_block_size: int,
+    paged_attention: bool,
+    j: int,


nit : use proper variable naming for J something like kv_block_idx.

quic-rishinr · 2026-09-09T05:19:42Z

+    # Gather each KV block once: block_index/updated/cache_kwargs only depend on `j`,
+    # not on q_block_idx, so hoist the read out of the q-block loop below.
+    kv_blocks = []
+    for j in range(num_kv_blocks):


rename the variable name j to kv_block_idx

quic-rishinr · 2026-09-09T05:20:22Z

+    # Gather each KV block once: block_index/updated/cache_kwargs only depend on `j`,
+    # not on head_block_idx/q_block_idx, so hoist the read out of the loops below.
+    kv_blocks = []
+    for j in range(num_kv_blocks):


use the kv_block_idx as variable name instead of j

quic-rishinr · 2026-09-09T06:40:43Z

        num_devices=4,
    )

+    # head qkv_paged_attention blocking


we should not be running paged attention for all the models. this would add another 1 hour to the current CI time which is not acceptable, lets reduce the configs and keep the testing limited for few priority models.

quic-rishinr · 2026-09-09T07:36:48Z

@athavale-shivani please fix the lint and Unit test failures

quic-rishinr · 2026-09-09T09:40:14Z


    @staticmethod
    def forward(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor):
+        if not torch.onnx.is_in_onnx_export():


@athavale-shivani why do you want this change in regular scatter method? this could increse the latency even if the paged attention is disabled. can we remove this?

It was added per @anujgupt-github earlier review comment flagging that the in-place scatter here mutates data directly, so a caller holding a reference to the pre-scatter tensor (e.g. for an eager-parity comparison) would see it silently change underneath them. It was applied to both CtxScatterFunc.forward and CtxScatterFuncPagedAttention.forward to match the 3D scatter variants (CtxScatterFunc3D, CtxScatterFunc3DGeneralized, CtxScatterFunc3DInt), which already clone before scattering. If you do not want CtxScatterFunc to change, I can remove it.

ochougul · 2026-09-09T09:51:11Z

CI-Ready

qraniumcitest · 2026-09-09T17:26:03Z

CI-Ready

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

…ttention forward() and fused_forward() forwarded arguments to their blocking-mode variants positionally without including the newly-added block_table/slot_id params, shifting every trailing argument two slots and leaving cos_cached/ sin_cached as None. This broke rotary embedding application for every DeepSeek-V3/Kimi-K2.5 forward pass, caught by test_kimi_k25_quickcheck_hf_qeff_vision_logits_parity. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

…pSeek-V3 fused_forward_h_blocking and fused_forward_kv_blocking accepted block_table and slot_id as parameters but never passed them into their generic_blocked_attention_interface call, unlike forward_full_kv_h_blocking which already does. Currently a no-op (generic_blocked_attention_interface's paged-write branch is gated behind 'if not is_mla', so MLA+paged-attention isn't wired end-to-end yet), but threading the params through now keeps the signature consistent and avoids silently dropping them. Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

vaibverm mentioned this pull request Jul 23, 2026

Adding PagedAttention support for CausalLM models #982

Closed

vaibverm force-pushed the PR_branch branch 2 times, most recently from 2569107 to fb508e1 Compare July 23, 2026 09:42

anujgupt-github requested review from kdulla and ochougul July 24, 2026 06:47

anujgupt-github requested changes Jul 24, 2026

View reviewed changes

kdulla requested changes Jul 24, 2026

View reviewed changes

anujgupt-github reviewed Jul 24, 2026

View reviewed changes

Comment thread QEfficient/customop/ctx_scatter_gather.py

anujgupt-github reviewed Jul 24, 2026

View reviewed changes

Comment thread QEfficient/blocking/attention_blocking.py

quic-hemagnih requested changes Jul 27, 2026

View reviewed changes

quic-rishinr reviewed Jul 30, 2026

View reviewed changes

vaibverm force-pushed the PR_branch branch from 800ce61 to c41035f Compare July 30, 2026 20:16

ochougul marked this pull request as draft August 4, 2026 16:39

quic-rishinr marked this pull request as ready for review August 9, 2026 06:08

quic-rishinr force-pushed the PR_branch branch from c41035f to 4764594 Compare August 9, 2026 06:08

quic-rishinr added the ci_ready label Aug 10, 2026

aditjadh and others added 16 commits September 9, 2026 10:42

[Formatting]: Change function definition to snake case

f88d608

Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

Non-paged blocking mode resolution covered for kv, qkv, hq, hkv, and …

b8f98eb

…hqkv Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

[FIX]:AttributeError: 'QEffGemma3DecoderWrapper' object has no attrib…

8ea6965

…ute 'qaic_config' Signed-off-by: Aditya Dhananjay Jadhav <aditjadh@qti.qualcomm.com>

Add tests for get_compilation_dims backward compat, paged-without-KV-…

82de7ff

…blocking ValueError, non-paged blocking numeric parity, and CPL > kv_block_size guard Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Remove unused get_seq_length_paged_attention

997d837

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

quic-rishinr force-pushed the PR_branch branch from ad880bf to c5675e5 Compare September 9, 2026 05:12

quic-rishinr requested changes Sep 9, 2026

View reviewed changes

Aditya Dhananjay Jadhav and others added 6 commits September 9, 2026 11:50

[Formatting]: Rename j as kv_block_idx in _read_kv_block func.

b755a94

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

Unit Test:Replace HQKV blocking config by KV blocking

70f0476

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

Fix ruff formatting drift

81830a6

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>

Add unit tests for kv, qkv and hqkv paged attention

905e6a2

Signed-off-by: Aditya Dhananjay Jadhav <your-email@qualcomm.com>

aditjadh force-pushed the PR_branch branch from 3f6936e to 905e6a2 Compare September 11, 2026 19:32

		return ""


		def _resolve_effective_blocking_mode(attention_cfg: Dict[str, Any], requested_mode: str) -> str:

		@@ -112,6 +112,8 @@ def forward(
		hidden_states: torch.Tensor,

Uh oh!

Conversation

vaibverm commented Jul 23, 2026

Summary of changes compared to BlockedKV:

Uh oh!

anujgupt-github commented Jul 24, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vaibverm Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kdulla left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anujgupt-github Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

vaibverm Jul 24, 2026 •

edited

Loading

anujgupt-github Jul 24, 2026 •

edited

Loading