feat(phi3-small): implement blocksparse attention - #581
Merged
Conversation
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.
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
block_sparseflag is set: attend to the local window ofblocksparse_num_local_blocksblocks of sizeblocksparse_block_sizeending at the current query block, plus earlier blocks selected byblocksparse_vert_stride.dense_attention_every_n_layersdense layers.block_sparseflag is threaded into the attention path so dense and sparse layers diverge correctly.L==1decode 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_causalandshort_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_patternandblocksparse_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_causalandlong_context_output_diverges_from_dense: long context follows the trained sparsity, not dense.decode_single_row_mask_follows_pattern: theL==1decode 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.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