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:19:58Z

+        batch, seq_len = position_ids.shape
+        num_kv_blocks, num_kv_heads, block_size, dh = k_out.shape
+        ctx_indices = torch.arange(block_size)[None, ...]
+        block_fill_len = position_ids.max(1, keepdim=True).values % block_size


Why don't you use slot_id here?
position_id.max % block_size should be same as slot_id, isnt it?
write** function is operating on slot_ids

This was switched from using (slot_id + seq_len) to (position_ids.max(1, keepdim=True).values % block_size) because of PyTorch and ONNX runtime generating NaN logits due to padding tokens even though it generated correct logits with QAIC compiled QPC run.

While reading blocks, we have 3 scenarios:

Block is invalid => do no read this block => invalid index assigned => gather_limit is 0

Block is fully valid => read the full block => gather_limit is kv_block_size

Block is partially valid => only some entries/tokens written in current block => read the block till valid tokens only => gather_limit = valid tokens length.

Since VLLM assigns slot_id at the start of prefill/decode run, so at the start of prefill slot_id is 0. If we write/scatter CPL = 4 tokens in a block with kv_block_size = 8, then we must read/gather only 4 tokens which can be achieved using (slot_id + seq_len) for recently "updated" block. Now, for QAIC runs padding tokens are not counted in seq_len while for PyTorch/ ONNX runtime runs padding tokens are part of seq_len. This leads to reading more tokens than needed for PyTorch and ONNX runtime cases and ends up with NaN logits. Using (position_ids.max(1, keepdim=True).values % block_size) for gather limit avoids this issue for all cases.

anujgupt-github · 2026-07-24T06:23:53Z

+            block_id = block_table[rows, block_index].unsqueeze(-1)
+            ctx_indices = torch.where(block_id < 0, invalid_scatter_index, ctx_indices)
+            block_id = block_id.unsqueeze(-1)
+            self.keys = CtxScatterFuncPagedAttention.apply(self.keys, block_id, ctx_indices, key_states)


On a CB decode step, key_states have shape [1, num_kv_heads, 1, dh] (single request) and block_table is [full_batch_size, num_kv_blocks]. The code therefore:

Reads block_table[torch.arange(1), block_index] — only row 0 of the block table, regardless of which request slot is active.
Scatters into self.keys at block_id = block_table[0, position_ids.max // block_size]
This means all continuously‑batched requests write into request‑0’s KV blocks.
Compare with mainline _QEffDynamicLayer.update() and _QEffDynamicLayer.write_only() in the same file. Both explicitly consume batch_index for the CB path. Paged write doesn't seem to do this

Block_table has shape [batch_size, num_kv_blocks] => if key_states have shape [1, num_kv_heads, 1, dh] then block_table will have shape [1, num_kv_blocks]. Hence, in this case we need to read K/V cache from physical location pointed to by block_index'th entry in the block_table and block table has only 1 row. The accuracy of the block_table is ensured by VLLM.

In CB case, each batch_index does not have separate allocated KV cache space. E.g. if 2 requests in the continuous batch share the same prefix, then only one of them will write to cache (or both of them can write to same location with same data, this is VLLM level optimization) and VLLM will ensure that KV cache for both those requests be read from the same physical KV blocks.

Let's take an example of num_kv_blocks=4 and BS=2. Let's say prefill wrote to 1 block only, then in decode case the block_table will look like [[0,1,-1,-1], [2,3,-1,-1]] in case no prefix is shared where -1 denotes invalid block which has not been written to yet. If for the same case, we share a prefix, the block table will look like [[0,1,-1,-1], [0,2,-1,-1]]. So, there might be cases where we read the same block for 2 different requests or different blocks and all this is managed by VLLM before the block_table for the current request is sent.

To tie this back to your question, if batch_size is 1 and num_kv_blocks is 4, then for first request the block_table can be [0,1,2,17] and for second request the block_table can be [4,15,22,80]. With block_index 0, we will read KV cache at physical block 0 for request 1 and physical block 4 for request 2 even though we read row 0 for both. Similarly, scatter will happen at physical block pointed to by (position_ids.max // kv_block_size)th entry in the block_table for that particular request.

The key idea here is that block_table is not fixed and it updates with input_ids and position_ids updates.

anujgupt-github · 2026-07-24T06:27:04Z

-    return compilation_batch_size, compilation_ctx_len, compilation_fbs
+    if compilation_num_kv_blocks := spec.get("num_kv_blocks", None):
+        compilation_num_kv_blocks = int(compilation_num_kv_blocks)
+    return compilation_batch_size, compilation_ctx_len, compilation_fbs, compilation_num_kv_blocks


Every caller in the codebase that unpacks 3 values will break with this change. You are updating only two call sites: cloud_ai_100_exec_kv and kv_offload_generate.
The remaining get_compilation_dims callers on the vision‑language single‑QPC generate paths and in some tests were not updated

Thanks for catching this. I have updated QEfficient/generation/vlm_generation.py and QEfficient/transformers/models/modeling_auto.py files. The following file was unpacking only 2 values, have updated it as well: QEfficient/transformers/models/modeling_auto.py

I was expecting to catch such issues in tests.

anujgupt-github · 2026-07-24T06:28:22Z

        compilation_fbs = int(compilation_fbs)
-    return compilation_batch_size, compilation_ctx_len, compilation_fbs
+    if compilation_num_kv_blocks := spec.get("num_kv_blocks", None):
+        compilation_num_kv_blocks = int(compilation_num_kv_blocks)


what happens without blocking?

Without kv_blocking, num_kv_blocks gets None value and none of the downstream "if num_kv_blocks ..." branches are triggered.

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.

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.

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


 def _resolve_effective_blocking_mode(attention_cfg: Dict[str, Any], requested_mode: str) -> str:
    mode = _normalize_attention_mode(requested_mode)
-    if mode == "":


Why did you remove this?
If user didn't specify blocking, why should below code get executed?

This has been corrected in the latest commit. Details in the comment above.

anujgupt-github · 2026-07-24T06:39:34Z

@@ -208,6 +208,121 @@ def blocked_kv_attention_forward(
    return attn_output, attn_weights




Six new attention forward functions are copy‑paste of the corresponding non‑paged ones.
The four paged functions differ only in the KV‑block fetch step (contiguous slice vs paged gather). Core functionality is identical. Instead of doubling the code, share a helper _fetch_kv_block(cache, layer_idx, j, num_kv_blocks, cache_kwargs, paged: bool) and keep one set of forwards

This was done to keep the paged and non-paged paths completely isolated from each other while developing the PagedAttention implementation for the first time and also to allow different attention calculation optimizations for each path, if needed. This separation was also done to follow QEfficient infrastructure style where different forwards are picked for different blocking methods.

Once PagedAttention is reviewed, I can share helper functions to keep one set of forward.

It would help if you can update this now. Will be easier to review

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

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.

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:49:44Z

+    return past_key_value is not None and hasattr(past_key_value, "read_only_pagedAttention")
+
+
 _STRATEGIES: Dict[BlockingMode, Callable] = {


Can create a second dictionary _STRATEGIES_PAGED that is checked if AttentionBlockingConfig has paged_attention flag True

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

    head_block_size: Optional[int] = None
    skip_kv: Optional[bool] = True
    num_batch_blocks: Optional[int] = None



kdulla · 2026-07-24T04:51:35Z

                comp_ctx_lengths=comp_ctx_lengths,
                batch_index=batch_index,
                position_ids=position_ids,
+                block_table=block_table,


Is this needed if we are not doing paged attention in this else condition?

This is needed to propagate the block_table and slot_id values form ONNX inputs down to attention forwards. In the else condition, these will just take the None values. In my opinion, the placeholders are still needed in the case where PagedAttention is supported. I would love to know if any alternative can work.

kdulla · 2026-07-24T04:51:59Z

                sliding_window=sliding_window,
            )

    strategy = _STRATEGIES.get(blocking_config.mode)


Can change this to check different strategies dictionary based on paged_attention flag or not

kdulla · 2026-07-24T04:54:08Z

 def _normalize_attention_mode(raw_mode: str) -> str:
    mode = raw_mode.lower()
-    if "h" in mode and "q" in mode and "kv" in mode:
+    if "h" in mode and "q" in mode and "kv" in mode and "paged" in mode:


See comments in attention_blocking, if we change paged_attention to a flag instead of additional BlockingModes we can avoid extra if conditions here

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 · 2026-07-24T04:59:02Z

    comp_ctx_lengths: Optional[torch.LongTensor] = None,
    batch_index: Optional[torch.LongTensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
+    block_table: Optional[torch.LongTensor] = None,


We only ever call this function when use_paged_kv_blocked is disabled, is this needed?

Updated the function to remove the unwanted args. Thanks for the good catch.

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

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?

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.

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?

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

+    """
+
+    @staticmethod
+    def forward(data: torch.Tensor, block_index: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor):


Is this in-place scatter safe if the same cache tensor is reused across eager parity checks?

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.

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?

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?

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

    return past_key_value is not None and hasattr(past_key_value, "read_only_blockedKV")


+def supports_paged_attention_blocked_kv(past_key_value: Optional[Cache]) -> bool:


Should paged behavior be selected by hasattr on the cache object rather than an explicit config signal?

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?

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.

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?

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

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

quic-hemagnih · 2026-07-27T08:52:44Z

        return self.layers[layer_idx].read_only_blockedKV(start_index, end_index, cache_kwargs)

+    def read_only_pagedAttention(self, block_index, updated, layer_idx, cache_kwargs):
+        # def read_only_pagedAttention(self, start_index, end_index, layer_idx, cache_kwargs):


remove this line

quic-rishinr · 2026-07-24T09:57:54Z

        self.full_batch_size = (
            full_batch_size if full_batch_size else self._fetch_full_batch_size()
        )  # Check and fetch full batch size if CB is enabled
+        self.num_kv_blocks = num_kv_blocks if num_kv_blocks else None


nit: since the default value is None we need not require the if condition here

quic-rishinr · 2026-07-30T09:27:20Z


+        if block_table is not None:
+            inputs["block_table"] = block_table
+            inputs["slot_id"] = np.zeros((prefill_logit_bs), dtype=np.int64)


Slot_id is initialized and never updated. Shouldnt we update the slot ID after each chunk similar to input ids and position ids in loop 851 to 875?

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?

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

…ormat cleanup Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

…n non-paged path Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

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.

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

quic-hemagnih requested changes Jul 27, 2026

View reviewed changes

quic-rishinr reviewed Jul 30, 2026

View reviewed changes

vaibverm added 8 commits July 30, 2026 15:15

Rebased PagedAttention support with latest Qeff for PR

c1b1140

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Added block_table and slot_id inputs + minor modelling_auto.py changes

2b0469c

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Working version with PagedAttention

0f054a6

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Minor fixes to specialization builder

0377888

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Added support for Qwen2.5_VL PagedAttention

caa8c8e

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

slot_id fix for Qwen2.5_VL PagedAttention decode

eb49d86

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Added support for Qwen3_VL PagedAttention

1a66341

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Removed commented code corrected in rebase

927b4cb

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

vaibverm added 10 commits July 30, 2026 15:15

Minor fix for enum bug

89c66d3

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Adding PagedAttetion support for Qwen3_VL_MOE

8d57575

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Adding PagedAttention specific test, unit tests and examples + lint/f…

20faa3c

…ormat cleanup Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Lint/format fixes

1f630f3

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Rebasing with latest Qeff mainline

ab81559

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Fix for get_compiler_dims returning 4 outputs

51e3cff

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Fix the faulty blocking modes resolution in blocking_configurator.py

7180f6e

Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Remove block_table and slot_id args from cache update function used i…

2436168

…n non-paged path Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Remove block_table and slot_id args from cache update function used i…

3f1c6d2

…n non-paged path Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

Remove block_table and slot_id args from cache update function used i…

c41035f

…n non-paged path Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>

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

		@@ -208,6 +208,121 @@ def blocked_kv_attention_forward(
		return attn_output, attn_weights

		return past_key_value is not None and hasattr(past_key_value, "read_only_pagedAttention")


		_STRATEGIES: Dict[BlockingMode, Callable] = {

		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,

		return past_key_value is not None and hasattr(past_key_value, "read_only_blockedKV")


		def supports_paged_attention_blocked_kv(past_key_value: Optional[Cache]) -> bool:

Uh oh!

Conversation

vaibverm commented Jul 23, 2026

Summary of changes compared to BlockedKV:

Uh oh!

anujgupt-github commented Jul 24, 2026

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!

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!

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!

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!

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

Uh oh!

vaibverm Jul 24, 2026 •

edited

Loading

vaibverm Jul 30, 2026 •

edited

Loading

anujgupt-github Jul 24, 2026 •

edited

Loading