Skip to content

feat(phi3-small): implement blocksparse attention - #581

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-511-phi3small-blocksparse
Jul 1, 2026
Merged

feat(phi3-small): implement blocksparse attention#581
inureyes merged 1 commit into
mainfrom
feature/issue-511-phi3small-blocksparse

Conversation

@inureyes

@inureyes inureyes commented Jul 1, 2026

Copy link
Copy Markdown
Member

Summary

Replaces the dense-attention fallback on Phi-3-small's non-dense layers with the trained blocksparse pattern (local blocks plus a vertical stride), while the periodic fully-dense layers keep dense attention. This restores long-context fidelity and the intended long-sequence efficiency; short context is unchanged.

Changes

  • Blocksparse mask / gather for layers where the per-layer block_sparse flag is set: attend to the local window of blocksparse_num_local_blocks blocks of size blocksparse_block_size ending at the current query block, plus earlier blocks selected by blocksparse_vert_stride.
  • Dense path preserved unchanged for the every-dense_attention_every_n_layers dense layers.
  • The already-computed per-layer block_sparse flag is threaded into the attention path so dense and sparse layers diverge correctly.
  • The L==1 decode path gathers the selected key blocks instead of allocating a full sparse mask each step.

Correctness (all covered by unit tests)

  • short_context_mask_equals_causal and short_context_output_matches_dense_fallback: when the sequence fits within the local window the blocksparse path reduces to the dense fallback.
  • blocksparse_mask_matches_reference_pattern and blocksparse_selected_key_blocks_are_local_plus_vertical: the selected key-block set matches the reference (local blocks + vertical stride) on a constructed case.
  • long_context_mask_is_sparser_than_causal and long_context_output_diverges_from_dense: long context follows the trained sparsity, not dense.
  • decode_single_row_mask_follows_pattern: the L==1 decode gather follows the pattern.
  • block_sparse_flag_marks_every_nth_layer_dense, config_parses_blocksparse_fields, config_defaults_apply_when_blocksparse_fields_absent.

Validation

  • cargo check --lib --tests --features metal,accelerate: compiles clean.
  • cargo test --release --features metal,accelerate phi3small: 10 tests pass.
  • Not validated here: a real Phi-3-small checkpoint generating coherently beyond the local-block window. No Phi-3-small checkpoint is present in this environment (models/ is unpopulated), so end-to-end long-context generation was not executed and should be confirmed on a real checkpoint before relying on it in production.

Closes #511

Replace the dense-attention fallback on phi3small's non-dense layers with the trained blocksparse pattern. Each block_sparse layer attends to the local window of blocksparse_num_local_blocks blocks of size blocksparse_block_size ending at the current query block, plus earlier blocks selected by blocksparse_vert_stride, while the periodic dense_attention_every_n_layers layers keep dense attention. The per-layer block_sparse flag is threaded into the attention path, and the L==1 decode path gathers the selected key blocks instead of building a full sparse mask each step, so short context reduces to the dense fallback while long context follows the trained sparsity.
@inureyes inureyes added the status:done Completed label Jul 1, 2026
@inureyes
inureyes merged commit 9760e27 into main Jul 1, 2026
5 checks passed
@inureyes
inureyes deleted the feature/issue-511-phi3small-blocksparse branch July 1, 2026 11:23
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions priority:low Low priority labels Sep 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

priority:low Low priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(phi3-small): implement blocksparse attention to replace the dense fallback

1 participant