feat(projection_kernel): add attention-head subspace affinity - #1721
Merged
jlarson4 merged 4 commits intoAug 28, 2026
Conversation
janmenjayap
marked this pull request as ready for review
August 24, 2026 16:06
janmenjayap
changed the title
6 tasks
jlarson4
reviewed
Aug 26, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Excellent work here @janmenjayap! This is a solid implementation of your plan from #1720. Just a couple comments to tighten things up before merging
- Add an MPS CPU fallback for principal-angle singular values while preserving result devices. - Make rank tolerance storage-aware without collapsing realistic FP16 and BF16 matrices, and reject material cosine-bound violations. - Preserve unsupported-weight errors, return independent contiguous masks, and expand numerical regression coverage.
- Add jaxtyping shape annotations and runtime-validation coverage for public tensor interfaces. - Clarify measured ranks, retained basis widths, and storage-aware tolerance behavior. - Restore alphabetical exports and document memory usage and quadratic pair-ranking costs.
6 tasks
Collaborator
|
Great work on this @janmenjayap, looks great, merging now. |
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
2 participants
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.
Description
Implements #1720.
Adds a basis-invariant Projection Kernel analysis surface in two reviewable commits:
numerical-rank metadata, roundoff-bound handling, and random-subspace moments;
hybrid-layer indices, forward/all layer masks, bounded-memory tiled scoring, and ranking.
The implementation promotes fp16/bf16 inputs to fp32 for stable SVDs, preserves float64,
detaches model weights from autograd, and reports role/layer/head context for malformed or
rank-deficient weights. The accompanying guide documents orientation, rank semantics,
limitations, GQA behavior, and the distinction from Composition Score, with runnable numerical
and GPT-2 TransformerBridge examples.
Validation
make check-format: passed.uv run mypy .: passed (392 source files).uv run build-docs: passed with no PK-specific warnings.uv build: source distribution and wheel built successfully.make test-prdid not exit successfully because of the threeintegration failures described below):
local HF token lacks access, and one existing GraniteMoeHybrid MPS parity test reproduced the
investigated PyTorch 2.10 MPS numerical divergence.
and notebook checks passed.
Checklist
docs/source/content/projection_kernel.md.Type of change