diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py index 045ede4b502..da48faced90 100644 --- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py +++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py @@ -481,7 +481,7 @@ def _update_intermediate_metadata( # - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1]. # These are within bounds only when the prefill has at least # d_conv tokens; shorter sequences (e.g. small CUDA-graph warmup - # buckets) would overrun the token axis, so _ssm_prefill clamps + # buckets) would overrun the token axis, so ssm_prefill clamps # the gather positions into range. The gathered state is unused. if real_count < max_count: self._intermediate_chunk_indices_buffer[real_count:max_count].fill_(0) @@ -506,7 +506,7 @@ def _update_intermediate_metadata( else: # No extraction: fill with safe defaults for CUDA graph warmup # (same rationale as padding comment above; abs_positions=d_conv may - # exceed a sub-d_conv warmup sequence, so _ssm_prefill clamps the + # exceed a sub-d_conv warmup sequence, so ssm_prefill clamps the # gather positions into range and the gathered state is unused) self._intermediate_chunk_indices_buffer[:max_count] = 0 self._intermediate_abs_positions_buffer[:max_count] = self.d_conv diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 73e0561fdbf..e374592a125 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -9,7 +9,7 @@ import logging import math from dataclasses import dataclass, replace -from typing import List, Optional, Tuple, Union +from typing import Optional, Tuple, Union import torch import torch.nn as nn @@ -18,9 +18,7 @@ from megatron.core import parallel_state from megatron.core.inference.contexts import BaseInferenceContext, DynamicInferenceContext from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( - tensor_get_slice_after, tensor_masked_update, - tensor_merge, ) from megatron.core.inference.utils import InferenceMode from megatron.core.packed_seq_params import PackedSeqParams @@ -32,6 +30,7 @@ scatter_intermediate_ssm, ) from megatron.core.ssm.ops.mamba_ssm import selective_state_update +from megatron.core.ssm.ssm_inference import SSMDynamicInferenceMixin from megatron.core.ssm.utils import _split_tensor_factory from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.tensor_parallel.gtp_api import HAVE_GTP @@ -47,7 +46,6 @@ deprecate_inference_params, is_causal_conv1d_min_version, is_mamba_min_version, - is_using_quantization_scales, log_single_rank, make_tp_sharded_tensor_for_checkpoint, ) @@ -141,7 +139,7 @@ class MambaMixerSubmodules: out_proj: Union[ModuleSpec, type] = None -class MambaMixer(MegatronModule): +class MambaMixer(SSMDynamicInferenceMixin, MegatronModule): """ Args: config: The config of the model. @@ -490,7 +488,7 @@ def forward( if in_inference_mode and inference_context is not None: if inference_context.is_dynamic_batching(): - return self._dynamic_inference(hidden_states, inference_context) + return self.ssm_dynamic_inference(hidden_states, inference_context) else: assert inference_context.is_static_batching() assert not self.config.batch_invariant_mode, ( @@ -501,7 +499,7 @@ def forward( conv_state, ssm_state = self._get_states_from_cache(inference_context, batch) if inference_context.seqlen_offset > 0: # The states are updated inplace - out, out_bias = self._decode(hidden_states, conv_state, ssm_state) + out, out_bias = self._static_decode(hidden_states, conv_state, ssm_state) return out, out_bias zxBCdt, _ = self.in_proj(hidden_states) @@ -514,7 +512,7 @@ def forward( "Training with packed sequences is not supported " "in the non-memory-efficient code path." ) - y = self._ssm_prefill(zxBCdt, conv_state=conv_state, ssm_state=ssm_state) + y = self._static_prefill(zxBCdt, conv_state=conv_state, ssm_state=ssm_state) else: assert ssm_state is None y = self._ssm_training(zxBCdt, packed_seq_params) @@ -523,215 +521,155 @@ def forward( return out, out_bias - def _dynamic_inference(self, hidden_states: torch.Tensor, context: DynamicInferenceContext): - """ - Executes dynamic inference by separating decode and prefill requests and - running them independently. - """ - sequence_packing_available, reason_for_no_sequence_packing = ( - _check_mamba_sequence_packing_support(for_inference_not_training=True) - ) - assert sequence_packing_available, reason_for_no_sequence_packing - - # Grab standard states - conv_state, ssm_state = context.mamba_states_cache(self.layer_number - self.pp_layer_offset) - - # Fetch intermediate states for speculative decoding - # (just buffers, existing data is overwritten) - int_conv_state = None - int_ssm_state = None - if context.num_speculative_tokens > 0: - int_conv_state, int_ssm_state = context.mamba_states_cache( - self.layer_number - self.pp_layer_offset, intermediate=True - ) - - padded_dims = context.padded_batch_dimensions - token_count = padded_dims.token_count - decode_req_count = padded_dims.decode_req_count - prefill_req_count = padded_dims.prefill_req_count + # ================================================================== + # Static / eager inference + # + # These methods implement legacy static-batching inference (and the + # non-memory-efficient training prefill fallback). They are deliberately + # kept separate from the dynamic inference hooks (`ssm_decode` / + # `ssm_prefill`) so that static-batching bookkeeping does not pollute the + # dynamic inference interface defined by `SSMDynamicInferenceMixin`. + # ================================================================== + def _static_decode( + self, hidden_states, conv_state, ssm_state + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Single-token static-batching decode step (updates state in place).""" + # assert self.ngroups_local_tp == 1, "Only support ngroups=1 for inference for now" + assert hidden_states.shape[0] == 1, "Only support decoding with 1 token at a time for now" - # Input projection + # (1, b, d_model) -> (1, b, proj_dim) zxBCdt, _ = self.in_proj(hidden_states) - y_decode = None - y_prefill = None - - # Decode - if decode_req_count > 0: - # For mixed batch, the decode tokens are at the start of zxBCdt - seq_len = 1 + context.num_speculative_tokens - decode_token_count = decode_req_count * seq_len - - zxBCdt_decode = zxBCdt[:decode_token_count] if prefill_req_count > 0 else zxBCdt - - # Reshape from [N*S, 1, d] to [N, S, d] for the 3D Triton kernels - zxBCdt_decode = zxBCdt_decode.squeeze(1).view(decode_req_count, seq_len, -1) - - y_decode = self._ssm_decode( - zxBCdt_decode, - conv_state, - ssm_state, - batch_indices=context.mamba_metadata.batch_indices_decode, - intermediate_conv_state=int_conv_state, - intermediate_ssm_state=int_ssm_state, - ) - - # Flatten back to [N*S, 1, d] to match merge logic - y_decode = y_decode.view(decode_token_count, 1, -1) - - # Prefill - if prefill_req_count > 0: - if decode_req_count > 0: - # If mixed, slice the prefill portion out of zxBCdt - zxBCdt_prefill = torch.empty_like(zxBCdt) - tensor_get_slice_after( - zxBCdt, - zxBCdt_prefill, - context.mamba_metadata.device_decode_prefill, - check_bounds=False, - ) - else: - zxBCdt_prefill = zxBCdt + assert self.cp.cp_size == 1, "Context parallel not supported for Mamba inference decode" - mamba_layer_idx = context.layer_map[self.layer_number - self.pp_layer_offset - 1] - y_prefill = self._dynamic_inference_prefill( - zxBCdt_prefill, context, conv_state, ssm_state, mamba_layer_idx=mamba_layer_idx - ) - - # Merge decode and prefill results if necessary - if y_decode is not None and y_prefill is not None: - y = torch.empty( - [token_count, 1, y_prefill.shape[-1]], - dtype=y_prefill.dtype, - device=y_prefill.device, - ) - tensor_merge( - y_decode, y_prefill, context.mamba_metadata.device_decode_prefill, output_tensor=y - ) - elif y_decode is not None: - y = y_decode - elif y_prefill is not None: - y = y_prefill - else: - raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests") + # Static batching has no slot remapping, so batch_indices is None. + y = self.ssm_decode(zxBCdt, conv_state=conv_state, ssm_state=ssm_state, batch_indices=None) - # Clear the outputs for padding tokens when using quantization scales - # to avoid corrupting amax calculations - if is_using_quantization_scales(self.config): - y[context.padding_slice] = 0.0 - - # Output projection + # y has shape (1, b, d_inner), which is what out_proj expects out, out_bias = self.out_proj(y) return out, out_bias - def _dynamic_inference_prefill( + def _static_prefill( self, zxBCdt: torch.Tensor, - context: DynamicInferenceContext, - conv_state: torch.Tensor, - ssm_state: torch.Tensor, - mamba_layer_idx: Optional[int] = None, + conv_state: Optional[torch.Tensor], + ssm_state: Optional[torch.Tensor], ) -> torch.Tensor: - """Helper to run dynamic inference prefill. - - All prefill requests (including chunked prefill) are processed together - through the unified varlen path. Uses precomputed metadata from - MambaMetadata.update() to avoid .item() calls and data-dependent - control flow, enabling CUDA graph compatibility. - - When padded_prefill_count > 0 but real_prefill_count == 0 (e.g. a - decode-only rank in expert parallelism that must match a mixed CUDA - graph), this function still executes the full kernel path. - The metadata reflects zero-length sequences (cu_seqlens all equal, - batch_indices all -1) so kernels produce a zero output tensor of the - correct padded shape, which is required by the merge logic in - _dynamic_inference. - - Intermediate state extraction (for Mamba prefix caching) is performed - inside _ssm_prefill via pre-allocated output buffers, making it fully - CUDA graph compatible. """ - metadata = context.mamba_metadata + Performs single-sequence SSM prefill for static-batching inference and the + non-memory-efficient (`use_mem_eff_path=False`) training fallback. - # Use precomputed metadata (no .item() calls, no stripping). - cu_seqlens = metadata.cu_seqlens - batch_indices = metadata.batch_indices_prefill - real_token_count = metadata.real_prefill_token_count - seq_idx = metadata.seq_idx + `conv_state` / `ssm_state` are `None` for the training fallback and + non-`None` for static-batching inference (updated in place). - # Pass full padded tensor — SSM kernel uses cu_chunk_seqlens for - # boundaries and never accesses tokens beyond the last boundary. - # Output y is initialized to zeros in _ssm_prefill so padding - # positions remain zero (safe for RMSNorm and downstream ops). + Args: + zxBCdt: The input tensor of shape (l, b, d), a concatenation of + z, x, B, C, and dt projections. + conv_state: The convolution state tensor, or `None` for training. + ssm_state: The selective scan state tensor, or `None` for training. - # Prepare intermediate extraction buffers (always passed, CUDA graph compat) - slot_allocator = context.mamba_slot_allocator - intermediate_chunk_indices = metadata.intermediate_chunk_indices - intermediate_abs_positions = metadata.intermediate_abs_positions - intermediate_real_count = metadata.intermediate_real_count - intermediate_ssm_out = None - intermediate_conv_out = None - if slot_allocator is not None and mamba_layer_idx is not None: - intermediate_ssm_out = slot_allocator.intermediate_ssm_out[mamba_layer_idx] - intermediate_conv_out = slot_allocator.intermediate_conv_out[mamba_layer_idx] + Returns: + Output tensor of shape (l, b, d). + """ + # transpose: l b pd --> b l pd + zxBCdt = rearrange(zxBCdt, "l b d -> b l d").contiguous() + + # (nheads_local_tpcp) + A = -torch.exp(self.cp.get_A_log().float()) - y_prefill = self._ssm_prefill( + z, xBC, dt = torch.split( zxBCdt, - conv_state=conv_state, - ssm_state=ssm_state, - seq_idx=seq_idx, - cu_seqlens=cu_seqlens, - batch_indices=batch_indices, - intermediate_chunk_indices=intermediate_chunk_indices, - intermediate_abs_positions=intermediate_abs_positions, - intermediate_real_count=intermediate_real_count, - intermediate_ssm_out=intermediate_ssm_out, - intermediate_conv_out=intermediate_conv_out, - cu_chunk_seqlens=metadata.cu_chunk_seqlens, - last_chunk_indices=metadata.last_chunk_indices, - seq_idx_for_varlen=metadata.seq_idx_for_varlen, - cu_seqlens_list=metadata.cu_seqlens_list, - real_token_count=real_token_count, - conv_seq_idx=metadata.conv_seq_idx, - conv_seq_start=metadata.conv_seq_start, + [ + self.cp.d_inner_local_tpcp, + self.cp.d_inner_local_tpcp + 2 * self.cp.ngroups_local_tpcp * self.d_state, + self.cp.nheads_local_tpcp, + ], + dim=-1, ) - return y_prefill - - def _decode( - self, hidden_states, conv_state, ssm_state, batch_indices: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Performs inference step for decoding.""" - # assert self.ngroups_local_tp == 1, "Only support ngroups=1 for inference for now" - is_dynamic_batching = batch_indices is not None + # Compute short convolution (single-sequence / non-varlen). + xBC = rearrange(xBC, "b l d -> b d l").contiguous() + if conv_state is not None: + # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv + # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. + conv_state.copy_(F.pad(xBC, (self.d_conv - xBC.shape[-1], 0))) # Update state (B D W) - if not is_dynamic_batching: - assert ( - hidden_states.shape[0] == 1 - ), "Only support decoding with 1 token at a time for now" + seqlen = xBC.size(2) + if causal_conv1d_fn is None: + xBC = self.act(self.cp.conv1d(xBC)[..., :seqlen]) + else: + assert self.activation in ["silu", "swish"] + xBC = causal_conv1d_fn( + x=xBC, + weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), + bias=self.cp.get_conv1d_bias(), + activation=self.activation, + ) + xBC = rearrange(xBC, "b d l -> b l d").contiguous() - # (1, b, d_model) -> (1, b, proj_dim) - zxBCdt, _ = self.in_proj(hidden_states) + x, B, C = torch.split( + xBC, + [ + self.cp.d_inner_local_tpcp, + self.cp.ngroups_local_tpcp * self.d_state, + self.cp.ngroups_local_tpcp * self.d_state, + ], + dim=-1, + ) - # Make batch size leading dimension since that is 1 - if is_dynamic_batching: - zxBCdt = zxBCdt.transpose(0, 1) + # TODO Vijay: fuse most of the transposes with the GEMMS + x = rearrange(x, "b l (h p) -> b l h p", p=self.headdim).contiguous() + dt = dt.contiguous() + B = rearrange(B, "b l (g n) -> b l g n", n=self.d_state).contiguous() + C = rearrange(C, "b l (g n) -> b l g n", n=self.d_state).contiguous() + z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous() - assert self.cp.cp_size == 1, "Context parallel not supported for Mamba inferenece decode" + # If `rmsnorm == False`, then the norm inside `mamba_chunk_scan_combined` will be used. + # In this case, if `cp_size > 1` then that norm could be performed on less heads than if + # `cp_size == 1` (groups of heads can be sharded across CP ranks), which would be + # mathematically incorrect, and potentially arithmetically unstable. + assert ( + self.cp.cp_size == 1 or self.rmsnorm + ), "Context parallel not supported for use_mem_eff_path==False and rmsnorm==False" - y = self._ssm_decode( - zxBCdt, conv_state=conv_state, ssm_state=ssm_state, batch_indices=batch_indices + initial_ssm_state = None + state_dtype_kwarg = ( + {"state_dtype": self.mamba_training_ssm_states_dtype} if MAMBA_HAS_STATE_DTYPE else {} + ) + y = mamba_chunk_scan_combined( + x, + dt, + A, + B, + C, + self.chunk_size, + D=( + rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim) + if self.D_has_hdim + else self.cp.get_D() + ), + z=z if not self.rmsnorm else None, + dt_bias=self.cp.get_dt_bias().float(), + dt_softplus=True, + return_final_states=ssm_state is not None, + initial_states=initial_ssm_state, + **state_dtype_kwarg, ) - # Restore sequence length as first dimension - if is_dynamic_batching: - y = y.transpose(0, 1) + if ssm_state is not None: + y, last_state = y + ssm_state.copy_(last_state) - # y has shape (1, b, d_inner), which is what out_proj expects - out, out_bias = self.out_proj(y) + y = rearrange(y, "b l h p -> l b (h p)").contiguous() + y = self.cp.post_conv_ssm(y) - return out, out_bias + if self.rmsnorm: + z = rearrange(z, "b l h p -> l b (h p)").contiguous() + z = self.cp.post_conv_ssm(z) + y = self.norm(y, z) + + return y def _ssm_training( self, zxBCdt: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None @@ -789,63 +727,70 @@ def _ssm_training( return y - def _ssm_prefill( + def ssm_prefill( self, zxBCdt: torch.Tensor, - conv_state: Optional[torch.Tensor], - ssm_state: Optional[torch.Tensor], - seq_idx: Optional[torch.Tensor] = None, - cu_seqlens: Optional[torch.Tensor] = None, - batch_indices: Optional[torch.Tensor] = None, - intermediate_chunk_indices: Optional[torch.Tensor] = None, - intermediate_abs_positions: Optional[torch.Tensor] = None, - intermediate_real_count: Optional[torch.Tensor] = None, - intermediate_ssm_out: Optional[torch.Tensor] = None, - intermediate_conv_out: Optional[torch.Tensor] = None, - cu_chunk_seqlens: Optional[torch.Tensor] = None, - last_chunk_indices: Optional[torch.Tensor] = None, - seq_idx_for_varlen: Optional[torch.Tensor] = None, - cu_seqlens_list: Optional[List[int]] = None, - real_token_count: Optional[int] = None, - conv_seq_idx: Optional[torch.Tensor] = None, - conv_seq_start: Optional[torch.Tensor] = None, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context: DynamicInferenceContext, ) -> torch.Tensor: """ - Performs SSM computation for inference prefill step. + Performs the variable-length SSM prefill for all dynamic-batching prefill requests. + + All varlen metadata (cu_seqlens, seq_idx, batch_indices, chunk boundaries, + intermediate extraction buffers, etc.) is read directly from + `context.mamba_metadata` / `context.mamba_slot_allocator` -- there is no + intermediate layer that unpacks the metadata into a long argument list. All + prefill requests (including chunked prefill) are processed together through a + single varlen kernel call; the precomputed metadata avoids `.item()` calls + and data-dependent control flow, enabling CUDA graph compatibility. + Intermediate state extraction (for Mamba prefix caching) is performed via + pre-allocated output buffers, also CUDA graph compatible. + + When padded_prefill_count > 0 but real_prefill_count == 0 (e.g. a decode-only + rank in expert parallelism that must match a mixed CUDA graph), the full kernel + path still runs; the metadata reflects zero-length sequences (cu_seqlens all + equal, batch_indices all -1) so the kernels produce a correctly-shaped zero + output tensor, which is required by the merge logic in `ssm_dynamic_inference`. Args: zxBCdt: The input tensor of shape (l, b, d), which is a concatenation of z, x, B, C, and dt projections. conv_state: The convolution state tensor for inference. ssm_state: The selective scan state tensor for inference. - seq_idx: A map from token index to request index for variable-length sequences. - cu_seqlens: Cumulative sequence lengths for variable-length sequences. - batch_indices: A map from batch id to position in the Mamba state tensors for - dynamic inference. - intermediate_chunk_indices: Pre-allocated tensor of chunk indices for - intermediate state extraction (fixed size, padded with 0). - intermediate_abs_positions: Pre-allocated tensor of absolute token - positions for conv state extraction (fixed size, padded with d_conv). - intermediate_real_count: int32[1] GPU tensor holding the number of - meaningful entries in the intermediate buffers this step. Read - inside the Triton scatter kernels so padded slots cost nothing. - intermediate_ssm_out: Output buffer for extracted SSM states - [max_intermediate_count, *ssm_shape]. - intermediate_conv_out: Output buffer for extracted conv states - [max_intermediate_count, *conv_shape]. - cu_chunk_seqlens: Precomputed chunk boundaries from MambaMetadata. - last_chunk_indices: Precomputed last chunk index per sequence. - seq_idx_for_varlen: Precomputed request ID per chunk. - cu_seqlens_list: Python list of cumulative sequence lengths (avoids .item()). - real_token_count: Number of real (non-padding) tokens. - conv_seq_idx: Precomputed per-token request ID for Triton conv1d. - conv_seq_start: Precomputed per-token request start for Triton conv1d. + context: The dynamic inference context supplying all varlen metadata. Returns: - Output tensor of shape (l, b, d). Intermediate states (if any) are - written directly to intermediate_ssm_out and intermediate_conv_out. + Output tensor of shape (l, b, d). Intermediate states (if any) are written + directly into the slot-allocator buffers held by `context`. """ - is_dynamic_batching = seq_idx is not None + assert ( + self.cp.cp_size == 1 + ), "Context parallel is not supported for MambaMixer dynamic inference prefill" + + metadata = context.mamba_metadata + slot_allocator = context.mamba_slot_allocator + + seq_idx = metadata.seq_idx + cu_seqlens = metadata.cu_seqlens + batch_indices = metadata.batch_indices_prefill + intermediate_chunk_indices = metadata.intermediate_chunk_indices + intermediate_abs_positions = metadata.intermediate_abs_positions + intermediate_real_count = metadata.intermediate_real_count + cu_chunk_seqlens = metadata.cu_chunk_seqlens + last_chunk_indices = metadata.last_chunk_indices + seq_idx_for_varlen = metadata.seq_idx_for_varlen + conv_seq_idx = metadata.conv_seq_idx + conv_seq_start = metadata.conv_seq_start + + # Wire the per-layer intermediate extraction buffers (prefix caching) when a + # slot allocator is present; otherwise extraction is disabled below. + intermediate_ssm_out = None + intermediate_conv_out = None + if slot_allocator is not None: + mamba_layer_idx = context.layer_map[self.layer_number - self.pp_layer_offset - 1] + intermediate_ssm_out = slot_allocator.intermediate_ssm_out[mamba_layer_idx] + intermediate_conv_out = slot_allocator.intermediate_conv_out[mamba_layer_idx] # transpose: l b pd --> b l pd zxBCdt = rearrange(zxBCdt, "l b d -> b l d").contiguous() @@ -863,74 +808,47 @@ def _ssm_prefill( dim=-1, ) - # Compute short convolution - xBC_pre_conv = None - if conv_state is not None and is_dynamic_batching: - assert batch_indices is not None - - # Extract initial conv states BEFORE saving new ones. - # causal_conv1d_varlen_states computes the final conv state from the - # input sequence and tensor_masked_update writes it into the conv_state - # buffer. If we read initial_conv_states after this write, restored - # requests see their own newly-computed states instead of the cached - # initial states from a previous request, corrupting the conv output. - initial_conv_states = conv_state[batch_indices, :, 1:] - - # Save final conv states from the input sequence - conv_varlen_states = causal_conv1d_varlen_states( - xBC.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] - ) - tensor_masked_update(conv_state, batch_indices, conv_varlen_states) + # Compute short convolution (unified varlen path over all prefill requests). + assert batch_indices is not None - # Conv state dtype might differ from params dtype, so cast xBC and weight / bias - # tensors to the conv state dtype for causal_conv1d_varlen_fn and then cast xBC - # back to the original dtype - xBC_dtype = xBC.dtype - conv_state_dtype = conv_state.dtype + # Extract initial conv states BEFORE saving new ones. + # causal_conv1d_varlen_states computes the final conv state from the + # input sequence and tensor_masked_update writes it into the conv_state + # buffer. If we read initial_conv_states after this write, restored + # requests see their own newly-computed states instead of the cached + # initial states from a previous request, corrupting the conv output. + initial_conv_states = conv_state[batch_indices, :, 1:] - xBC = xBC.to(conv_state_dtype) - conv_weight = rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w").to( - conv_state_dtype - ) - conv_bias = self.cp.get_conv1d_bias().to(conv_state_dtype) + # Save final conv states from the input sequence + conv_varlen_states = causal_conv1d_varlen_states( + xBC.squeeze(0), cu_seqlens, state_len=conv_state.shape[-1] + ) + tensor_masked_update(conv_state, batch_indices, conv_varlen_states) - xBC_pre_conv = xBC if intermediate_conv_out is not None else None - from megatron.core.ssm.ops.causal_conv1d_varlen import causal_conv1d_varlen_fn + # Conv state dtype might differ from params dtype, so cast xBC and weight / bias + # tensors to the conv state dtype for causal_conv1d_varlen_fn and then cast xBC + # back to the original dtype + xBC_dtype = xBC.dtype + conv_state_dtype = conv_state.dtype - xBC_out = causal_conv1d_varlen_fn( - x=xBC.squeeze(0).contiguous(), - weight=conv_weight, - bias=conv_bias, - cu_seqlens=cu_seqlens, - initial_states=initial_conv_states, - activation=self.activation, - precomputed_seq_idx=conv_seq_idx, - precomputed_seq_start=conv_seq_start, - ) - xBC = xBC_out.to(xBC_dtype).unsqueeze(0) - else: - # Non-dynamic-batching path (static batching / training fallback) - xBC = rearrange(xBC, "b l d -> b d l").contiguous() - if conv_state is not None: - # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv - # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise. - conv_state.copy_( - F.pad(xBC, (self.d_conv - xBC.shape[-1], 0)) - ) # Update state (B D W) - - seqlen = xBC.size(2) - if causal_conv1d_fn is None: - xBC = self.act(self.cp.conv1d(xBC)[..., :seqlen]) - else: - assert self.activation in ["silu", "swish"] - xBC = causal_conv1d_fn( - x=xBC, - weight=rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w"), - bias=self.cp.get_conv1d_bias(), - activation=self.activation, - seq_idx=seq_idx, - ) - xBC = rearrange(xBC, "b d l -> b l d").contiguous() + xBC = xBC.to(conv_state_dtype) + conv_weight = rearrange(self.cp.get_conv1d_weight(), "d 1 w -> d w").to(conv_state_dtype) + conv_bias = self.cp.get_conv1d_bias().to(conv_state_dtype) + + xBC_pre_conv = xBC if intermediate_conv_out is not None else None + from megatron.core.ssm.ops.causal_conv1d_varlen import causal_conv1d_varlen_fn + + xBC_out = causal_conv1d_varlen_fn( + x=xBC.squeeze(0).contiguous(), + weight=conv_weight, + bias=conv_bias, + cu_seqlens=cu_seqlens, + initial_states=initial_conv_states, + activation=self.activation, + precomputed_seq_idx=conv_seq_idx, + precomputed_seq_start=conv_seq_start, + ) + xBC = xBC_out.to(xBC_dtype).unsqueeze(0) x, B, C = torch.split( xBC, @@ -957,173 +875,141 @@ def _ssm_prefill( self.cp.cp_size == 1 or self.rmsnorm ), "Context parallel not supported for use_mem_eff_path==False and rmsnorm==False" - if is_dynamic_batching: - # Unified varlen SSM path: all prefill requests through single kernel call - initial_ssm_state = ssm_state[batch_indices] - - x = x.squeeze(0) - dt = dt.squeeze(0) - A = A.squeeze(0) - B = B.squeeze(0) - C = C.squeeze(0) - z = z.squeeze(0) - # Initialize with zeros so padding positions (beyond cu_chunk_seqlens - # boundaries) remain zero, which is safe for RMSNorm and downstream ops. - y = torch.zeros_like(x) - - if cu_chunk_seqlens is not None: - # Use precomputed chunk metadata (CUDA graph compatible, no .item()) - pass - else: - # Fallback: build chunk metadata from cu_seqlens (non-precomputed) - chunk_boundaries = [0] - last_chunk_indices_list = [] - num_seqs = cu_seqlens.numel() - 1 - for i in range(num_seqs): - start = cu_seqlens[i].item() - end = cu_seqlens[i + 1].item() - pos = start + self.chunk_size - while pos < end: - chunk_boundaries.append(pos) - pos += self.chunk_size - chunk_boundaries.append(end) - last_chunk_indices_list.append(len(chunk_boundaries) - 2) - - cu_chunk_seqlens = cu_seqlens.new_tensor(chunk_boundaries) - last_chunk_indices = cu_seqlens.new_tensor(last_chunk_indices_list) - - seq_idx_for_varlen = None - if seq_idx is not None: - chunk_starts = cu_chunk_seqlens[:-1] - seq_idx_for_varlen = seq_idx[0, chunk_starts].contiguous() - - # Batch-invariant decode replays the partial prefill tail, so keep - # the cached SSM state at the last complete chunk boundary. - if self.config.batch_invariant_mode: - prefill_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) - tail_lens = prefill_lens % self.chunk_size - has_boundary = prefill_lens >= self.chunk_size - # A partial tail uses the preceding full chunk's state. - boundary_chunk_indices = ( - last_chunk_indices.to(torch.long) - (tail_lens > 0).to(torch.long) - ).clamp(min=0) - - # Extraction is enabled when the slot allocator wired buffers in via - # the caller. When enabled, the chunk scan returns its raw states so - # our Triton kernels do a fused gather+conditional-scatter directly, - # skipping the dense intermediate tensor and the padded-slot writes. - extract_intermediates = ( - not self.config.batch_invariant_mode - and intermediate_chunk_indices is not None - and intermediate_ssm_out is not None - ) - ssm_varlen_result = mamba_chunk_scan_combined_varlen( - x=x, - dt=dt, - A=A, - B=B, - C=C, - chunk_size=self.chunk_size, - cu_chunk_seqlens=cu_chunk_seqlens, - last_chunk_indices=last_chunk_indices, - seq_idx=seq_idx_for_varlen, - out=y, - D=( - rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim) - if self.D_has_hdim - else self.cp.get_D() - ), - z=z if (self.config.batch_invariant_mode or not self.rmsnorm) else None, - dt_bias=self.cp.get_dt_bias().float(), - initial_states=initial_ssm_state, - return_raw_states=self.config.batch_invariant_mode or extract_intermediates, - dt_softplus=True, - dt_limit=(0.0, float("inf")), - state_dtype=ssm_state.dtype, - ) - - if self.config.batch_invariant_mode or extract_intermediates: - ssm_varlen_states, raw_ssm_states = ssm_varlen_result - else: - ssm_varlen_states = ssm_varlen_result - raw_ssm_states = None + # Unified varlen SSM path: all prefill requests through single kernel call + initial_ssm_state = ssm_state[batch_indices] + + x = x.squeeze(0) + dt = dt.squeeze(0) + A = A.squeeze(0) + B = B.squeeze(0) + C = C.squeeze(0) + z = z.squeeze(0) + # Initialize with zeros so padding positions (beyond cu_chunk_seqlens + # boundaries) remain zero, which is safe for RMSNorm and downstream ops. + y = torch.zeros_like(x) + + if cu_chunk_seqlens is not None: + # Use precomputed chunk metadata (CUDA graph compatible, no .item()) + pass + else: + # Fallback: build chunk metadata from cu_seqlens (non-precomputed) + chunk_boundaries = [0] + last_chunk_indices_list = [] + num_seqs = cu_seqlens.numel() - 1 + for i in range(num_seqs): + start = cu_seqlens[i].item() + end = cu_seqlens[i + 1].item() + pos = start + self.chunk_size + while pos < end: + chunk_boundaries.append(pos) + pos += self.chunk_size + chunk_boundaries.append(end) + last_chunk_indices_list.append(len(chunk_boundaries) - 2) + + cu_chunk_seqlens = cu_seqlens.new_tensor(chunk_boundaries) + last_chunk_indices = cu_seqlens.new_tensor(last_chunk_indices_list) + + seq_idx_for_varlen = None + if seq_idx is not None: + chunk_starts = cu_chunk_seqlens[:-1] + seq_idx_for_varlen = seq_idx[0, chunk_starts].contiguous() + + # Batch-invariant decode replays the partial prefill tail, so keep + # the cached SSM state at the last complete chunk boundary. + if self.config.batch_invariant_mode: + prefill_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) + tail_lens = prefill_lens % self.chunk_size + has_boundary = prefill_lens >= self.chunk_size + # A partial tail uses the preceding full chunk's state. + boundary_chunk_indices = ( + last_chunk_indices.to(torch.long) - (tail_lens > 0).to(torch.long) + ).clamp(min=0) + + # Extraction is enabled when the slot allocator wired buffers in via + # the caller. When enabled, the chunk scan returns its raw states so + # our Triton kernels do a fused gather+conditional-scatter directly, + # skipping the dense intermediate tensor and the padded-slot writes. + extract_intermediates = ( + not self.config.batch_invariant_mode + and intermediate_chunk_indices is not None + and intermediate_ssm_out is not None + ) + ssm_varlen_result = mamba_chunk_scan_combined_varlen( + x=x, + dt=dt, + A=A, + B=B, + C=C, + chunk_size=self.chunk_size, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx_for_varlen, + out=y, + D=( + rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim) + if self.D_has_hdim + else self.cp.get_D() + ), + z=z if (self.config.batch_invariant_mode or not self.rmsnorm) else None, + dt_bias=self.cp.get_dt_bias().float(), + initial_states=initial_ssm_state, + return_raw_states=self.config.batch_invariant_mode or extract_intermediates, + dt_softplus=True, + dt_limit=(0.0, float("inf")), + state_dtype=ssm_state.dtype, + ) - y = y.unsqueeze(0) - z = z.unsqueeze(0) + if self.config.batch_invariant_mode or extract_intermediates: + ssm_varlen_states, raw_ssm_states = ssm_varlen_result + else: + ssm_varlen_states = ssm_varlen_result + raw_ssm_states = None - if self.config.batch_invariant_mode: - boundary_mask = has_boundary.view(-1, 1, 1, 1) - cache_states = torch.where( - boundary_mask, raw_ssm_states[boundary_chunk_indices], initial_ssm_state - ) - else: - cache_states = ssm_varlen_states - - tensor_masked_update(ssm_state, batch_indices, cache_states) - if self.config.batch_invariant_mode: - self._get_batch_invariant_decoder().seed( - x, - z.squeeze(0), - dt, - B, - C, - cu_seqlens, - batch_indices, - max_requests=ssm_state.shape[0], - ) + y = y.unsqueeze(0) + z = z.unsqueeze(0) - if extract_intermediates: - # Fused gather+conditional-scatter for SSM: read row - # raw_ssm_states[chunk_indices[i]] into intermediate_ssm_out[i], - # only for i < real_count. - scatter_intermediate_ssm( - raw_ssm_states, - intermediate_chunk_indices, - intermediate_real_count, - intermediate_ssm_out, - ) - # Same pattern for conv: gather a length-d_conv window ending at - # abs_positions[i] (clamped into the valid token range) from - # xBC_pre_conv and scatter (transposed) into intermediate_conv_out[i], - # only for i < real_count. - scatter_intermediate_conv( - xBC_pre_conv, - intermediate_abs_positions, - intermediate_real_count, - intermediate_conv_out, - d_conv=intermediate_conv_out.shape[-1], - ) - else: - # Non-dynamic-batching path (static batching) - initial_ssm_state = None - state_dtype_kwarg = ( - {"state_dtype": self.mamba_training_ssm_states_dtype} - if MAMBA_HAS_STATE_DTYPE - else {} + if self.config.batch_invariant_mode: + boundary_mask = has_boundary.view(-1, 1, 1, 1) + cache_states = torch.where( + boundary_mask, raw_ssm_states[boundary_chunk_indices], initial_ssm_state ) - y = mamba_chunk_scan_combined( + else: + cache_states = ssm_varlen_states + + tensor_masked_update(ssm_state, batch_indices, cache_states) + if self.config.batch_invariant_mode: + self._get_batch_invariant_decoder().seed( x, + z.squeeze(0), dt, - A, B, C, - self.chunk_size, - D=( - rearrange(self.cp.get_D().float(), "(h p) -> h p", p=self.headdim) - if self.D_has_hdim - else self.cp.get_D() - ), - z=z if not self.rmsnorm else None, - dt_bias=self.cp.get_dt_bias().float(), - dt_softplus=True, - return_final_states=ssm_state is not None, - initial_states=initial_ssm_state, - **state_dtype_kwarg, + cu_seqlens, + batch_indices, + max_requests=ssm_state.shape[0], ) - if ssm_state is not None: - y, last_state = y - ssm_state.copy_(last_state) + if extract_intermediates: + # Fused gather+conditional-scatter for SSM: read row + # raw_ssm_states[chunk_indices[i]] into intermediate_ssm_out[i], + # only for i < real_count. + scatter_intermediate_ssm( + raw_ssm_states, + intermediate_chunk_indices, + intermediate_real_count, + intermediate_ssm_out, + ) + # Same pattern for conv: gather a length-d_conv window ending at + # abs_positions[i] (clamped into the valid token range) from + # xBC_pre_conv and scatter (transposed) into intermediate_conv_out[i], + # only for i < real_count. + scatter_intermediate_conv( + xBC_pre_conv, + intermediate_abs_positions, + intermediate_real_count, + intermediate_conv_out, + d_conv=intermediate_conv_out.shape[-1], + ) y = rearrange(y, "b l h p -> l b (h p)").contiguous() y = self.cp.post_conv_ssm(y) @@ -1167,7 +1053,7 @@ def train(self, mode: bool = True): self._A_neg_exp_cache_stale = True return super().train(mode) - def _ssm_decode( + def ssm_decode( self, zxBCdt: torch.Tensor, conv_state: torch.Tensor, diff --git a/megatron/core/ssm/ssm_inference.py b/megatron/core/ssm/ssm_inference.py new file mode 100644 index 00000000000..a850b51ebfa --- /dev/null +++ b/megatron/core/ssm/ssm_inference.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared dynamic-batching inference scaffolding for linear-attention mixers. + +A growing family of mixers in Megatron behave like "linear attention" / SSM +recurrences for inference purposes: they carry a small per-request recurrent +state (a short-convolution state plus a matrix-valued SSM state) instead of a +growing KV cache. Mamba was the first; Gated Delta Net / Gated Delta Product +(GDP) and friends are the same shape of computation with different kernels. + +All of these variants share an *identical* request-level control flow for the +dynamic inference engine: + + 1. Fetch this layer's (conv_state, ssm_state) slabs from the context. + 2. Project the packed input (`in_proj`). + 3. Split the packed batch into a decode partition (1 token per request, + placed first) and a prefill partition (variable length, placed after). + The kernels cannot mix the two, so they run independently. + 4. Run the decode and prefill kernels on their respective partitions. + 5. Merge the two partitions back into packed token order. + 6. Apply the output projection (`out_proj`). + +Only the kernels in step 4 differ between variants. This mixin owns the shared +control flow (steps 1-3, 5, 6 and the orchestration) and delegates the +variant-specific work to two hooks, `ssm_decode` and `ssm_prefill`. New +linear-attention variants should subclass this mixin and implement those two +hooks rather than re-deriving the decode/prefill bookkeeping. + +Both hooks are given the `DynamicInferenceContext` directly and read whatever +per-step metadata they need from `context.mamba_metadata` / +`context.mamba_slot_allocator` themselves; there is deliberately no +intermediate "unpack the metadata into a long argument list" layer. + +Speculative decoding is supported by the shared orchestration: the decode path +reshapes tokens into `[batch, seq_len, d]`, fetches intermediate state buffers +from the context, and passes them to `ssm_decode`. Variants that do not yet +support speculative decoding should assert `seq_len == 1` inside their +`ssm_decode` implementation. + +Chunked prefill and prefix caching are handled entirely inside `ssm_prefill` +via `context.mamba_metadata` and `context.mamba_slot_allocator`; the mixin +orchestration is unaware of them. + +Note: static-batching ("legacy") inference is intentionally *not* part of this +interface. Concrete mixers keep any static/eager inference path separate so +it does not pollute the dynamic decode/prefill hooks defined here. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from megatron.core.inference.contexts import DynamicInferenceContext +from megatron.core.inference.contexts.attention_context.triton.tensor_ops import ( + tensor_get_slice_after, + tensor_merge, +) +from megatron.core.utils import is_using_quantization_scales + + +class SSMDynamicInferenceMixin: + """Mixin providing the shared decode/prefill orchestration for the dynamic + inference engine. Concrete mixers implement the two `ssm_*` hooks below.""" + + # ------------------------------------------------------------------ + # Hooks implemented by concrete mixers. + # ------------------------------------------------------------------ + def ssm_decode( + self, + zxBCdt: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + batch_indices: torch.Tensor, + intermediate_conv_state: torch.Tensor = None, + intermediate_ssm_state: torch.Tensor = None, + ) -> torch.Tensor: + """Run the single-token-per-request decode kernels. + + Args: + zxBCdt: `[decode_req_count, seq_len, proj_dim]` projected decode tokens, + where `seq_len = 1 + num_speculative_tokens`. + conv_state: `[num_slots, conv_channels, d_conv]` conv state cache. + ssm_state: `[num_slots, *ssm_shape]` SSM state cache. + batch_indices: `[decode_req_count]` slot index per decode request + (`-1` marks padding slots). + intermediate_conv_state: Optional buffer for storing conv states at + intermediate sequence steps (speculative decoding). + intermediate_ssm_state: Optional buffer for storing SSM states at + intermediate sequence steps (speculative decoding). + + Returns `[decode_req_count, seq_len, d_inner]`; updates state in place. + Variants that do not yet support speculative decoding should assert + `seq_len == 1` inside their implementation. + """ + raise NotImplementedError + + def ssm_prefill( + self, + zxBCdt: torch.Tensor, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + context: DynamicInferenceContext, + ) -> torch.Tensor: + """Run the variable-length prefill kernels for all prefill requests. + + The implementation reads its varlen metadata (`cu_seqlens`, + `batch_indices_prefill`, `seq_idx`, chunk boundaries, intermediate + extraction buffers, etc.) directly from `context.mamba_metadata` and + `context.mamba_slot_allocator` and processes every prefill request in + one varlen call, writing the resulting final states back into the caches. + + Returns `[prefill_token_count, 1, d_inner]`; updates state in place. + """ + raise NotImplementedError + + # ------------------------------------------------------------------ + # Shared orchestration. + # ------------------------------------------------------------------ + def ssm_dynamic_inference( + self, hidden_states: torch.Tensor, context: DynamicInferenceContext + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Execute one dynamic inference step for a linear-attention mixer. + + Separates decode and prefill requests, runs them through the + variant-specific kernels independently, and merges the results back + into packed token order. + """ + # Grab standard states. + conv_state, ssm_state = context.mamba_states_cache(self.layer_number - self.pp_layer_offset) + + # Fetch intermediate state buffers for speculative decoding. + # These are pre-allocated output buffers; existing data is overwritten. + int_conv_state = None + int_ssm_state = None + if context.num_speculative_tokens > 0: + int_conv_state, int_ssm_state = context.mamba_states_cache( + self.layer_number - self.pp_layer_offset, intermediate=True + ) + + padded_dims = context.padded_batch_dimensions + token_count = padded_dims.token_count + decode_req_count = padded_dims.decode_req_count + prefill_req_count = padded_dims.prefill_req_count + + # Input projection over the full packed batch. + zxBCdt, _ = self.in_proj(hidden_states) + + y_decode = None + y_prefill = None + + # --- Decode partition (placed first in the packed batch) --------- + if decode_req_count > 0: + seq_len = 1 + context.num_speculative_tokens + decode_token_count = decode_req_count * seq_len + zxBCdt_decode = zxBCdt[:decode_token_count] if prefill_req_count > 0 else zxBCdt + # Reshape from [N*S, 1, d] to [N, S, d] for the decode kernels. + zxBCdt_decode = zxBCdt_decode.squeeze(1).view(decode_req_count, seq_len, -1) + y_decode = self.ssm_decode( + zxBCdt_decode, + conv_state, + ssm_state, + batch_indices=context.mamba_metadata.batch_indices_decode, + intermediate_conv_state=int_conv_state, + intermediate_ssm_state=int_ssm_state, + ) + # Flatten back to [N*S, 1, d] to match the merge logic. + y_decode = y_decode.view(decode_token_count, 1, -1) + + # --- Prefill partition ------------------------------------------- + if prefill_req_count > 0: + if decode_req_count > 0: + # Mixed batch: gather the prefill tokens out of the packed tensor. + zxBCdt_prefill = torch.empty_like(zxBCdt) + tensor_get_slice_after( + zxBCdt, + zxBCdt_prefill, + context.mamba_metadata.device_decode_prefill, + check_bounds=False, + ) + else: + zxBCdt_prefill = zxBCdt + y_prefill = self.ssm_prefill(zxBCdt_prefill, conv_state, ssm_state, context) + + # --- Merge back into packed token order -------------------------- + if y_decode is not None and y_prefill is not None: + y = torch.empty( + [token_count, 1, y_prefill.shape[-1]], + dtype=y_prefill.dtype, + device=y_prefill.device, + ) + tensor_merge( + y_decode, y_prefill, context.mamba_metadata.device_decode_prefill, output_tensor=y + ) + elif y_decode is not None: + y = y_decode + elif y_prefill is not None: + y = y_prefill + else: + raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests") + + # Zero padding positions to avoid corrupting quantization amax calculations. + if is_using_quantization_scales(self.config): + y[context.padding_slice] = 0.0 + + return self.out_proj(y) diff --git a/tests/unit_tests/ssm/ops/test_ssm_kernel.py b/tests/unit_tests/ssm/ops/test_ssm_kernel.py index 62c25f43f0f..5067b06028a 100644 --- a/tests/unit_tests/ssm/ops/test_ssm_kernel.py +++ b/tests/unit_tests/ssm/ops/test_ssm_kernel.py @@ -119,17 +119,18 @@ def setUp(self): self.mixer.D = nn.Parameter(torch.ones(self.nheads, device=self.device)) # Bind methods - self.mixer._ssm_prefill = MambaMixer._ssm_prefill.__get__(self.mixer, MambaMixer) - self.mixer._ssm_decode = MambaMixer._ssm_decode.__get__(self.mixer, MambaMixer) + self.mixer.ssm_prefill = MambaMixer.ssm_prefill.__get__(self.mixer, MambaMixer) + self.mixer.ssm_decode = MambaMixer.ssm_decode.__get__(self.mixer, MambaMixer) def test_ssm_prefill_padding_isolation(self): """ Tests that ssm_prefill only updates states for the real request and that padding request states remain untouched. - _ssm_prefill expects inputs pre-stripped to real tokens only - (stripping is done by _dynamic_inference_prefill). This test - passes only the real tokens and verifies that only the active + ssm_prefill reads all varlen metadata from the DynamicInferenceContext + and expects `zxBCdt` pre-stripped to real tokens only (stripping is + done upstream). This test passes only the real tokens, wires the + metadata through a mock context, and verifies that only the active request's state is modified. """ num_requests = 48 @@ -153,15 +154,29 @@ def test_ssm_prefill_padding_isolation(self): num_requests, self.nheads, self.headdim, self.d_state, device=self.device ) - # Run - self.mixer.norm = MagicMock(side_effect=lambda x, z: x * z) - output = self.mixer._ssm_prefill( - zxBCdt=zxBCdt, - conv_state=conv_state, - ssm_state=ssm_state, + # Mock the dynamic inference context. Leaving the chunk metadata (and + # extraction buffers) unset exercises the non-precomputed fallback path, + # which rebuilds chunk boundaries from cu_seqlens; no slot allocator means + # intermediate-state extraction (prefix caching) is disabled. + mamba_metadata = SimpleNamespace( seq_idx=seq_idx, cu_seqlens=cu_seqlens, - batch_indices=batch_indices, + batch_indices_prefill=batch_indices, + intermediate_chunk_indices=None, + intermediate_abs_positions=None, + intermediate_real_count=None, + cu_chunk_seqlens=None, + last_chunk_indices=None, + seq_idx_for_varlen=None, + conv_seq_idx=None, + conv_seq_start=None, + ) + context = SimpleNamespace(mamba_metadata=mamba_metadata, mamba_slot_allocator=None) + + # Run + self.mixer.norm = MagicMock(side_effect=lambda x, z: x * z) + output = self.mixer.ssm_prefill( + zxBCdt=zxBCdt, conv_state=conv_state, ssm_state=ssm_state, context=context ) # Output should have real_seq_len tokens