[FEA]: Follow up work to "Segmented TopK using Thread-Block-Clusters" #10023

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

CUB

Is your feature request related to a problem? Please describe.

There are things that we do not want to resolve in #9224 to get it merged ASAP.

Describe the solution you'd like

• Improve code quality with abstraction:
  • Use BlockLoad for non-bulk loads.
  • Ring buffers for slots/stages/mbarriers?
  • TMA abstraction (Work on BlockLoadToShared or a separate piepelining abstraction)
• Docs:
  • Add API examples for determinism & tie-breaking requirements, now that clusters supports it.
• Dispatch:
  • Logging
  • Can we cache certain CUDA runtime queries in static variables?
  • Add support for launching via driver API launcher factory.
  • Fix MaxPotentialDynamicSmemBytes() and then use launcher_factory.max_dynamic_smem_size_for
• Arg Fw:
  • Try to unify verification and traits in the argument annotation framework. Concretely, segment_size_is_integral and segment_sizes_handle_ok would come to mind.
  • Should bounded_offset_t in cub/detail/segmented_params.cuh protect against negative upper bound instead of needing to catch this in each interface separately?
  • Add a simple way to get element type from Arg Fw traits (currently value/element type depending on wrapper type)
  • Make sure CUB primitives align in their Arg Fw usage.
• Benchmarking/Tuning:
  • Make sure items-per-thread for scans is uneven to avoid smem bank conflicts.
  • Add scan algo(s) to tuning policy.
  • Make baseline tunable (currently its policy_selector_from_types is implemented in a way that assumes arch-independent tuning)
  • Rename benchmark files for "baseline" backend to have the backend name in the file name once we have a proper backend name.
  • Improve tuning for edge cases: DeviceSegmentedTopK cluster backend tunings are only used for small k #11358 and The copy-fast-path of the DeviceSegmentedTopK cluster backend is barely tuned #11359
• Testing:
  • Clean up tests (compile time and code size). See Implement Segmented TopK using Thread-Block-Clusters #9224 (review)
• Feature:
  • Add a fallback through the cluster-based implementation (single block, but with streaming/ping-pong/determinism) for pre-Hopper.
• Perf:
  • Have the histograms in a striped layout in smem so the load for the scan is cheaper (no bank conflicts). Can we do sth about the key loads for the selection scan as well?
  • Try to improve the implementation of filter-pass-early exit by checking more rarely and reusing results from earlier atomics w/ 2x __syncthreads_or().
  • Use local mbarriers instead of cluster.barrier for communication: [FEA]: Use mbarriers to improve barrier/communication latency in cluster segmented TopK #9994
  • Try to use tensor memory to further expand the number of resident (not re-streamed) keys for large segments.
  • Can we save on register pressure by not keeping layout.block_keys_in = d_key_segments_it[segment_id] around without regressing cases where that de-reference is expensive (needs new benchmark). Maybe even make the behavior depend on the type of iterator.
  • Try using reusing atomic results and synchronizing them via __syncthreads_or() to avoid reads in stopping condition of filters.
  • Try packing radix pass state (bucket id + early stop flag) into a single 32b DSMEM load/store.
  • Re-assess 64-bit type usage Implement Segmented TopK using Thread-Block-Clusters #9224 (comment)
  • Avoid using smem for edges when not using bulk copies Implement Segmented TopK using Thread-Block-Clusters #9224 (comment)

Describe alternatives you've considered

No response

Additional context

No response