diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index b8d49db97..286c51de9 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -534,6 +534,9 @@ class MambaMixer2(torch.nn.Module): layer_cache, MambaPool.SpeculativeState ), "layer_cache must be SpeculativeState for speculative decoding" draft_token_num = metadata.draft_token_num + self.intermediate_state_indices = torch.arange( + num_decodes, dtype=torch.int32, device=state_indices_tensor_d.device + ) # Reshape for batch processing hidden_states_B_C_d_reshaped = hidden_states_B_C_d.view( @@ -548,6 +551,7 @@ class MambaMixer2(torch.nn.Module): self.activation, conv_state_indices=state_indices_tensor_d[:num_decodes], intermediate_conv_window=layer_cache.intermediate_conv_window[0], + intermediate_state_indices=self.intermediate_state_indices, retrieve_next_token=metadata.retrieve_next_token, retrieve_next_sibling=metadata.retrieve_next_sibling, retrieve_parent_token=metadata.retrieve_parent_token, @@ -621,6 +625,7 @@ class MambaMixer2(torch.nn.Module): intermediate_states_buffer=layer_cache.intermediate_ssm, cache_steps=draft_token_num, retrieve_parent_token=metadata.retrieve_parent_token, + intermediate_state_indices=self.intermediate_state_indices, ) else: selective_state_update( diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index 97039d675..f238d51b4 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py @@ -56,6 +56,14 @@ else: is not None } ) +@triton.heuristics( + { + "HAS_INTERMEDIATE_STATE_INDICES": lambda args: args[ + "intermediate_state_indices_ptr" + ] + is not None + } +) @triton.jit(do_not_specialize=["T"]) def _selective_scan_update_kernel( # Pointers to matrices @@ -74,6 +82,7 @@ def _selective_scan_update_kernel( intermediate_states_buffer, cache_steps, retrieve_parent_token_ptr, + intermediate_state_indices_ptr, # Matrix dimensions batch, T, @@ -130,6 +139,7 @@ def _selective_scan_update_kernel( DISABLE_STATE_UPDATE: tl.constexpr, CACHE_INTERMEDIATE_STATES: tl.constexpr, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr, + HAS_INTERMEDIATE_STATE_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, ): pid_m = tl.program_id(axis=0) @@ -177,7 +187,12 @@ def _selective_scan_update_kernel( cache_idx = -1 if CACHE_INTERMEDIATE_STATES: - if HAS_STATE_BATCH_INDICES: + if HAS_INTERMEDIATE_STATE_INDICES: + intermediate_state_idx = tl.load(intermediate_state_indices_ptr + pid_b).to( + tl.int64 + ) + cache_idx = intermediate_state_idx + elif HAS_STATE_BATCH_INDICES: cache_idx = state_batch_idx else: cache_idx = pid_b @@ -250,7 +265,7 @@ def _selective_scan_update_kernel( if state_batch_idx != pad_slot_id: cache_ptr_base = ( intermediate_states_buffer - + state_batch_idx * cache_steps * nheads * dim * dstate + + cache_idx * cache_steps * nheads * dim * dstate + current_step_idx * nheads * dim * dstate + pid_h * dim * dstate ) @@ -300,6 +315,7 @@ def selective_state_update( intermediate_states_buffer=None, cache_steps=None, retrieve_parent_token=None, + intermediate_state_indices=None, ): """ Argument: @@ -324,6 +340,8 @@ def selective_state_update( intermediate_states_buffer: Buffer to cache intermediate states cache_steps: Total number of steps in the buffer retrieve_parent_token: (batch, T) tensor of parent token indices for EAGLE tree attention + intermediate_state_indices: (batch,) tensor of indices for intermediate_states_buffer operations. + If provided, uses these indices instead of state_batch_indices for the buffer. """ if state.dim() == 3: state = state.unsqueeze(1) @@ -426,6 +444,7 @@ def selective_state_update( intermediate_states_buffer, cache_steps if cache_steps is not None else 0, retrieve_parent_token, + intermediate_state_indices, batch, T, nheads, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 261895b01..a870358d1 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1376,7 +1376,13 @@ class ServerArgs: else: self.quantization = model_config.quantization self.moe_runner_backend = "flashinfer_cutlass" - if not self.disable_radix_cache: + + if not self.disable_radix_cache and self.speculative_algorithm is not None: + logger.warning( + "Disabling radix cache since speculative decoding for NemotronHForCausalLM is not supported with radix cache yet." + ) + self.disable_radix_cache = True + elif not self.disable_radix_cache: logger.warning( "Disabling overlap schedule since MambaRadixCache is not compatible with " "overlap schedule currently, try to use --disable-radix-cache if overlap schedule is necessary"