Added wrapper for setting devFuncAttributeMaxDynamicSharedMemorySize

csrc/attention/attention_kernels.cu

@@ -542,31 +542,52 @@ __global__ void paged_attention_v2_reduce_kernel(
 
 } // namespace vllm
 
-#ifndef USE_ROCM
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                                  \
-  cudaFuncSetAttribute(                                                                       \
-    (void*)vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>,            \
-    cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem_size);                            \
-  vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>                      \
-  <<<grid, block, shared_mem_size, stream>>>(                                                 \
-    out_ptr,                                                                                  \
-    query_ptr,                                                                                \
-    key_cache_ptr,                                                                            \
-    value_cache_ptr,                                                                          \
-    head_mapping_ptr,                                                                         \
-    scale,                                                                                    \
-    block_tables_ptr,                                                                         \
-    context_lens_ptr,                                                                         \
-    max_num_blocks_per_seq,                                                                   \
-    alibi_slopes_ptr,                                                                         \
-    q_stride,                                                                                 \
-    kv_block_stride,                                                                          \
-    kv_head_stride);
-#else
+// #ifndef USE_ROCM
+// #define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                                  \
+//   cudaFuncSetAttribute(                                                                       \
+//     (void*)vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>,            \
+//     cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem_size);                            \
+//   vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>                      \
+//   <<<grid, block, shared_mem_size, stream>>>(                                                 \
+//     out_ptr,                                                                                  \
+//     query_ptr,                                                                                \
+//     key_cache_ptr,                                                                            \
+//     value_cache_ptr,                                                                          \
+//     head_mapping_ptr,                                                                         \
+//     scale,                                                                                    \
+//     block_tables_ptr,                                                                         \
+//     context_lens_ptr,                                                                         \
+//     max_num_blocks_per_seq,                                                                   \
+//     alibi_slopes_ptr,                                                                         \
+//     q_stride,                                                                                 \
+//     kv_block_stride,                                                                          \
+//     kv_head_stride);
+// #else
+// #define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                                  \
+//   hipFuncSetAttribute(                                                                       \
+//     (void*)vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>,            \
+//     hipFuncAttributeMaxDynamicSharedMemorySize, shared_mem_size);                            \
+//   vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>                      \
+//   <<<grid, block, shared_mem_size, stream>>>(                                                 \
+//     out_ptr,                                                                                  \
+//     query_ptr,                                                                                \
+//     key_cache_ptr,                                                                            \
+//     value_cache_ptr,                                                                          \
+//     head_mapping_ptr,                                                                         \
+//     scale,                                                                                    \
+//     block_tables_ptr,                                                                         \
+//     context_lens_ptr,                                                                         \
+//     max_num_blocks_per_seq,                                                                   \
+//     alibi_slopes_ptr,                                                                         \
+//     q_stride,                                                                                 \
+//     kv_block_stride,                                                                          \
+//     kv_head_stride);
+// #endif
+
 #define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                                  \
-  hipFuncSetAttribute(                                                                       \
-    (void*)vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>,            \
-    hipFuncAttributeMaxDynamicSharedMemorySize, shared_mem_size);                            \
+  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                                       \
+    ((void*)vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>),            \
+    shared_mem_size);                                                                         \
   vllm::paged_attention_v1_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS>                      \
   <<<grid, block, shared_mem_size, stream>>>(                                                 \
     out_ptr,                                                                                  \
@@ -582,7 +603,6 @@ __global__ void paged_attention_v2_reduce_kernel(
     q_stride,                                                                                 \
     kv_block_stride,                                                                          \
     kv_head_stride);
-#endif
 
 // TODO(woosuk): Tune NUM_THREADS.
 template<

csrc/cuda_compat.h

@@ -13,7 +13,15 @@
 #endif
 
 #ifndef USE_ROCM
-  #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane);
+  #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane)
 #else
   #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane)
-#endif
+#endif
+
+#ifndef USE_ROCM
+  #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \
+    cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL)
+#else
+  #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \
+    hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL)
+#endif