perf: accelerate fp16 HNSW search with AMX#7933
Conversation
📝 WalkthroughWalkthroughAdds an AMX-FP16 batched dot-product kernel, metric-aware distance batching, and 16-wide neighbor distance computation in graph search. Build detection, runtime dispatch, scalar fallbacks, numerical checks, recall tests, and benchmark coverage are included. ChangesAMX-FP16 Batched Vector Search
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant beam_search_loop
participant DistCalculator
participant FlatFloatDistanceCalc
participant dot_f16_batch_16
beam_search_loop->>DistCalculator: Request distances for 16 neighbor IDs
DistCalculator->>FlatFloatDistanceCalc: Dispatch distance_batch_16
FlatFloatDistanceCalc->>dot_f16_batch_16: Compute fp16 Dot batch
dot_f16_batch_16-->>FlatFloatDistanceCalc: Return 16 distances
FlatFloatDistanceCalc-->>beam_search_loop: Apply neighbor acceptance logic
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
rust/lance-index/src/vector/flat/storage.rs-413-413 (1)
413-413: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDrop the
_prefix now that_distance_typeis read. The binary constructors now store the parameter into thedistance_typefield, so it is no longer an unused binding. Rename the parameter (and this assignment) todistance_typein bothnew_binary(Line 413) andnew_binary_from_id(Line 431).As per coding guidelines: "Reserve
_-prefixed names for truly unused bindings; if a variable is read, remove the underscore."♻️ Proposed rename
fn new_binary( vectors: &'a FixedSizeListArray, query: ArrayRef, - _distance_type: DistanceType, + distance_type: DistanceType, ) -> Self { ... - distance_type: _distance_type, + distance_type, distance_fn: hamming,Apply the same to
new_binary_from_id(Line 431).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/vector/flat/storage.rs` at line 413, Rename the distance-type parameter from _distance_type to distance_type, and update the distance_type field assignment in both new_binary and new_binary_from_id; preserve the existing constructor behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Other comments:
In `@rust/lance-index/src/vector/flat/storage.rs`:
- Line 413: Rename the distance-type parameter from _distance_type to
distance_type, and update the distance_type field assignment in both new_binary
and new_binary_from_id; preserve the existing constructor behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 3da444c0-c467-448a-9f36-17e0ab854fd4
📒 Files selected for processing (9)
rust/lance-index/src/vector/flat/storage.rsrust/lance-index/src/vector/graph.rsrust/lance-index/src/vector/storage.rsrust/lance-linalg/build.rsrust/lance-linalg/src/distance.rsrust/lance-linalg/src/distance/dot_f16.rsrust/lance-linalg/src/simd.rsrust/lance-linalg/src/simd/amx_fp16.crust/lance-linalg/src/simd/amx_fp16.rs
I'm not a maintainer this is just my thoughts...Relatively clean change but not sure if this should be the direction of an OSS package. Narrow gain for significant change that's a hot loop for every user:
Broad change:
|
2 participants
Summary
DistCalculator::distance_batch_16hook with a scalar-compatible default implementation._tile_dpfp16ps) kernel.The batched graph loop preserves visited-node ordering and applies neighbor acceptance in the same order as the scalar path. Non-FP16-Dot backends retain the scalar-compatible default.
Validation
cargo fmt --all -- --checkcargo clippy -p lance-linalg --features fp16kernels --testscargo clippy -p lance-index --features lance-linalg/fp16kernels --testscargo test -p lance-linalg --features fp16kernels dot_f16 -- --nocapturecargo test -p lance-index --features lance-linalg/fp16kernels f16 -- --nocaptureLANCE_DISABLE_AMX=1 cargo test -p lance-index --features lance-linalg/fp16kernels f16 -- --nocaptureLANCE_NGRAM_NUM_PARTITIONS=128 cargo test -p lance-index --libThe FP16 HNSW test reports recall@10 of 0.993 on both AMX and fallback paths.
Performance
Measured on an Intel Xeon 6972P using flat FP16 + Dot HNSW search. The comparison below is AMX-FP16 enabled versus
LANCE_DISABLE_AMX=1using the same post-patch batched HNSW traversal and FP16 fallback path:The gain grows with dimension as fixed tile/setup overhead is amortized, and narrows at full machine concurrency as memory bandwidth and scheduling become the bottleneck. These figures isolate the AMX kernel and dispatch benefit within the post-patch batched path; they are not a three-way pre-patch / post-patch-AMX-off / post-patch-AMX-on comparison. FP16 HNSW recall@10 remained 0.993 in both AMX and fallback runs.