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 9, 14:45 UTC
[Performance] Tune Mamba selective scan kernel for B200 (#32873)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
danisereb committed
Jan 26, 2026 at 15:56 UTC
f4a0921c9c1194abf0a5e53218cfe5fb110eedac
2 files changed
+26
-11
vllm/model_executor/layers/mamba/mamba_mixer2.py
+5
@@ -41,6 +41,7 @@ from vllm.model_executor.model_loader.weight_utils import (
41
sharded_weight_loader,
42
)
43
from vllm.model_executor.utils import set_weight_attrs
44
+from vllm.platforms import current_platform
45
from vllm.utils.torch_utils import direct_register_custom_op
46
from vllm.v1.attention.backend import AttentionMetadata
47
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
@@ -502,6 +503,9 @@ class MambaMixer2(MambaBase, CustomOp):
503
dim=-1,
504
)
505
506
+ # Check if running on Blackwell (SM100+) for kernel tuning
507
+ self.is_blackwell = current_platform.is_device_capability_family(100)
508
+
509
def forward_native(
510
self,
511
hidden_states: torch.Tensor,
@@ -883,6 +887,7 @@ class MambaMixer2(MambaBase, CustomOp):
887
state_batch_indices=state_indices_tensor_d_input,
888
dst_state_batch_indices=state_indices_tensor_d_output,
889
out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
890
+ is_blackwell=self.is_blackwell,
891
)
892
893
def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+21
-11
@@ -286,6 +286,7 @@ def selective_state_update(
286
out=None,
287
num_accepted_tokens=None,
288
cu_seqlens=None,
289
+ is_blackwell=False,
290
):
291
"""
292
Argument:
@@ -391,17 +392,26 @@ def selective_state_update(
392
if dst_state_batch_indices is not None
393
else (0, 0)
394
)
394
- # We don't want autotune since it will overwrite the state
395
- # We instead tune by hand.
396
- BLOCK_SIZE_M, num_warps = (
397
- (32, 4)
398
- if dstate <= 16
399
- else (
400
- (16, 4)
401
- if dstate <= 32
402
- else ((8, 4) if dstate <= 64 else ((4, 4) if dstate <= 128 else ((4, 8))))
403
- )
404
- )
395
+ # We don't want autotune since it will overwrite the state.
396
+ # We instead tune by hand based on dstate.
397
+
398
+ # Default
399
+ BLOCK_SIZE_M, num_warps = 4, 8
400
+
401
+ if dstate <= 16:
402
+ BLOCK_SIZE_M, num_warps = 32, 4
403
+ elif dstate <= 32:
404
+ BLOCK_SIZE_M, num_warps = 16, 4
405
+ elif dstate <= 64:
406
+ BLOCK_SIZE_M, num_warps = 8, 4
407
+ else:
408
+ # dstate > 64
409
+ if is_blackwell:
410
+ # Optimized for B200 with dstate>64
411
+ BLOCK_SIZE_M, num_warps = 32, 8
412
+ elif dstate <= 128:
413
+ BLOCK_SIZE_M, num_warps = 4, 4
414
+
415
tie_hdim = (
416
A.stride(-1) == 0
417
and A.stride(-2) == 0