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
There was a problem hiding this comment.
Many thanks @divyegala for the comprehensive PR!
I double-verified the functional correctness of the kernel from this PR. The kernel also matches one-to-one the performance with the one from gitlab.
dantegd
left a comment
dantegd
left a comment
There was a problem hiding this comment.
Could we make the path selection and workspace allocation agree on whether the cuTile launch is guaranteed to succeed? When minClusterDistanceCompute.cu selects FusedCutile, it doesn’t allocate the CUTLASS KVP scratch or mutex workspace. But if try_fused_1nn_tile() returns false, cpp/src/distance/detail/fused_distance_nn.cuh falls through to CUTLASS and expects those buffers to exist. That leaves us without a valid fallback if launcher construction fails, even though the try_ API suggests falling back is supported. Would it make sense to build or probe the launcher before committing to FusedCutile, or keep the fallback workspace allocated for as long as this path can still return false?
| cutlass_kvp_scratch, | ||
| stream); | ||
| break; | ||
| case cuvs::distance::DistanceType::InnerProduct: break; |
There was a problem hiding this comment.
Could we route InnerProduct to a real non-cuTile implementation when cuTile isn't built or available?
On CUDA 12, try_fused_1nn_tile is compiled as false, so this function initializes the scratch output to (0, max), hits this empty case, and then unpacks that untouched sentinel. The new input_fp32_fused test includes InnerProduct unconditionally, so the supported CUDA 12 configuration returns incorrect results. What do you think about selecting the unfused path upstream for InnerProduct when cuTile is unavailable, or implementing the fallback here?
There was a problem hiding this comment.
The unfused path does not support InnerProduct for some reason. Let me see if we can support or just throw for now.
| static_assert(std::is_same_v<IdxT, int> || std::is_same_v<IdxT, int64_t>); | ||
|
|
||
| constexpr int strict_pitch_elements = 16 / sizeof(DataT); | ||
| const bool use_strict_abi = k % strict_pitch_elements == 0; |
There was a problem hiding this comment.
Both the strict and relaxed signatures declare base_addr_divisible_by=16 for matrices, norms, indices, and distances, but this selection only checks row pitch. Valid RAFT subviews or DLPack byte offsets can therefore launch under a false compiler assumption; the int64 loop below also makes xn + max_i32 and nearest_dist + max_i32 non-16-byte-aligned on its second batch. cuTile documents violating an exported signature assumption as undefined behavior: “Launching the exported kernel with an array that doesn’t satisfy this assumption would then result in undefined behavior.”
Should we validate base alignment for every non-null array before launching, or export a base-alignment-1 fallback? What do you think about making alignment part of eligibility before choosing either ABI?
There was a problem hiding this comment.
Oh I did not think at all about subviews or offsets, good catch. For now, let's just validate base pointer alignment. I can add to the follow-up issue to validate 1-byte alignment and what kind of cubin is produced by cutile?
6 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++.