Introduce fused GEMM + 1-NN primitive using cuTile - #2249
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
divyegala
changed the title
divyegala
added
feature request
|
/ok to test 2fdd7f7 |
|
/ok to test 2fb6253 |
|
/ok to test e28ecd9 |
|
/ok to test 169309d |
| if metric_code == METRIC_L2_EXPANDED: | ||
| # The A norm is constant across centroids. Reduce | ||
| # 0.5 * ||y||^2 - dot(x, y), then recover full L2 once. | ||
| score = (0.5 * b_norm)[None, :] - accumulator |
There was a problem hiding this comment.
I guess for L2, b_normand a_norm below are not yet squared? Because for cosine, we dropped the squared root on b_norm
| if require_tma_friendly_pitch | ||
| else (1, 1) | ||
| ), | ||
| shape_divisible_by=(1, 1), |
There was a problem hiding this comment.
Should we make shape_divisible_by=(2, 1) for sm_80 and sm_89 when elem_dtype is float16 and for strict ABI? This should already improve those sm's for 13.4 and float16
There was a problem hiding this comment.
If this will be fixed in 13.5 anyway, I think we can just wait and document the requirement to make sure we don't have too many conditionals lying around.
| else: | ||
| # Defer the A-norm division until after selecting the | ||
| # winning centroid. | ||
| score = -(accumulator / b_norm[None, :]) |
There was a problem hiding this comment.
Nit: I guess accumulator / (-b_norm)[None, :] might reduce slightly the number of flops as accumulator has the shape [tm, tn] while b_norm has [tn]
|
|
||
| if constexpr (std::is_same_v<LabelT, IdxT>) { | ||
| auto labels_view = raft::make_device_vector_view<IdxT, IdxT>(labels, n_rows); | ||
| cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>( |
There was a problem hiding this comment.
One small thing to fix along the way: we can use the public API here: min_cluster_and_distance
5 participants
This PR adds infrastructure built on top of existing JIT LTO architecture to generate kernels using
cutile-pythonat build time, and embed them in the C++ library to make them callable from C++.