Pine/VoxCPM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: stabilize CPU SDPA mask broadcasting

Browse Source 
Use an explicit broadcastable attention mask shape during MiniCPM incremental decoding so CPU runtimes avoid a PyTorch SDPA dimension error without changing attention semantics.

Made-with: Cursor
This commit is contained in:
刘鑫
2026-04-13 15:38:53 +08:00
parent 5510503182
commit b1584aec7c
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/voxcpm/modules/minicpm4/model.py
+3 -1
View File
@@ -196,7 +196,9 @@ class MiniCPMAttention(nn.Module):
key_cache[:, :, position_id, :] = key_states key_cache[:, :, position_id, :] = key_states
value_cache[:, :, position_id, :] = value_states value_cache[:, :, position_id, :] = value_states
attn_mask = torch.arange(key_cache.size(2), device=key_cache.device) <= position_id # Use an explicit broadcastable mask shape for SDPA. A 1D mask can
# trigger a CPU-side dimension bug in some PyTorch versions.
attn_mask = (torch.arange(key_cache.size(2), device=key_cache.device) <= position_id).view(1, 1, 1, -1)
# ref: https://github.com/pytorch/pytorch/issues/163597 # ref: https://github.com/pytorch/pytorch/issues/163597
# there is a bug in MPS for non-contiguous tensors, so we need to make them contiguous # there is a bug in MPS for non-contiguous tensors, so we need to make them contiguous