[ROCm] Split skinny_gemms_int8.cu into per-N translation units

.github/workflows/build-rocm-wheels.yml

@@ -121,7 +121,13 @@ jobs:
         env:
           PYTORCH_ROCM_ARCH: ${{ env.PYTORCH_ROCM_ARCH }}
           VLLM_TARGET_DEVICE: rocm
-          MAX_JOBS: 1
+          # MAX_JOBS=2 was reverted (adf8d91564) due to OOM on 16 GB runners.
+          # The heavy TUs have since been split into per-N shards (see
+          # csrc/rocm/skinny_gemms_int4/, skinny_gemms/, skinny_gemms_int8/,
+          # attention/, and csrc/libtorch_stable/attention/paged_attention_v1/).
+          # Peak RSS per shard is now <=3.2 GB (down from 12.6 GB for int4).
+          # Measured worst-case concurrent peak at MAX_JOBS=4: ~13 GB.
+          MAX_JOBS: 4
           SCCACHE_GHA_ENABLED: "true"
           SCCACHE_IDLE_TIMEOUT: 0
         run: |

CMakeLists.txt

@@ -644,7 +644,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
     "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
+    "csrc/libtorch_stable/mamba/selective_scan_fwd/shard_fp16.cu"
+    "csrc/libtorch_stable/mamba/selective_scan_fwd/shard_bf16.cu"
     "csrc/libtorch_stable/attention/paged_attention_v1.cu"
+    "csrc/libtorch_stable/attention/paged_attention_v1/shard_small.cu"
+    "csrc/libtorch_stable/attention/paged_attention_v1/shard_medium.cu"
+    "csrc/libtorch_stable/attention/paged_attention_v1/shard_large.cu"
     "csrc/libtorch_stable/attention/paged_attention_v2.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
     "csrc/libtorch_stable/cache_kernels.cu"
@@ -1295,15 +1300,32 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_ROCM_EXT_SRC
     "csrc/rocm/torch_bindings.cpp"
     "csrc/rocm/skinny_gemms.cu"
+    "csrc/rocm/skinny_gemms/instantiate_n1.cu"
+    "csrc/rocm/skinny_gemms/instantiate_n2.cu"
+    "csrc/rocm/skinny_gemms/instantiate_n3.cu"
+    "csrc/rocm/skinny_gemms/instantiate_n4.cu"
+    "csrc/rocm/skinny_gemms/instantiate_n5.cu"
     "csrc/rocm/skinny_gemms_int8.cu"
+    "csrc/rocm/skinny_gemms_int8/instantiate_n1.cu"
+    "csrc/rocm/skinny_gemms_int8/instantiate_n2.cu"
+    "csrc/rocm/skinny_gemms_int8/instantiate_n3.cu"
+    "csrc/rocm/skinny_gemms_int8/instantiate_n4.cu"
+    "csrc/rocm/skinny_gemms_int8/instantiate_n5.cu"
     "csrc/rocm/skinny_gemms_int4.cu"
+    "csrc/rocm/skinny_gemms_int4/instantiate_n1.cu"
+    "csrc/rocm/skinny_gemms_int4/instantiate_n2.cu"
+    "csrc/rocm/skinny_gemms_int4/instantiate_n3.cu"
+    "csrc/rocm/skinny_gemms_int4/instantiate_n4.cu"
+    "csrc/rocm/skinny_gemms_int4/instantiate_n5.cu"
     "csrc/rocm/skinny_gemms_w8a8.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n1.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n2.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n3.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n4.cu"
     "csrc/rocm/skinny_gemms_w8a8/instantiate_n5.cu"
-    "csrc/rocm/attention.cu")
+    "csrc/rocm/attention.cu"
+    "csrc/rocm/attention/dispatch_auto.cu"
+    "csrc/rocm/attention/dispatch_fp8.cu")
 
   set(VLLM_ROCM_FLAGS ${VLLM_GPU_FLAGS})
 

csrc/libtorch_stable/attention/paged_attention_v1.cu

@@ -16,175 +16,61 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "../torch_utils.h"
-#include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
 
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
+#include "paged_attention_v1/common.h"
 
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
-  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
-      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
-                                              BLOCK_SIZE, NUM_THREADS,      \
-                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
-      shared_mem_size);                                                     \
-  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
-                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
-      <<<grid, block, shared_mem_size, stream>>>(                           \
-          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
-          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
-          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
-          k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks,      \
-          blocksparse_vert_stride, blocksparse_block_size,                  \
-          blocksparse_head_sliding_step);
-
-// TODO(woosuk): Tune NUM_THREADS.
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
-          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
-          int NUM_THREADS = 128>
-void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
+void paged_attention_v1(
+    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::stable::Tensor&
+        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
+    torch::stable::Tensor&
+        value_cache,       // [num_blocks, num_heads, head_size, block_size]
+    int64_t num_kv_heads,  // [num_heads]
+    double scale,
+    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::stable::Tensor& seq_lens,      // [num_seqs]
+    int64_t block_size, int64_t max_seq_len,
     const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
-  int num_seqs = query.size(0);
-  int num_heads = query.size(1);
+    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
+    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step) {
   int head_size = query.size(2);
-  int max_num_blocks_per_seq = block_tables.size(1);
-  int q_stride = query.stride(0);
-  int kv_block_stride = key_cache.stride(0);
-  int kv_head_stride = key_cache.stride(1);
-
-  // NOTE: alibi_slopes is optional.
-  const float* alibi_slopes_ptr =
-      alibi_slopes
-          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
-          : nullptr;
-
-  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
-  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
-  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
-  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
-  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
-  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
-
-  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
-  int padded_max_seq_len =
-      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
-  int logits_size = padded_max_seq_len * sizeof(float);
-  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
-  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
-  // Keep that in sync with the logic here!
-  int shared_mem_size = std::max(logits_size, outputs_size);
-
-  dim3 grid(num_heads, num_seqs, 1);
-  dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
   switch (head_size) {
-    // NOTE(woosuk): To reduce the compilation time, we only compile for the
-    // head sizes that we use in the model. However, we can easily extend this
-    // to support any head size which is a multiple of 16.
     case 32:
-      LAUNCH_PAGED_ATTENTION_V1(32);
-      break;
     case 64:
-      LAUNCH_PAGED_ATTENTION_V1(64);
-      break;
     case 80:
-      LAUNCH_PAGED_ATTENTION_V1(80);
+      paged_attention_v1_small(
+          out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
+          seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
+          k_scale, v_scale, tp_rank, blocksparse_local_blocks,
+          blocksparse_vert_stride, blocksparse_block_size,
+          blocksparse_head_sliding_step);
       break;
     case 96:
-      LAUNCH_PAGED_ATTENTION_V1(96);
-      break;
     case 112:
-      LAUNCH_PAGED_ATTENTION_V1(112);
-      break;
     case 120:
-      LAUNCH_PAGED_ATTENTION_V1(120);
-      break;
     case 128:
-      LAUNCH_PAGED_ATTENTION_V1(128);
+      paged_attention_v1_medium(
+          out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
+          seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
+          k_scale, v_scale, tp_rank, blocksparse_local_blocks,
+          blocksparse_vert_stride, blocksparse_block_size,
+          blocksparse_head_sliding_step);
       break;
     case 192:
-      LAUNCH_PAGED_ATTENTION_V1(192);
-      break;
     case 256:
-      LAUNCH_PAGED_ATTENTION_V1(256);
+      paged_attention_v1_large(
+          out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
+          seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
+          k_scale, v_scale, tp_rank, blocksparse_local_blocks,
+          blocksparse_vert_stride, blocksparse_block_size,
+          blocksparse_head_sliding_step);
       break;
     default:
       STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
       break;
   }
 }
-
-#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
-  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
-                              IS_BLOCK_SPARSE>(                              \
-      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
-      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
-      blocksparse_local_blocks, blocksparse_vert_stride,                     \
-      blocksparse_block_size, blocksparse_head_sliding_step);
-
-#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
-  if (is_block_sparse) {                                                   \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
-  } else {                                                                 \
-    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
-  }
-
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
-// 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
-  }
-
-void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
-        value_cache,       // [num_blocks, num_heads, head_size, block_size]
-    int64_t num_kv_heads,  // [num_heads]
-    double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step) {
-  const bool is_block_sparse = (blocksparse_vert_stride > 1);
-
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
-                             CALL_V1_LAUNCHER_BLOCK_SIZE)
-}
-
-#undef MAX
-#undef MIN
-#undef DIVIDE_ROUND_UP