Skip to content

fix: limit step attention to active cache - #395

Open
ousamabenyounes wants to merge 1 commit into
OpenBMB:mainfrom
ousamabenyounes:fix/issue-293
Open

fix: limit step attention to active cache#395
ousamabenyounes wants to merge 1 commit into
OpenBMB:mainfrom
ousamabenyounes:fix/issue-293

Conversation

@ousamabenyounes

Copy link
Copy Markdown

Summary

  • limit step attention to the populated KV-cache prefix
  • remove the now-redundant boolean attention mask
  • retain contiguous tensors for the documented MPS compatibility path

Fixes #293.

Test verification (RED → GREEN)

The regression test intercepts the SDPA boundary, records the KV sequence
length and mask presence, and still calls the real PyTorch SDPA implementation.

With the new test on upstream/main (production fix absent):

FAILED tests/test_minicpm4_attention.py::test_forward_step_only_attends_to_populated_cache
1 failed, 2 warnings in 2.75s

With the production fix restored:

1 passed, 2 warnings in 2.74s

The production-only revert was repeated after implementation and reproduced
the same failure; restoring the patch returned the test to green.

Full local suite validation:

  • baseline suite: 56 passed
  • patched suite: 57 passed
  • Black: passed for both changed files
  • source distribution and wheel build: passed
  • changed-line coverage: 100% (2 executable lines, 0 missing)

No local CUDA benchmark was run or claimed.

Slice the static KV cache before SDPA so autoregressive steps only attend to populated positions and avoid a full-length boolean mask.

Generated by Ora Studio
Vibe coded by ousamabenyounes

Co-Authored-By: Ora Agent <noreply@oratelecom.net>
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.

MiniCPM4 forward_step Attention Performance Fix

1 participant