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 14, 05:30 UTC
[V1][Mamba1] - FP32 SSM Kernel Support (#23506)
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
Asaf Joseph Gardin committed
Sep 2, 2025 at 06:53 UTC
2b41cbbf030dd6cf4d5441fe679ca1c9add0d0e6
3 files changed
+65
-32
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+47
-25
@@ -27,11 +27,12 @@
27
28
template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_,
29
bool kIsVariableB_, bool kIsVariableC_,
30
- bool kHasZ_, bool kVarlen_, typename input_t_, typename weight_t_>
30
+ bool kHasZ_, bool kVarlen_, typename input_t_, typename weight_t_, typename state_t_>
31
struct Selective_Scan_fwd_kernel_traits {
32
static_assert(kNItems_ % 4 == 0);
33
using input_t = input_t_;
34
using weight_t = weight_t_;
35
+ using state_t = state_t_;
36
static constexpr int kNThreads = kNThreads_;
37
// Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
38
static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
@@ -132,7 +133,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
133
input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride;
134
weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * kNRows * params.C_d_stride;
135
input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride;
135
- input_t *ssm_states = reinterpret_cast<input_t *>(params.ssm_states_ptr) +
136
+ typename Ktraits::state_t *ssm_states = reinterpret_cast<typename Ktraits::state_t *>(params.ssm_states_ptr) +
137
cache_index * params.ssm_states_batch_stride +
138
dim_id * kNRows * params.ssm_states_dim_stride;
139
@@ -261,7 +262,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
262
if (threadIdx.x == 0) {
263
smem_running_prefix[state_idx] = prefix_op.running_prefix;
264
if (chunk == n_chunks - 1) {
264
- ssm_states[state_idx * params.ssm_states_dstate_stride] = input_t(prefix_op.running_prefix.y);
265
+ ssm_states[state_idx * params.ssm_states_dstate_stride] = typename Ktraits::state_t(prefix_op.running_prefix.y);
266
}
267
}
268
#pragma unroll
@@ -310,7 +311,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
311
}
312
}
313
313
-template<int kNThreads, int kNItems, typename input_t, typename weight_t>
314
+template<int kNThreads, int kNItems, typename input_t, typename weight_t, typename state_t>
315
void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
316
// Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block
317
// processing 1 row.
@@ -321,7 +322,7 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
322
BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
323
BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {
324
BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] {
324
- using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, kVarlen, input_t, weight_t>;
325
+ using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, kVarlen, input_t, weight_t, state_t>;
326
constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);
327
dim3 grid(params.batch, params.dim / kNRows);
328
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
@@ -341,59 +342,78 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
342
});
343
}
344
344
-template<typename input_t, typename weight_t>
345
+template<typename input_t, typename weight_t, typename state_t>
346
void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) {
347
348
#ifndef USE_ROCM
349
if (params.seqlen <= 128) {
349
- selective_scan_fwd_launch<32, 4, input_t, weight_t>(params, stream);
350
+ selective_scan_fwd_launch<32, 4, input_t, weight_t, state_t>(params, stream);
351
} else if (params.seqlen <= 256) {
351
- selective_scan_fwd_launch<32, 8, input_t, weight_t>(params, stream);
352
+ selective_scan_fwd_launch<32, 8, input_t, weight_t, state_t>(params, stream);
353
} else if (params.seqlen <= 512) {
353
- selective_scan_fwd_launch<32, 16, input_t, weight_t>(params, stream);
354
+ selective_scan_fwd_launch<32, 16, input_t, weight_t, state_t>(params, stream);
355
} else if (params.seqlen <= 1024) {
355
- selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
356
+ selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream);
357
} else {
357
- selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
358
+ selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream);
359
}
360
#else
361
if (params.seqlen <= 256) {
361
- selective_scan_fwd_launch<64, 4, input_t, weight_t>(params, stream);
362
+ selective_scan_fwd_launch<64, 4, input_t, weight_t, state_t>(params, stream);
363
} else if (params.seqlen <= 512) {
363
- selective_scan_fwd_launch<64, 8, input_t, weight_t>(params, stream);
364
+ selective_scan_fwd_launch<64, 8, input_t, weight_t, state_t>(params, stream);
365
} else if (params.seqlen <= 1024) {
365
- selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
366
+ selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream);
367
} else {
367
- selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
368
+ selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream);
369
}
370
#endif
371
}
372
372
-template void selective_scan_fwd_cuda<at::BFloat16, float>(SSMParamsBase ¶ms, cudaStream_t stream);
373
-template void selective_scan_fwd_cuda<at::Half, float>(SSMParamsBase ¶ms, cudaStream_t stream);
374
-template void selective_scan_fwd_cuda<float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
373
+template void selective_scan_fwd_cuda<at::BFloat16, float, at::BFloat16>(SSMParamsBase ¶ms, cudaStream_t stream);
374
+template void selective_scan_fwd_cuda<at::BFloat16, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
375
+template void selective_scan_fwd_cuda<at::Half, float, at::Half>(SSMParamsBase ¶ms, cudaStream_t stream);
376
+template void selective_scan_fwd_cuda<at::Half, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
377
+template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
378
379
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
380
378
-#define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \
381
+#define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...) \
382
if (ITYPE == at::ScalarType::Half) { \
383
using input_t = at::Half; \
384
using weight_t = float; \
382
- __VA_ARGS__(); \
385
+ if (STYPE == at::ScalarType::Half) { \
386
+ using state_t = at::Half; \
387
+ __VA_ARGS__(); \
388
+ } else if (STYPE == at::ScalarType::Float) { \
389
+ using state_t = float; \
390
+ __VA_ARGS__(); \
391
+ } else { \
392
+ AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
393
+ } \
394
} else if (ITYPE == at::ScalarType::BFloat16) { \
395
using input_t = at::BFloat16; \
396
using weight_t = float; \
386
- __VA_ARGS__(); \
397
+ if (STYPE == at::ScalarType::BFloat16) { \
398
+ using state_t = at::BFloat16; \
399
+ __VA_ARGS__(); \
400
+ } else if (STYPE == at::ScalarType::Float) { \
401
+ using state_t = float; \
402
+ __VA_ARGS__(); \
403
+ } else { \
404
+ AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
405
+ } \
406
} else if (ITYPE == at::ScalarType::Float) { \
407
using input_t = float; \
408
using weight_t = float; \
409
+ using state_t = float; \
410
__VA_ARGS__(); \
411
} else { \
412
AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
413
}
414
415
396
-template<typename input_t, typename weight_t>
416
+template<typename input_t, typename weight_t, typename state_t>
417
void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream);
418
419
void set_ssm_params_fwd(SSMParamsBase ¶ms,
@@ -648,7 +668,9 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
668
669
// Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
670
at::Tensor out = delta;
651
- TORCH_CHECK(ssm_states.scalar_type() == input_type);
671
+ // ssm_states can now be either the same as input_type or float32
672
+ auto state_type = ssm_states.scalar_type();
673
+ TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float);
674
TORCH_CHECK(ssm_states.is_cuda());
675
TORCH_CHECK(ssm_states.stride(-1) == 1);
676
@@ -670,7 +692,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
692
693
const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
694
auto stream = at::cuda::getCurrentCUDAStream().stream();
673
- DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
674
- selective_scan_fwd_cuda<input_t, weight_t>(params, stream);
695
+ DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] {
696
+ selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream);
697
});
698
}
tests/models/language/generation/test_hybrid.py
+6
-1
@@ -65,6 +65,11 @@ V0_UNSUPPORTED_MODELS = [
65
"LiquidAI/LFM2-1.2B",
66
]
67
68
+FP32_STATE_MODELS = [
69
+ "state-spaces/mamba-130m-hf",
70
+ "Zyphra/Zamba2-1.2B-instruct",
71
+]
72
+
73
# Avoid OOM
74
MAX_NUM_SEQS = 4
75
@@ -434,7 +439,7 @@ def test_full_cuda_graph(
439
)
440
441
437
-@pytest.mark.parametrize("model", ["Zyphra/Zamba2-1.2B-instruct"])
442
+@pytest.mark.parametrize("model", FP32_STATE_MODELS)
443
@pytest.mark.parametrize("max_tokens", [64])
444
@pytest.mark.parametrize("num_logprobs", [5])
445
def test_fp32_state(
vllm/model_executor/layers/mamba/mamba_utils.py
+12
-6
@@ -30,12 +30,8 @@ class MambaStateDtypeCalculator:
30
mamba_cache_dtype: MambaDType,
31
mamba_ssm_cache_dtype: MambaDType,
32
) -> tuple[torch.dtype, ...]:
33
- # TODO (tdoublep) requires kernel changes
34
- if mamba_cache_dtype == "float32" or mamba_ssm_cache_dtype == "float32":
35
- raise ValueError("fp32 state for mamba1 is not yet supported")
36
- else:
37
- return MambaStateDtypeCalculator.mamba2_state_dtype(
38
- model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype)
33
+ return cls._mamba_state_dtype(model_dtype, mamba_cache_dtype,
34
+ mamba_ssm_cache_dtype)
35
36
@classmethod
37
def mamba2_state_dtype(
@@ -43,6 +39,16 @@ class MambaStateDtypeCalculator:
39
model_dtype: Union[ModelDType, torch.dtype],
40
mamba_cache_dtype: MambaDType,
41
mamba_ssm_cache_dtype: MambaDType,
42
+ ) -> tuple[torch.dtype, ...]:
43
+ return cls._mamba_state_dtype(model_dtype, mamba_cache_dtype,
44
+ mamba_ssm_cache_dtype)
45
+
46
+ @classmethod
47
+ def _mamba_state_dtype(
48
+ cls,
49
+ model_dtype: Union[ModelDType, torch.dtype],
50
+ mamba_cache_dtype: MambaDType,
51
+ mamba_ssm_cache_dtype: MambaDType,
52
) -> tuple[torch.dtype, ...]:
53
conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype,
54
model_dtype)