import logging import math import torch from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.lightning_attn import ( BailingLinearKernel, linear_decode_forward_triton, ) from sglang.srt.layers.attention.linear.linear_metadata import ( BailingLinearMetadata, ) from sglang.srt.layers.attention.linear.seg_la import SegLaMeta, seg_la_fwd from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.runtime_context import get_parallel logger = logging.getLogger(__name__) class LightningAttentionBackend(MambaAttnBackendBase): """ Note about the init: - If no spec decoding - FlashAttentionBackend will be init once when the server starts. - If spec decoding - FlashAttentionBackend will be init once for the target worker - FlashAttentionMultiStepBackend will be once for the draft worker - It will spawn num_steps FlashAttentionBackend for the draft worker Note about CUDA Graph: - We only support CUDA Graph for Decode (Normal Decode and Draft Decode) and Target Verify. - We don't support CUDA Graph for Extend and Draft Extend. - When server init, init_cuda_graph_state will be called first and then init_cuda_graph_capture will be called. - For each forward batch, init_replay_cuda_graph will be called first and then replay the graph. """ def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) # seg_la processes draft tokens as a chain -- it has no parent-indices # plumbing for tree-shaped drafts, so spec v2 tree verify (topk > 1) would # commit wrong mamba states silently. Fail fast instead of mis-decoding. if self.topk > 1: raise NotImplementedError( "Lightning (seg_la) linear-attention backend does not support " f"speculative decoding with topk > 1 (got topk={self.topk}); " "seg_la verifies a draft tree as a chain. Use " "--speculative-eagle-topk 1." ) # lightning attn does not need conv cache, but to keep the interface for mamba cache self.conv_states_shape = ( model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape ) assert not ( model_runner.sliding_window_size is not None and model_runner.model_config.is_encoder_decoder ), "Sliding window and cross attention are not supported together" # extra metadata for handling speculative decoding topk > 1, extended draft decode and verify self.max_context_len = model_runner.model_config.context_len self.device = model_runner.device self.decode_cuda_graph_metadata = {} self.kv_cache_dtype = model_runner.kv_cache_dtype self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype self.BLOCK = ( model_runner.model_config.block if hasattr(model_runner.model_config, "block") else 256 ) total_num_heads = model_runner.model_config.hf_config.num_attention_heads num_hidden_layers = model_runner.model_config.hf_config.num_hidden_layers self.tp_slope = LightningAttentionBackend._build_slope_tensor( total_num_heads, num_hidden_layers, self.device ) self.linear_backend = getattr( model_runner.model_config.hf_config, "linear_backend", "seg_la" ) logger.info( f"linear_backend for linear attention in hybrid_linear_backend: {self.linear_backend}" ) def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False, ): # seq_lens_cpu is unused by the underlying _replay_metadata for # non-target-verify modes; pass it through for compatibility. bs = forward_batch.batch_size metadata = self._replay_metadata( bs, forward_batch.req_pool_indices, forward_batch.forward_mode, forward_batch.spec_info, forward_batch.seq_lens_cpu if not in_capture else None, ) self.forward_metadata = BailingLinearMetadata.prepare_decode( metadata.query_start_loc, metadata.mamba_cache_indices, bs, forward_batch.seq_lens, ) def init_forward_metadata(self, forward_batch: ForwardBatch): metadata = self._forward_metadata(forward_batch) self.forward_metadata = BailingLinearMetadata.prepare_mixed( metadata.query_start_loc, metadata.mamba_cache_indices, forward_batch, ) @staticmethod def _build_slope_tensor( n_attention_heads: int, num_hidden_layers: int, device="cuda" ): def get_slopes(n): def get_slopes_power_of_2(n): start = 2 ** (-(2 ** -(math.log2(n) - 3))) ratio = start return [start * ratio**i for i in range(n)] if math.log2(n).is_integer(): return get_slopes_power_of_2(n) else: closest_power_of_2 = 2 ** math.floor(math.log2(n)) return ( get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] ) slopes = torch.tensor( get_slopes(n_attention_heads), dtype=torch.float32 ).reshape(n_attention_heads, 1, 1) tp_heads = n_attention_heads // get_parallel().attn_tp_size tp_rank = get_parallel().attn_tp_rank if num_hidden_layers <= 1: slope_rate_list = [slopes * (1 + 1e-5)] else: slope_rate_list = [ slopes * (1 - layer_id / (num_hidden_layers - 1) + 1e-5) for layer_id in range(num_hidden_layers) ] tp_slope = [ slope_rate_list[layer_id][tp_rank * tp_heads : (tp_rank + 1) * tp_heads] .contiguous() .to(device) for layer_id in range(num_hidden_layers) ] return tp_slope def _prefill_and_mix_infer( self, q, k, v, kv_cache, state_indices_tensor, forward_batch, layer, metadata, ): hidden = [] for _prefill_idx in range(metadata.num_prefills): if _prefill_idx >= forward_batch.extend_start_loc.shape[0]: break if _prefill_idx >= state_indices_tensor.shape[0]: break _start = forward_batch.extend_start_loc[_prefill_idx] if _prefill_idx + 1 < forward_batch.extend_start_loc.shape[0]: _end = forward_batch.extend_start_loc[_prefill_idx + 1] else: if ( forward_batch.extend_seq_lens is not None and _prefill_idx < forward_batch.extend_seq_lens.shape[0] and metadata.num_decodes > 0 ): seq_len = forward_batch.extend_seq_lens[_prefill_idx] _end = _start + seq_len else: _end = q.shape[0] slot_id = state_indices_tensor[_prefill_idx] qs = q[_start:_end].transpose(0, 1).contiguous() ks = k[_start:_end].transpose(0, 1).contiguous() vs = v[_start:_end].transpose(0, 1).contiguous() slice_layer_cache = kv_cache[slot_id, ...] out_slice = BailingLinearKernel.jit_linear_forward_prefix( qs, ks, vs, slice_layer_cache, self.tp_slope[layer.layer_id], self.BLOCK, layer_idx=layer.layer_id, ) hidden.append(out_slice.contiguous()) if metadata.num_decodes > 0: hidden.append( self._decode_infer( q, k, v, kv_cache, state_indices_tensor, metadata, layer ) ) if not hidden: return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype) hidden = torch.concat(hidden, dim=0).contiguous() return hidden def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, metadata, layer): num_prefill_tokens = metadata.num_prefill_tokens num_prefills = metadata.num_prefills q = q[num_prefill_tokens:].unsqueeze(2).contiguous() k = k[num_prefill_tokens:].unsqueeze(2).contiguous() v = v[num_prefill_tokens:].unsqueeze(2).contiguous() slot_id = state_indices_tensor[num_prefills:] assert slot_id.shape[0] == q.shape[0], ( f"slot_id length {slot_id.shape[0]} does not match decode batch size {q.shape[0]}. " "This indicates a bug in the upstream logic that should be investigated." ) hidden = linear_decode_forward_triton( q, k, v, kv_cache, self.tp_slope[layer.layer_id], slot_id, 32 ) return hidden def _linear_attention_entry( self, q, k, v, kv_cache, state_indices_tensor, metadata, layer, mask=None, temp_cache=None, intermediate_state_indices=None, ): q_offsets = metadata.query_start_loc seg_meta = SegLaMeta( batch_size=metadata.batch_size, q_offsets=metadata.query_start_loc, s_offsets=state_indices_tensor, q_lengths=q_offsets.diff(), s_scales=metadata.has_initial_states, max_q_length=None, mask=mask, ) hidden = seg_la_fwd( q=q, k=k, v=v, s=kv_cache, decay_scales=self.tp_slope[layer.layer_id], meta=seg_meta, caches=temp_cache, cache_indices=intermediate_state_indices, decouple=True, ) return hidden def forward_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, save_kv_cache=True, **kwargs, ): layer_id = layer.layer_id if layer else kwargs["layer_id"] metadata = self.forward_metadata if self.kv_cache_dtype_str != "auto" and layer.k_scale is not None: q = q.to(self.kv_cache_dtype) cache_indices = self.forward_metadata.mamba_cache_indices mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id) ssm_states = mamba_cache_params.temporal if self.linear_backend == "minimax": o = self._prefill_and_mix_infer( q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), k, v, ssm_states, cache_indices, forward_batch, layer, metadata, ) elif self.linear_backend == "seg_la": intermediate_state_indices = ( torch.arange( cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device, ) if forward_batch.forward_mode.is_target_verify() else None ) o = self._linear_attention_entry( q, k, v, ssm_states, cache_indices, metadata, layer, temp_cache=( mamba_cache_params.intermediate_ssm if forward_batch.forward_mode.is_target_verify() else None ), intermediate_state_indices=intermediate_state_indices, ) else: raise ValueError( f"linear backend: {self.linear_backend} is not support for now" ) if ( not forward_batch.forward_mode.is_target_verify() and forward_batch.mamba_track_mask is not None ): # save mamba cache for extra buffer mamba_track_mask = forward_batch.mamba_track_mask mamba_track_indices = forward_batch.mamba_track_indices dst_masked = mamba_track_indices[mamba_track_mask] src_masked = metadata.mamba_cache_indices[mamba_track_mask] ssm_states[dst_masked] = ssm_states[src_masked] return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) def forward_decode( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, save_kv_cache=True, **kwargs, ) -> torch.Tensor: layer_id = layer.layer_id if layer else kwargs["layer_id"] # Use precomputed metadata across all layers metadata = self.forward_metadata if self.kv_cache_dtype_str != "auto": q = q.to(self.kv_cache_dtype) # Do linear attention cache_indices = self.forward_metadata.mamba_cache_indices mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id) ssm_states = mamba_cache_params.temporal if self.linear_backend == "minimax": o = self._decode_infer(q, k, v, ssm_states, cache_indices, metadata, layer) elif self.linear_backend == "seg_la": o = self._linear_attention_entry( q, k, v, ssm_states, cache_indices, metadata, layer ) else: raise ValueError( f"linear backend: {self.linear_backend} is not support for now" ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)