[MUSA][8/N] Port CUDA kernels that are compatible with MUSA (#17946)

Signed-off-by: yafeng.li <yafeng.li@mthreads.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai>
2026-04-23 18:04:58 -07:00
parent c0166355ae
commit 74c2e5bacd
15 changed files with 1905 additions and 66 deletions
@@ -7,11 +7,18 @@

 #include <cstdint>

+#ifndef USE_MUSA
 __forceinline__ __device__ int get_lane_id() {
  int lane_id;
  asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
  return lane_id;
 }
+#else
+constexpr int WarpSize = 32;
+__forceinline__ __device__ int get_lane_id() {
+  return threadIdx.x % WarpSize;
+}
+#endif

 int ceil_div(int a, int b) {
  return (a + b - 1) / b;