From e50159546d48d619d63ff33dea3613920181d75f Mon Sep 17 00:00:00 2001 From: Kushal Dulla Date: Tue, 28 Jul 2026 13:48:55 +0530 Subject: [PATCH] Updated reshape before CtxChunkScatterBatch to use static NH Signed-off-by: Kushal Dulla --- .../blocking/blocked_attention_forwards.py | 122 ++++++++++++++++++ QEfficient/transformers/cache_utils.py | 5 +- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 2 +- 3 files changed, 125 insertions(+), 4 deletions(-) diff --git a/QEfficient/blocking/blocked_attention_forwards.py b/QEfficient/blocking/blocked_attention_forwards.py index 3b2164b9b..4916f0572 100644 --- a/QEfficient/blocking/blocked_attention_forwards.py +++ b/QEfficient/blocking/blocked_attention_forwards.py @@ -248,6 +248,128 @@ def blocked_kv_attention_forward( def blocked_kv_attention_forward_decode_headpar_batch( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + num_kv_blocks: int, + cache_kwargs: Dict[str, Any], + layer_idx: int, + past_key_value: Cache, + ctx_len: int, + *, + use_causal_mask: bool = False, + sliding_window: Optional[int] = None, + skip_kv: bool = False, + position_bias: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + **kwargs, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Batch-folded decode: k_cache/v_cache [1, BH, T, D] where BH=B*num_kv_heads is static. + Used when B > 1 (decode, non-chunk_kv) — one core/device per (batch, kv-head) + pair, so no inner `split` dimension is needed (unlike blocked_kv_attention_forward_headpar_offline). + query: [B, NQH, 1, D], k_cache/v_cache: [1, BH, T, D], position_ids: [B, 1] + """ + batch_size, num_heads, seq_len, head_dim = query.shape + assert seq_len == 1, "blocked_kv_attention_forward_decode_headpar_batch is decode-only (seq_len must be 1)." + num_kv_groups = getattr(module, "num_key_value_groups", None) + num_kv_heads = num_heads // num_kv_groups + BH = batch_size * num_kv_heads # static at compile time + position_ids = cache_kwargs.get("position_ids") + num_kv_blocks = max(1, num_kv_blocks) + kv_block_size = -(-ctx_len // num_kv_blocks) + current_position = position_ids.max(dim=-1).values + + # Reshape query: [B, NQH, 1, D] -> [1, BH, num_kv_groups, D] + query_flat = query.reshape(batch_size, num_kv_heads, num_kv_groups, seq_len, head_dim).reshape( + 1, BH, num_kv_groups * seq_len, head_dim + ) + + max_blocks: list = [] + sum_blocks: list = [] + out_blocks: list = [] + + for j in range(num_kv_blocks): + start_index = j * kv_block_size + kv_len_block = (ctx_len - start_index) if j == num_kv_blocks - 1 else kv_block_size + end_index = start_index + kv_len_block + + skip_future = None + if skip_kv: + skip_future = (torch.tensor(start_index, device=query.device) > current_position).all() + if not torch.onnx.is_in_onnx_export() and not torch.jit.is_tracing(): + if skip_future.item(): + break + + # Read K: [B, num_kv_heads, T_block, D] -> [1, BH, T_block, D] + k_block = past_key_value.read_only_blocked_K_batch(start_index, end_index, layer_idx, cache_kwargs) + + attn_weights_block = torch.matmul(query_flat, k_block.transpose(3, 2)) * scaling + + # Causal mask: [B, 1, T_block] -> [B, num_kv_heads, 1, T_block] -> [1, BH, 1, T_block] + causal_mask = ( + torch.arange(kv_len_block, device=query.device).view(1, 1, kv_len_block) + > (position_ids - start_index)[:, :, None] + ) + # [B, 1, T_block] -> [B, num_kv_heads, seq_len, T_block] -> [1, BH, seq_len, T_block] + causal_mask = ( + causal_mask.unsqueeze(1) + .expand(batch_size, num_kv_heads, seq_len, kv_len_block) + .reshape(1, BH, seq_len, kv_len_block) + ) + # [1, BH, seq_len, T_block] -> [1, BH, num_kv_groups*seq_len, T_block] + # (tile, not interleave — query_flat's m-axis is rep-major/seq-minor: m = r*seq_len + q_pos) + causal_mask = ( + causal_mask.unsqueeze(3) + .expand(1, BH, seq_len, num_kv_groups, kv_len_block) + .reshape(1, BH, seq_len * num_kv_groups, kv_len_block) + ) + + attn_weights_block = torch.where( + causal_mask, torch.full_like(attn_weights_block, float(MIN_MASKED_ATTENTION_VALUE)), attn_weights_block + ) + + # causal_mask = causal_mask.repeat(1, 1, num_kv_groups, 1) + # attn_weights_block = attn_weights_block.masked_fill(causal_mask, float(MIN_MASKED_ATTENTION_VALUE)) + + max_block = attn_weights_block.max(dim=-1).values + exp_block = torch.exp(attn_weights_block - max_block.unsqueeze(-1)) + if skip_kv and (torch.onnx.is_in_onnx_export() or torch.jit.is_tracing()): + max_block = torch.where( + skip_future, torch.full_like(max_block, float(MIN_MASKED_ATTENTION_VALUE)), max_block + ) + exp_block = torch.where(skip_future, torch.zeros_like(exp_block), exp_block) + + # Read V: [B, num_kv_heads, T_block, D] -> [1, BH, T_block, D] + v_block = past_key_value.read_only_blocked_V_batch(start_index, end_index, layer_idx, cache_kwargs) + # sum_block = exp_block.sum(dim=-1) + sum_block = torch.einsum("btdn->btd", exp_block) + out_block = torch.matmul(exp_block, v_block) + if skip_kv and (torch.onnx.is_in_onnx_export() or torch.jit.is_tracing()): + sum_block = torch.where(skip_future, torch.zeros_like(sum_block), sum_block) + out_block = torch.where(skip_future, torch.zeros_like(out_block), out_block) + max_blocks.append(max_block) + sum_blocks.append(sum_block) + out_blocks.append(out_block) + + max_stacked = torch.stack(max_blocks) + sum_stacked = torch.stack(sum_blocks) + out_stacked = torch.stack(out_blocks) + block_max = max_stacked.max(dim=0).values + block_weight = torch.exp(max_stacked - block_max.unsqueeze(0)) + block_sum = torch.einsum("nbtd->btd", (block_weight * sum_stacked)) + # block_sum = (block_weight * sum_stacked).sum(dim=0) + block_out = torch.einsum("nbtdk->btdk", (block_weight.unsqueeze(-1) * out_stacked)) + # block_out = (block_weight.unsqueeze(-1) * out_stacked).sum(dim=0) + output = block_out / block_sum.unsqueeze(-1) # [1, BH, num_kv_groups*seq_len, D] + return output.reshape(batch_size, num_kv_heads, num_kv_groups, seq_len, head_dim).reshape( + batch_size, num_heads, seq_len, head_dim + ), None + + +def blocked_kv_attention_forward_decode_headpar_batch_unroll( module: nn.Module, query: torch.Tensor, key: torch.Tensor, diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index b6ddf4f33..5bda05de1 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -494,9 +494,8 @@ def write_only_batch(self, key_states, value_states, cache_kwargs): self.values = self.values.reshape(1, BH, self.keys.shape[2], D) self._mark_initialized(self.keys) position_ids = cache_kwargs.get("position_ids") - NKV = BH / position_ids.shape[0] - pos_folded = position_ids.unsqueeze(1).repeat(1, NKV, 1) - pos_folded = pos_folded.reshape(1, BH, -1) + NKV = (BH / position_ids.shape[0]).int().item() + pos_folded = position_ids.unsqueeze(1).repeat(1, NKV, QL).reshape(1, BH, QL) key_folded = key_states.reshape(1, BH, -1, D) value_folded = value_states.reshape(1, BH, -1, D) diff --git a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index a87758b7e..f21bc268e 100644 --- a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -1272,7 +1272,7 @@ def get_dummy_inputs( ) if batch_fold: - kv_cache_shape = [1, kv_cache_shape[0] * kv_cache_shape[1], *kv_cache_shape[2:]] + kv_cache_shape = [1, kv_cache_shape[0] * kv_cache_shape[1], kv_cache_shape[2], kv_cache_shape[3]] lang_inputs["past_key_values"] = [[] for _ in range(self.model.config.text_config.num_hidden_layers)] for i in range(self.model.config.text_config.num_hidden_layers):