Skip to content

Fix GPT-OSS blocked prefill cache reads and head-par prefill blocking - #1302

Draft
smedhe wants to merge 1 commit into
quic:mainfrom
smedhe:gptoss_blocked_headpar
Draft

Fix GPT-OSS blocked prefill cache reads and head-par prefill blocking#1302
smedhe wants to merge 1 commit into
quic:mainfrom
smedhe:gptoss_blocked_headpar

Conversation

@smedhe

@smedhe smedhe commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Fix GPT-OSS blocked cache read helpers to use the owning layer cache tensors directly.
  • Avoid duplicate cache updates during blocked prefill for non-sliding GPT-OSS attention.
  • Default n_rep_chunk to 1 in attention blocking config.
  • Update the GPT-OSS blocked prefill head-par example to handle logits shaped [B, V], [B, 1, V], or [B, S, V].
  • Pass MoE prefill packed chunk sizing through the attention blocking config.
  • Add unit coverage for blocked K/V cache read dispatch.

Signed-off-by: smedhe <smedhe@qti.qualcomm.com>
@smedhe
smedhe marked this pull request as draft September 2, 2026 06:48
"num_kv_blocks": 2,
"num_q_blocks": 2,
"ctx_len": args.ctx_len,
"moe_config": {"expert_prefill_chunk_size": args.moe_prefill_packed_chunk_size},

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be expert_parallel_chunk_size

query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)

blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig())
use_blocking = blocking_config is not None and blocking_config.mode.is_prefill and (self.sliding_window is None)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should raise an error if we do a blocking mode that is not prefill specific, not just skip it like this code does.

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.

2 participants