@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8f89381fc6

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 3, 20:15 UTC

[Hybrid] Warmup Mamba2 SSD kernel (#39822)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Thomas Parnell committed May 12, 2026 at 14:46 UTC 8f89381fc6b2d54591a7a560e20ee5211ce1ac33

3 files changed +110 -4

vllm/config/model.py

+2 -2

@@ -1345,7 +1345,7 @@ class ModelConfig:

1345 )

1346 raise AssertionError(f"Unsupported block type: {block_type}")

1347

1348 - def get_mamba_chunk_size(self) -> int | None:

1348 + def get_mamba_chunk_size(self) -> int:

1349 """

1350 Returns the mamba chunk size if it exists

1351 """

@@ -1356,7 +1356,7 @@ class ModelConfig:

1356 chunk_size = getattr(self.hf_text_config, "chunk_size", None)

1357

1358 # Since Mamba1 does not have a chunk notion

1359 - # we use a default chunk size of 1024.

1359 + # we use a default chunk size of 2048.

1360 if chunk_size is None:

1361 chunk_size = 2048

1362

vllm/model_executor/layers/mamba/gdn_linear_attn.py

+2 -1

@@ -397,6 +397,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):

397 )

398

399 self.chunk_gated_delta_rule = ChunkGatedDeltaRule()

400 + self._prefill_kernels_warmed_up = False

401 self.enable_packed_recurrent_decode = (

402 envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE

403 )

@@ -905,7 +906,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):

906 which has fixed kernel parameters (no autotuning), so only the

907 prefill (chunked) path needs warming up.

908 """

908 - if hasattr(self, "_prefill_kernels_warmed_up"):

909 + if self._prefill_kernels_warmed_up:

910 return

911 self._prefill_kernels_warmed_up = True

912

vllm/model_executor/layers/mamba/mamba_mixer2.py

+106 -1

@@ -14,6 +14,7 @@ from vllm.distributed import (

14 tensor_model_parallel_all_reduce,

15 )

16 from vllm.forward_context import ForwardContext, get_forward_context

17 +from vllm.logger import init_logger

18 from vllm.model_executor.custom_op import CustomOp, PluggableLayer

19 from vllm.model_executor.layers.linear import (

20 ColumnParallelLinear,

@@ -54,6 +55,8 @@ from vllm.v1.attention.backend import AttentionMetadata

55 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata

56 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

58 +logger = init_logger(__name__)

59 +

60 # Added by the IBM Team, 2024

@@ -473,6 +476,8 @@ class MambaMixer2(MambaBase, PluggableLayer):

476 intermediate_size, n_groups, self.use_rms_norm, eps=rms_norm_eps

477 )

478

479 + self._ssd_kernels_warmed_up = False

480 +

481 # - get hidden_states, B and C after depthwise convolution.

482 self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split(

483 hidden_states_B_C,

@@ -557,6 +562,104 @@ class MambaMixer2(MambaBase, PluggableLayer):

562

563 return output

564

565 + def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None:

566 + """Run a minimal SSD forward pass to trigger Triton autotuning

567 + while GPU memory is still plentiful (before SSM cache allocation).

568 + """

569 + if self._ssd_kernels_warmed_up:

570 + return

571 + self._ssd_kernels_warmed_up = True

572 + logger.info_once("Warming up Mamba2 SSD Triton kernels...")

573 +

574 + device = projected_states.device

575 + dtype = projected_states.dtype

576 +

577 + nheads = self.num_heads // self.tp_size

578 + ngroups = self.n_groups // self.tp_size

579 + headdim = self.head_dim

580 + dstate = self.ssm_state_size

581 +

582 + if self.model_config is None:

583 + return

584 + chunk_size = self.model_config.get_mamba_chunk_size()

585 +

586 + # Triton's autotuner includes tensor dtypes in its cache key,

587 + # so state_dtype must match what real inference uses.

588 + _, ssm_state_dtype = self.get_state_dtype()

589 +

590 + # SSD kernel autotune keys depend on dtype and head dimensions,

591 + # not on sequence length or batch size, so a single shape suffices.

592 + seqlen = chunk_size

593 + batch = 1

594 + nchunks = seqlen // chunk_size # = 1

595 +

596 + x = torch.randn(seqlen, nheads, headdim, device=device, dtype=dtype)

597 + dt = torch.randn(seqlen, nheads, device=device, dtype=dtype)

598 + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype)

599 + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype)

600 + cu_seqlens = torch.tensor([0, seqlen], device=device, dtype=torch.int32)

601 + cu_chunk_seqlens = torch.tensor(

602 + [i * chunk_size for i in range(nchunks + 1)],

603 + device=device,

604 + dtype=torch.int32,

605 + )

606 + last_chunk_indices = torch.tensor(

607 + [nchunks - 1], device=device, dtype=torch.int32

608 + )

609 + seq_idx = torch.zeros(nchunks, device=device, dtype=torch.int32)

610 + out = torch.empty(seqlen, nheads, headdim, device=device, dtype=dtype)

611 +

612 + # Two kernels (_state_passing_fwd, _chunk_scan_fwd) use

613 + # HAS_INITSTATES as a constexpr, producing separate compiled

614 + # binaries. Warm up both code paths so neither triggers

615 + # JIT compilation during inference.

616 + for use_initial_states in (False, True):

617 + initial_states = (

618 + torch.randn(

619 + batch,

620 + nheads,

621 + headdim,

622 + dstate,

623 + device=device,

624 + dtype=ssm_state_dtype,

625 + )

626 + if use_initial_states

627 + else None

628 + )

629 + try:

630 + mamba_chunk_scan_combined_varlen(

631 + x=x,

632 + dt=dt,

633 + A=self.A,

634 + B=B,

635 + C=C,

636 + chunk_size=chunk_size,

637 + cu_seqlens=cu_seqlens,

638 + cu_chunk_seqlens=cu_chunk_seqlens,

639 + last_chunk_indices=last_chunk_indices,

640 + seq_idx=seq_idx,

641 + out=out,

642 + D=self.D,

643 + z=None,

644 + dt_bias=self.dt_bias,

645 + initial_states=initial_states,

646 + dt_softplus=True,

647 + dt_limit=(0.0, float("inf")),

648 + state_dtype=ssm_state_dtype,

649 + )

650 + except Exception:

651 + logger.warning(

652 + "Mamba2 SSD kernel warmup failed for layer %s "

653 + "(initial_states=%s). First inference may experience "

654 + "latency spike or OOM due to autotuner.",

655 + self.prefix,

656 + use_initial_states,

657 + exc_info=True,

658 + )

659 +

660 + logger.debug("Mamba2 SSD kernel warmup completed for layer %s", self.prefix)

661 + torch.accelerator.empty_cache()

662 +

663 def conv_ssm_forward(

664 self,

665 projected_states: torch.Tensor,

@@ -608,7 +711,9 @@ class MambaMixer2(MambaBase, PluggableLayer):

711 num_decode_tokens = attn_metadata.num_decode_tokens

712

713 if attn_metadata is None:

611 - # profile run

714 + # V1 profile run -- warm up SSD kernels so that autotuning

715 + # completes before SSM cache allocation.

716 + self._warmup_ssd_kernels(projected_states)

717 hidden_states_B_C = (

718 hidden_states_B_C.transpose(0, 1).clone().transpose(0, 1)

719 ).contiguous()