chore(core): correct the sink attention Used by comment
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Found during review of PR #1779.
## Problem
`fast_scaled_dot_product_attention_with_sinks` carries `Used by: GptOss` at all three of its declaration sites, but four families call it: GPT-OSS, Laguna, DeepSeek-V4 and Falcon-OCR. The same comment calls a sink a "per-head attention bias for first position". A sink is an extra per-head logit that joins the softmax denominator and has no value row: MLX's fallback in https://github.com/ml-explore/mlx/blob/main/mlx/fast.cpp concatenates it in front of the scores, applies the softmax, and slices it off. `docs/code-guidelines.md` ("Shared Function Comments") makes the `// Used by:` comment the record of what a change to a shared function can break, so the stale list hides three callers.
## Evidence
- Comment sites: `src/lib/mlxcel-core/src/lib.rs:1272-1273`, `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h:1026-1027`, `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:3742-3743`.
- Callers: `src/models/gpt_oss.rs:764`, `src/models/laguna_layers.rs:183`, `src/models/deepseek_v4_attention.rs:346`, `src/models/falcon_ocr.rs:435`. These are the only call sites in `src/`, `examples/`, `tests/` and `benches/`.
- The other `Used by: GptOss` lines (`lib.rs:633`, `mlx_cxx_bridge.h:522`, `mlx_cxx_bridge.cpp:1511`) belong to the GPT-OSS SwiGLU activation and are correct; leave them.
## Proposed fix
At each of the three sites, replace the two comment lines with the text below (`///` in `lib.rs`, `//` in the C++ files). Comment-only; no code changes.
```
Fast SDPA with optional attention sinks: an extra per-head logit added to the softmax denominator, with no value row
Used by: GptOss, Laguna, DeepSeekV4, FalconOcr
```
## Acceptance criteria
- [ ] All three comments name the four families and describe the sink as an extra softmax logit.
- [ ] `grep -rn 'fast_scaled_dot_product_attention_with_sinks(' src/models` still lists exactly the four callers above.
## Verification
```bash
grep -rn 'fast_scaled_dot_product_attention_with_sinks(' src/models
grep -n -B3 'fast_scaled_dot_product_attention_with_sinks(' src/lib/mlxcel-core/src/lib.rs src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp
cargo fmt --all -- --check
cargo build --release --features metal,accelerate
```
Contributor guide
Research direction
Start with the three named declaration sites: src/lib/mlxcel-core/src/lib.rs, mlx_cxx_bridge.h, and mlx_cxx_bridge.cpp, then compare them with the four caller locations in src/models. Verify that only the sink-function comments change, all four model families are named, the sink definition is accurate, and the listed grep, formatting, and build checks pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- documentation
- Issue type
- Documentation
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 90/100