@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a547aeb828
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 08:30 UTC

feat(rocm-support): support mamba2 on rocm (#18565)

Signed-off-by: Islam Almersawi <islam.almersawi@openinnovation.ai> Co-authored-by: Islam Almersawi <islam.almersawi@openinnovation.ai>

almersawi committed May 27, 2025 at 11:07 UTC a547aeb8283b6ba757912fe43169ec38adb9ca45
5 files changed +60 -49
CMakeLists.txt
+2 -2
@@ -232,6 +232,8 @@ endif()
232 #
233
234 set(VLLM_EXT_SRC
235 + "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
236 + "csrc/mamba/causal_conv1d/causal_conv1d.cu"
237 "csrc/cache_kernels.cu"
238 "csrc/attention/paged_attention_v1.cu"
239 "csrc/attention/paged_attention_v2.cu"
@@ -287,8 +289,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
289 FetchContent_MakeAvailable(cutlass)
290
291 list(APPEND VLLM_EXT_SRC
290 - "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
291 - "csrc/mamba/causal_conv1d/causal_conv1d.cu"
292 "csrc/quantization/aqlm/gemm_kernels.cu"
293 "csrc/quantization/awq/gemm_kernels.cu"
294 "csrc/permute_cols.cu"
csrc/mamba/causal_conv1d/causal_conv1d.cu
+4 -6
@@ -13,6 +13,10 @@
13 #include <cub/block/block_load.cuh>
14 #include <cub/block/block_store.cuh>
15
16 +#ifdef USE_ROCM
17 + namespace cub = hipcub;
18 +#endif
19 +
20 #include "static_switch.h"
21
22
@@ -501,15 +505,9 @@ void causal_conv1d_fwd_launch(ConvParamsBase &params, cudaStream_t stream) {
505 auto kernel = &causal_conv1d_fwd_kernel<Ktraits>;
506
507 if (kSmemSize >= 48 * 1024) {
504 - #ifndef USE_ROCM
505 - C10_CUDA_CHECK(cudaFuncSetAttribute(
506 - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
507 - #else
508 - // There is a slight signature discrepancy in HIP and CUDA "FuncSetAttribute" function.
508 C10_CUDA_CHECK(cudaFuncSetAttribute(
509 (void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
510 std::cerr << "Warning (causal_conv1d fwd launch): attempting to set maxDynamicSharedMemorySize on an AMD GPU which is currently a non-op (in ROCm versions <= 6.1). This might lead to undefined behavior. \n" << std::endl;
512 - #endif
511 }
512 kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
513
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+1 -1
@@ -321,7 +321,7 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
321 auto kernel = &selective_scan_fwd_kernel<Ktraits>;
322 if (kSmemSize >= 48 * 1024) {
323 C10_CUDA_CHECK(cudaFuncSetAttribute(
324 - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
324 + (void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
325 }
326 kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
327 C10_CUDA_KERNEL_LAUNCH_CHECK();
csrc/torch_bindings.cpp
+35 -35
@@ -482,41 +482,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
482 " Tensor page_table, float scale) -> ()");
483 ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
484
485 - // Mamba selective scan kernel
486 - ops.def(
487 - "selective_scan_fwd(Tensor! u, Tensor! delta,"
488 - "Tensor! A, Tensor! B, Tensor! C,"
489 - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
490 - "bool delta_softplus,"
491 - "Tensor? query_start_loc,"
492 - "Tensor? cache_indices,"
493 - "Tensor? has_initial_state,"
494 - "Tensor! ssm_states,"
495 - "int pad_slot_id) -> ()");
496 - ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
497 -
498 - ops.def(
499 - "causal_conv1d_update(Tensor! x,"
500 - "Tensor! conv_state,"
501 - "Tensor! weight,"
502 - "Tensor? bias_,"
503 - "bool silu_activation,"
504 - "Tensor? cache_seqlens_,"
505 - "Tensor? conv_state_indices,"
506 - "int pad_slot_id) -> ()");
507 - ops.impl("causal_conv1d_update", torch::kCUDA, &causal_conv1d_update);
508 -
509 - ops.def(
510 - "causal_conv1d_fwd(Tensor! x, Tensor! weight,"
511 - "Tensor? bias_,"
512 - "Tensor!? conv_states,"
513 - "Tensor? query_start_loc,"
514 - "Tensor? cache_indices,"
515 - "Tensor? has_initial_state,"
516 - "bool silu_activation,"
517 - "int pad_slot_id) -> ()");
518 - ops.impl("causal_conv1d_fwd", torch::kCUDA, &causal_conv1d_fwd);
519 -
485 // Compute NVFP4 block quantized tensor.
486 ops.def(
487 "scaled_fp4_quant(Tensor! output, Tensor input,"
@@ -584,6 +549,41 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
549 ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
550 &dynamic_scaled_int8_quant);
551
552 + // Mamba selective scan kernel
553 + ops.def(
554 + "selective_scan_fwd(Tensor! u, Tensor! delta,"
555 + "Tensor! A, Tensor! B, Tensor! C,"
556 + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
557 + "bool delta_softplus,"
558 + "Tensor? query_start_loc,"
559 + "Tensor? cache_indices,"
560 + "Tensor? has_initial_state,"
561 + "Tensor! ssm_states,"
562 + "int pad_slot_id) -> ()");
563 + ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
564 +
565 + ops.def(
566 + "causal_conv1d_update(Tensor! x,"
567 + "Tensor! conv_state,"
568 + "Tensor! weight,"
569 + "Tensor? bias_,"
570 + "bool silu_activation,"
571 + "Tensor? cache_seqlens_,"
572 + "Tensor? conv_state_indices,"
573 + "int pad_slot_id) -> ()");
574 + ops.impl("causal_conv1d_update", torch::kCUDA, &causal_conv1d_update);
575 +
576 + ops.def(
577 + "causal_conv1d_fwd(Tensor! x, Tensor! weight,"
578 + "Tensor? bias_,"
579 + "Tensor!? conv_states,"
580 + "Tensor? query_start_loc,"
581 + "Tensor? cache_indices,"
582 + "Tensor? has_initial_state,"
583 + "bool silu_activation,"
584 + "int pad_slot_id) -> ()");
585 + ops.impl("causal_conv1d_fwd", torch::kCUDA, &causal_conv1d_fwd);
586 +
587 #ifndef USE_ROCM
588 // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
589 ops.def(
vllm/model_executor/layers/mamba/mamba2_metadata.py
+18 -5
@@ -5,10 +5,9 @@ from dataclasses import dataclass
5 import torch
6
7 from vllm.attention.backends.abstract import AttentionMetadata
8 -from vllm.attention.backends.flash_attn import FlashAttentionMetadata
8 from vllm.attention.backends.placeholder_attn import (
9 PlaceholderAttentionMetadata)
11 -from vllm.attention.backends.xformers import XFormersMetadata
10 +from vllm.platforms import current_platform
11
12
13 @dataclass
@@ -23,6 +22,21 @@ class Mamba2Metadata:
22 chunk_offsets: torch.Tensor
23
24
25 +def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:
26 + """Returns the appropriate metadata classes for the current platform."""
27 + if current_platform.is_rocm():
28 + from vllm.attention.backends.rocm_flash_attn import (
29 + ROCmFlashAttentionMetadata)
30 + return (ROCmFlashAttentionMetadata, PlaceholderAttentionMetadata)
31 + elif current_platform.is_cuda():
32 + from vllm.attention.backends.flash_attn import FlashAttentionMetadata
33 + from vllm.attention.backends.xformers import XFormersMetadata
34 + return (FlashAttentionMetadata, XFormersMetadata,
35 + PlaceholderAttentionMetadata)
36 + raise ValueError(
37 + f"Unsupported platform for Mamba2: {current_platform.device_type}")
38 +
39 +
40 def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
41 chunk_size: int,
42 total_seqlens: int):
@@ -78,9 +92,8 @@ def prepare_mamba2_metadata(
92
93 # Compute seq_idx, chunk_indices and chunk_offsets for prefill only
94 if num_prefills > 0:
81 - if (isinstance(attn_metadata,
82 - (FlashAttentionMetadata, XFormersMetadata,
83 - PlaceholderAttentionMetadata))
95 + attn_metadata_instances = get_platform_metadata_classes()
96 + if (isinstance(attn_metadata, attn_metadata_instances)
97 and attn_metadata.context_lens_tensor is not None):
98 has_initial_states = \
99 attn_metadata.context_lens_tensor[:num_prefills] > 0 #[batch,]