diff --git a/lightllm/common/attention/backends/_flashinfer.py b/lightllm/common/attention/backends/_flashinfer.py new file mode 100644 index 0000000000..cae6a29ca7 --- /dev/null +++ b/lightllm/common/attention/backends/_flashinfer.py @@ -0,0 +1,10 @@ +import os + +import torch + + +def set_flashinfer_envs(): + capability = torch.cuda.get_device_capability() + arch = f"{capability[0]}.{capability[1]}" + os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" + return diff --git a/lightllm/common/attention/backends/_shared/linear/gdn.py b/lightllm/common/attention/backends/_shared/linear/gdn.py new file mode 100644 index 0000000000..6d673d52a0 --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/gdn.py @@ -0,0 +1,56 @@ +"""GDN tensor-parallel dimensions and input views shared by native and replay kernels.""" + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + + +class GDNMetadata: + def _init_linear_layer_metadata(self, network_config, tp_world_size): + self.mtp_step = get_env_start_args().mtp_step + self.num_v_heads = network_config["linear_num_value_heads"] + self.num_k_heads = network_config["linear_num_key_heads"] + self.head_k_dim = network_config["linear_key_head_dim"] + self.head_v_dim = network_config["linear_value_head_dim"] + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] + self.activation = network_config["hidden_act"] + self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size + self.tp_ba_dim = self.num_v_heads * 2 // tp_world_size + self.tp_num_k_heads = self.num_k_heads // tp_world_size + self.tp_num_v_heads = self.num_v_heads // tp_world_size + self.tp_key_dim = self.key_dim // tp_world_size + self.tp_value_dim = self.value_dim // tp_world_size + if self.num_v_heads % self.num_k_heads: + raise ValueError("num_v_heads must be divisible by num_k_heads") + self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads + ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} + self.ssm_state_dtype = ssm_dtype_dict.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + return + + def _split_qkvzba(self, mixed_qkvzba): + qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim + z_end = qkv_dim + self.tp_value_dim + b_end = z_end + self.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return (mixed_qkv, z, b, a) + + def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): + if decode: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + batch_size = mixed_qkv.shape[0] + query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) + else: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + seq_len = query.shape[0] + query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) diff --git a/lightllm/common/attention/backends/_shared/linear/layout.py b/lightllm/common/attention/backends/_shared/linear/layout.py new file mode 100644 index 0000000000..1ada562ea0 --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/layout.py @@ -0,0 +1,67 @@ +from dataclasses import dataclass, field + +import torch + +from lightllm.common.attention.backends.base import AttentionExecutionContext + +from .mtp_state_params import build_dynamic_mtp_layout + + +@dataclass(slots=True) +class LinearDecodeContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor | None = None + b_ssm_buffer_idx: torch.Tensor | None = None + b1_mtp_cu_q_seq_len: torch.Tensor | None = None + b_cu_seqlens: torch.Tensor | None = None + accepted_by_state: dict[torch.Tensor, torch.Tensor] = field(default_factory=dict) + + def accepted_tokens(self, state_index: torch.Tensor) -> torch.Tensor: + """Gather once per State per update; Graph captures the gather for every replay. + + The cache retains the source tensors until the next update, keeping identity + keys valid for this context. HOLD sequences must select the initial convolution state. + """ + key = state_index + if key not in self.accepted_by_state: + self.accepted_by_state[key] = torch.where( + self.b_conv_buffer_idx == self.batch.hold_req_id, + 1, + state_index[self.b_conv_buffer_idx] + 1, + ) + return self.accepted_by_state[key] + + +def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: bool, *, replay: bool) -> None: + batch = context.batch + context.accepted_by_state.clear() + if draft_step == 0: + context.b_conv_buffer_idx = batch.b_req_idx + context.b_ssm_buffer_idx = batch.b_req_idx + context.b1_mtp_cu_q_seq_len = None + if replay: + context.b_cu_seqlens = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + return + width = draft_step + 1 + if dynamic: + context.b1_mtp_cu_q_seq_len, context.b_conv_buffer_idx = build_dynamic_mtp_layout( + b_req_idx=batch.b_req_idx, + b_mtp_index=batch.b_mtp_index, + hold_req_id=batch.hold_req_id, + ) + else: + if batch.batch_size % width: + raise ValueError("linear decode batch size must be divisible by draft_step + 1") + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch.batch_size + 1, + width, + dtype=torch.int32, + device=batch.device, + ) + context.b_conv_buffer_idx = batch.b_req_idx.view(-1, width)[:, 0].contiguous() + if replay: + context.b_ssm_buffer_idx = context.b_conv_buffer_idx + context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len + else: + offsets = torch.arange(width, dtype=batch.b_req_idx.dtype, device=batch.device) + context.b_ssm_buffer_idx = context.b_conv_buffer_idx[:, None] * width + offsets diff --git a/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py index b81bf9ec16..993ef6f1b2 100644 --- a/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py +++ b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py @@ -14,6 +14,7 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( batch_size, hold_req_id, BLOCK_SIZE: tl.constexpr, + SELECT_ACCEPTED: tl.constexpr = True, ): offsets = tl.arange(0, BLOCK_SIZE) token_mask = offsets < batch_size @@ -30,21 +31,18 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( # replay arbitrary compact per-request widths at the same token batch size. tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask) tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=token_mask) - tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) + if SELECT_ACCEPTED: + tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) tl.debug_barrier() # mtp_index restarts at zero on the first row of every request group. is_start = valid_row & (mtp_index == 0) sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1 - accepted_state_index = tl.load( - req_to_mtp_state_index + req_idx, - mask=is_start, - other=0, - ) - tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start) tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start) - tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) + if SELECT_ACCEPTED: + accepted_state_index = tl.load(req_to_mtp_state_index + req_idx, mask=is_start, other=0) + tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) def build_dynamic_mtp_linear_att_state_params( @@ -85,11 +83,10 @@ def build_dynamic_mtp_linear_att_state_params( per-request query lengths. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda - assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape - assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32 + _validate_layout_inputs(b_req_idx, b_mtp_index) + if req_to_mtp_state_index.device != b_req_idx.device or req_to_mtp_state_index.ndim != 1: + raise ValueError("MTP state indices must be a vector on the request device") batch_size = b_req_idx.shape[0] - assert batch_size > 0 b1_cu_q_seq_len = torch.empty((batch_size + 1,), dtype=torch.int32, device=b_req_idx.device) b_conv_buffer_idx = torch.empty_like(b_req_idx) @@ -109,3 +106,34 @@ def build_dynamic_mtp_linear_att_state_params( num_stages=1, ) return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + + +def build_dynamic_mtp_layout(b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, hold_req_id: int): + _validate_layout_inputs(b_req_idx, b_mtp_index) + batch_size = b_req_idx.numel() + cu = torch.empty(batch_size + 1, dtype=torch.int32, device=b_req_idx.device) + requests = torch.empty_like(b_req_idx) + _build_dynamic_mtp_linear_att_state_params_kernel[(1,)]( + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_mtp_state_index=None, + out_cu_q_seq_len=cu, + out_conv_buffer_idx=requests, + out_num_accepted_tokens=None, + batch_size=batch_size, + hold_req_id=hold_req_id, + BLOCK_SIZE=triton.next_power_of_2(batch_size + 1), + SELECT_ACCEPTED=False, + num_warps=8, + num_stages=1, + ) + return cu, requests + + +def _validate_layout_inputs(requests, positions): + if not requests.is_cuda or positions.device != requests.device: + raise ValueError("MTP layout requires inputs on the same CUDA device") + if requests.ndim != 1 or requests.shape != positions.shape or requests.numel() == 0: + raise ValueError("MTP layout requires matching nonempty vectors") + if requests.dtype != torch.int32 or positions.dtype != torch.int32: + raise ValueError("MTP layout requires int32 request and position indices") diff --git a/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py index b7f138c844..450c695375 100644 --- a/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix +from .replayssm_delta import _apply_accepted_prefix def check_replay_sync_mark( diff --git a/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py index 2e491a038b..42aad6edc2 100644 --- a/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix +from .replayssm_delta import _apply_accepted_prefix @triton.jit diff --git a/lightllm/common/attention/backends/dense/fa3/backend.py b/lightllm/common/attention/backends/dense/fa3/backend.py new file mode 100644 index 0000000000..644f70e281 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/backend.py @@ -0,0 +1,218 @@ +import dataclasses + +import torch +import triton + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers +from lightllm.utils.envs_utils import enable_env_vars +from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune + + +@dataclasses.dataclass(slots=True) +class Fa3PrefillContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class Fa3DecodeContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class Fa3AttBackend(BaseAttentionBackend[Fa3PrefillContext | Fa3DecodeContext, DenseInputs, torch.Tensor]): + requires_cpu_cache_sync = enable_env_vars("LIGHTLLM_FA3_CPU_CACHE_SYNC") + provider = "fa3" + dependency = "sgl_kernel.flash_attn" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_attention_sink = True + supports_sliding_window = True + supports_noncausal_decode = True + max_head_size = 256 + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.causal = self.uses_causal_attention() + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.page_table = torch.empty( + (context.batch.batch_size, triton.cdiv(context.batch.max_kv_seq_len, self.infer_page_size)), + dtype=torch.int32, + device=context.batch.device, + ) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=context.batch.b_req_idx, + page_size=self.infer_page_size, + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_paged( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ) -> torch.Tensor: + if layer.params.sliding_window is not None: + window_size = layer.params.sliding_window + else: + window_size = (-1, -1) + if layer.params.sink is not None: + sink_weight: torch.Tensor = layer.params.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq**0.5 + o = flash_attn_with_kvcache( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o + + def _update_decode(self, context): + context.causal = self.uses_causal_attention() + draft_step = self.decode_draft_step + if self.uses_dynamic_spec_verify_layout(): + b_att_req_idx = self._decode_init_dynamic_spec_verify_state(context, draft_step) + elif draft_step > 0: + b_att_req_idx = self._decode_init_fixed_spec_decode_state(context, draft_step) + else: + b_att_req_idx = self._decode_init_normal_decode_state(context) + self._decode_init_page_table(context, b_att_req_idx) + + def _decode_init_dynamic_spec_verify_state(self, context, draft_step: int) -> torch.Tensor: + b_mark_mtp_shared_group = build_mtp_shared_group_markers(context.batch.b_req_idx, hold_req_id=self.hold_req_id) + b_q_seq_len, b_kv_seq_len, b_att_req_idx, context.b_att_seq_len = build_dynamic_spec_fa3_decode_params( + b_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_mtp_shared_group=b_mark_mtp_shared_group, + att_batch_size=context.batch.batch_size, + hold_req_id=self.hold_req_id, + ) + self._decode_init_spec_decode_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = draft_step + 1 + return b_att_req_idx + + def _decode_init_fixed_spec_decode_state(self, context, draft_step: int) -> torch.Tensor: + mtp_size = draft_step + 1 + assert context.batch.batch_size % mtp_size == 0, ( + "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={context.batch.batch_size}, draft_step={draft_step}." + ) + b_q_seq_len = torch.full( + (context.batch.b_seq_len.shape[0] // mtp_size,), + fill_value=mtp_size, + dtype=torch.int32, + device=context.batch.b_seq_len.device, + ) + b_kv_seq_len = context.batch.b_seq_len[draft_step::mtp_size] + b_att_req_idx = context.batch.b_req_idx[draft_step::mtp_size] + context.b_att_seq_len = b_kv_seq_len.contiguous() + self._decode_init_spec_decode_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = mtp_size + return b_att_req_idx + + def _decode_init_normal_decode_state(self, context) -> torch.Tensor: + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.b_att_seq_len = context.batch.b_seq_len + context.decode_max_q_seq_len = 1 + return context.batch.b_req_idx + + def _decode_init_spec_decode_cu_seqlens(self, context, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): + b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) + context.cu_seqlens_q = b1_cu_q_seq_len.int() + context.cu_seqlens_k = b1_cu_kv_seq_len.int() + + def _decode_init_page_table(self, context, b_att_req_idx): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + width = triton.cdiv(context.capacity.max_kv_seq_len, self.infer_page_size) + context.page_table = context.workspace.get("fa3_page_table", b_att_req_idx.numel() * width, torch.int32)[ + : b_att_req_idx.numel() * width + ].view(b_att_req_idx.numel(), width) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=b_att_req_idx, + page_size=self.infer_page_size, + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._decode_normal_decode_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _decode_normal_decode_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ): + if layer.params.sliding_window is not None: + window_size = layer.params.sliding_window + else: + window_size = (-1, -1) + if layer.params.sink is not None: + sink_weight: torch.Tensor = layer.params.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq**0.5 + o = flash_attn_with_kvcache_autotune( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=None, + max_seqlen_q=context.decode_max_q_seq_len, + max_seqlen_k=context.decode_max_kv_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/fp8.py b/lightllm/common/attention/backends/dense/fa3/fp8.py new file mode 100644 index 0000000000..6fd5aade3e --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/fp8.py @@ -0,0 +1,154 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend, Fa3DecodeContext, Fa3PrefillContext +from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( + q_per_head_fp8_quant, + q_per_head_static_fp8_quant, +) +from lightllm.utils.log_utils import init_logger +from lightllm.utils.sgl_utils import flash_attn_with_kvcache + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3PrefillContext(Fa3PrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3DecodeContext(Fa3DecodeContext): + pass + + +class Fp8Fa3AttBackend(Fa3AttBackend): + provider = "fa3" + supports_sliding_window = False + cache_formats = frozenset({"fp8kv_sph"}) + + @property + def infer_page_size(self) -> int: + return 1 + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8Fa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8Fa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + super()._update_prefill(context) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_fp8_prefill_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _prefill_fp8_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ) -> torch.Tensor: + q_head_num = q.shape[1] + q_head_dim = q.shape[2] + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + batch_size = context.batch.b_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.params.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.params.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, q_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale, + k_descale=layer.params.key_scale.view(1, -1).expand(batch_size, -1), + v_descale=layer.params.value_scale.view(1, -1).expand(batch_size, -1), + return_softmax_lse=False, + ) + return o + + def _update_decode(self, context): + super()._update_decode(context) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._decode_fp8_decode_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _decode_fp8_decode_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ): + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + q_head_num = q.shape[1] + att_batch_size = context.b_att_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.params.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.params.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, k_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale.view(att_batch_size, k_head_num), + k_descale=layer.params.key_scale.view(1, -1).expand(att_batch_size, -1), + v_descale=layer.params.value_scale.view(1, -1).expand(att_batch_size, -1), + return_softmax_lse=False, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/neo.py b/lightllm/common/attention/backends/dense/fa3/neo.py new file mode 100644 index 0000000000..cebe3a015d --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/neo.py @@ -0,0 +1,80 @@ +import dataclasses +import inspect +import warnings + +import torch + +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend, Fa3PrefillContext + +_NEO_FA3_INSTALL_HINT = ( + "Install the Neo FA3 build from source in the current Python environment:\n" + " git clone --recursive https://github.com/WANDY666/flash-attention.git\n" + " cd flash-attention/hopper\n" + " FLASH_ATTENTION_FORCE_BUILD=TRUE python -m pip install --no-build-isolation --no-deps --force-reinstall .\n" + "This replaces a standalone flash_attn_3 / flash_attn_interface installation; LightLLM's regular" + " FA3 uses the separate sgl_kernel package." +) +try: + from flash_attn_interface import flash_attn_with_kvcache as flash_attn_with_kvcache_neo + + _sig = inspect.signature(flash_attn_with_kvcache_neo) + if "image_token_end" not in _sig.parameters: + raise ImportError("flash_attn_interface is missing image_token_end support (need the Neo build)") + HAS_FLASH_ATTN_INTERFACE = True +except ImportError as exc: + warnings.warn( + f"Neo FA3 is unavailable: {exc}. Automatic Neo prefill selection will fall back to Triton. " + "To select Triton explicitly, use --llm_prefill_att_backend triton.\n" + _NEO_FA3_INSTALL_HINT + ) + flash_attn_with_kvcache_neo = None + HAS_FLASH_ATTN_INTERFACE = False + + +@dataclasses.dataclass(slots=True) +class NeoFa3PrefillContext(Fa3PrefillContext): + pass + + +class NeoFa3AttBackend(Fa3AttBackend): + @classmethod + def check_support(cls, spec, caps=None, execution=None) -> None: + super().check_support(spec, caps, execution) + if spec.phase.value == "prefill" and not HAS_FLASH_ATTN_INTERFACE: + raise UnsupportedAttention("Neo FA3 requires the image-token flash attention interface") + + provider = "fa3" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NeoFa3PrefillContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + q = inputs.query + k = layer.params.key + v = layer.params.value + if not HAS_FLASH_ATTN_INTERFACE: + raise ImportError("Neo prefill requires FA3 with image_token_end support.\n" + _NEO_FA3_INSTALL_HINT) + if context.batch.b_image_token_end is None: + raise ValueError("Neo prefill requires b_image_token_end to describe image attention spans") + o = flash_attn_with_kvcache_neo( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=1.0 / q.shape[-1] ** 0.5, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=None, + v_descale=None, + return_softmax_lse=False, + image_token_end=context.batch.b_image_token_end, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/varlen.py b/lightllm/common/attention/backends/dense/fa3/varlen.py new file mode 100644 index 0000000000..f2636ead31 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/varlen.py @@ -0,0 +1,70 @@ +import torch + +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class Fa3VarlenAttentionOp(BaseVarlenAttentionOp): + provider = "fa3" + dependency = "sgl_kernel.flash_attn" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + max_head_size = 256 + head_alignment = 8 + supports_graph = True + + def __init__(self): + self._op = None + + def forward(self, inputs, *, out): + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + tokens, _, head_dim = q.shape + if tokens == 0: + return out + if self._op is None: + import sgl_kernel.flash_attn + + self._op = torch.ops.sgl_kernel.fwd.default + softmax_scale = head_dim**-0.5 + window_size = (-1, -1) + self._op( + q=q, + k=k, + v=v, + k_new=None, + v_new=None, + q_v=None, + out=out, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + cu_seqlens_k_new=None, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + page_table=None, + kv_batch_idx=None, + leftpad_k=None, + rotary_cos=None, + rotary_sin=None, + seqlens_rotary=None, + q_descale=None, + k_descale=None, + v_descale=None, + softmax_scale=softmax_scale, + is_causal=False, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=0, + softcap=0.0, + is_rotary_interleaved=False, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0, + sinks=None, + sparse_mask_fine=None, + only_qv=False, + ) + + return out diff --git a/lightllm/common/attention/backends/dense/flashinfer/backend.py b/lightllm/common/attention/backends/dense/flashinfer/backend.py new file mode 100644 index 0000000000..78d1c937de --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/backend.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch +import triton + +from lightllm.common.attention.api import BackendResources +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.decode import BatchDecodeWithPagedKVCacheWrapper + from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class FlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class FlashInferDecodeContext(AttentionExecutionContext): + kv_last_page_len_buffer: torch.Tensor = None + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + decode_wrapper: BatchDecodeWithPagedKVCacheWrapper | None = None + + +class FlashInferAttBackend( + BaseAttentionBackend[FlashInferPrefillContext | FlashInferDecodeContext, DenseInputs, torch.Tensor] +): + provider = "flashinfer" + dependency = "flashinfer" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_expert_parallel = False + cache_formats = frozenset({"None"}) + workspace_buffer_key = "flashinfer_fp" + workspace_buffer_size = 512 * 1024 * 1024 + + def __init__(self, model): + set_flashinfer_envs() + super().__init__(model=model) + tp_world_size = get_dp_world_size() + self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size + self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1) + head_dim = model.config["hidden_size"] // model.config["num_attention_heads"] + self.head_dim = model.config.get("head_dim", head_dim) + self.max_seq_length = model.max_seq_length + self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) + self.q_data_type = model.data_type + self.kv_data_type = model.data_type + + def configure_spec(self, spec): + self.tp_q_head_num = spec.num_q_heads + self.tp_kv_head_num = spec.cache.key_heads + self.head_dim = spec.head_size + + def resource_requirements(self, capacity, *, capture): + return BackendResources( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return FlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return FlashInferDecodeContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_paged( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ) -> torch.Tensor: + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + context.prefill_wrapper.run( + q, + ( + k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def _decode_refresh_cuda_graph_decode_plan(self, context, max_kv_len: int): + from flashinfer.decode import fast_decode_plan + + uniform_kv_indptr_cpu = torch.arange( + context.batch.batch_size + 1, dtype=torch.int32, device="cpu" + ) * triton.cdiv(max_kv_len, self.infer_page_size) + fast_decode_plan( + context.decode_wrapper, + indptr=context.kv_starts, + indices=context.kv_indices, + last_page_len=context.kv_last_page_len_buffer, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_kv_head_num, + head_dim=self.head_dim, + page_size=self.infer_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + global_override_indptr_cpu=uniform_kv_indptr_cpu, + ) + return + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._decode_normal_decode_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _decode_normal_decode_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ): + o_tensor = alloc_func(q.shape, q.dtype) + context.decode_wrapper.run( + q, + ( + k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def _prepare_prefill(self, context): + import flashinfer + + batch = context.batch + q_starts = batch.b1_cu_q_seq_len.int() + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.infer_page_size + indices = context.workspace.get("flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=kv_starts[:-1], + max_token_len=batch.max_kv_seq_len, + out_page_indices=indices, + page_size=self.infer_page_size, + ) + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + qo_indptr_buf=q_starts, + paged_kv_indptr_buf=kv_starts, + paged_kv_indices_buf=indices, + paged_kv_last_page_len_buf=last_page_len, + ) + context.prefill_wrapper.plan( + q_starts, + kv_starts, + indices, + last_page_len, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.infer_page_size, + causal=self.causal, + pos_encoding_mode="NONE", + logits_soft_cap=0.0, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + ) + + def _prepare_decode(self, context): + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.infer_page_size + if context.kv_starts is None: + context.kv_starts, context.kv_last_page_len_buffer = starts, last_page_len + context.kv_indices = context.workspace.get( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + else: + context.kv_starts.copy_(starts, non_blocking=True) + context.kv_last_page_len_buffer.copy_(last_page_len, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.infer_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + "NHD", + use_cuda_graph=context.capture, + use_tensor_cores=True, + paged_kv_indptr_buffer=context.kv_starts, + paged_kv_indices_buffer=context.kv_indices, + paged_kv_last_page_len_buffer=context.kv_last_page_len_buffer, + ) + context.decode_wrapper.plan( + context.kv_starts, + context.kv_indices, + context.kv_last_page_len_buffer, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.infer_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + ) + else: + self._decode_refresh_cuda_graph_decode_plan(context, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/dense/flashinfer/fp8.py b/lightllm/common/attention/backends/dense/flashinfer/fp8.py new file mode 100644 index 0000000000..12a7f54e88 --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/fp8.py @@ -0,0 +1,99 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.dense.flashinfer.backend import ( + FlashInferAttBackend, + FlashInferDecodeContext, + FlashInferPrefillContext, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferPrefillContext(FlashInferPrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferDecodeContext(FlashInferDecodeContext): + pass + + +class Fp8FlashInferAttBackend(FlashInferAttBackend): + provider = "flashinfer" + cache_formats = frozenset({"fp8kv_spt"}) + + def __init__(self, model): + set_flashinfer_envs() + super().__init__(model=model) + self.kv_data_type = torch.float8_e4m3fn + + @property + def infer_page_size(self) -> int: + return 1 + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8FlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8FlashInferDecodeContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_fp8_prefill_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _prefill_fp8_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ) -> torch.Tensor: + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.params.key_scale + v_descale = layer.params.value_scale + context.prefill_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._decode_fp8_decode_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _decode_fp8_decode_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ): + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.params.key_scale + v_descale = layer.params.value_scale + context.decode_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/sdpa/varlen.py b/lightllm/common/attention/backends/dense/sdpa/varlen.py new file mode 100644 index 0000000000..0997e37841 --- /dev/null +++ b/lightllm/common/attention/backends/dense/sdpa/varlen.py @@ -0,0 +1,52 @@ +import torch +import torch.nn.functional as F + +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp +from lightllm.platform import Platform + + +class SdpaVarlenAttentionOp(BaseVarlenAttentionOp): + """Fallback with a device-to-host boundary sync and a Python loop per sequence. + + This path cannot be captured in a CUDA Graph. + """ + + provider = "sdpa" + supported_platforms = frozenset({Platform.CPU, Platform.CUDA}) + + def forward(self, inputs, *, out): + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + output_buffer = out + assert q.ndim == k.ndim == v.ndim == output_buffer.ndim == 3 + assert cu_seqlens is not None and cu_seqlens.ndim == 1 + cu_seqlens = cu_seqlens.detach().to("cpu") + B = cu_seqlens.numel() - 1 + + with torch.no_grad(): + for b in range(B): + s = int(cu_seqlens[b]) + e = int(cu_seqlens[b + 1]) + L = e - s + if L <= 0: + continue + if max_seqlen: + assert L <= max_seqlen + + # [L, H, D] -> [1, H, L, D] + q_ = q[s:e].permute(1, 0, 2).unsqueeze(0) + k_ = k[s:e].permute(1, 0, 2).unsqueeze(0) + v_ = v[s:e].permute(1, 0, 2).unsqueeze(0) + + sequence_output = F.scaled_dot_product_attention( + q_, + k_, + v_, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + # [1, H, L, D] -> [L, H, D] + output_buffer[s:e].copy_(sequence_output.squeeze(0).permute(1, 0, 2)) + + return output_buffer diff --git a/lightllm/common/attention/backends/dense/triton/backend.py b/lightllm/common/attention/backends/dense/triton/backend.py new file mode 100644 index 0000000000..f4e3871b88 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/backend.py @@ -0,0 +1,192 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers + + +@dataclasses.dataclass(slots=True) +class TritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass(slots=True) +class TritonDecodeContext(AttentionExecutionContext): + b_mark_mtp_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class TritonAttBackend(BaseAttentionBackend[TritonPrefillContext | TritonDecodeContext, DenseInputs, torch.Tensor]): + provider = "triton" + supports_alibi = True + supports_sliding_window = True + supports_speculative_window = False + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return TritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return TritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + if layer.params.alibi is not None: + return self._prefill_alibi(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + out = alloc_func(q.shape, q.dtype) + from lightllm.common.basemodel.triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd + + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + layer.params.alibi, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + ) + return out + + def _prefill_paged(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) + + if layer.params.sliding_window is not None: + sliding_window = layer.params.sliding_window + else: + sliding_window = (-1, -1) + out = alloc_func(q.shape, q.dtype) + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + sliding_window=sliding_window, + ) + return out + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + draft_step = self.decode_draft_step + if draft_step > 0: + context.b_mark_mtp_shared_group = build_mtp_shared_group_markers( + context.batch.b_req_idx, hold_req_id=self.hold_req_id + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + if layer.params.alibi is not None: + return self._decode_alibi(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + draft_step = self.decode_draft_step + q_head_num = q.shape[1] + k_head_num = k.shape[1] + if draft_step > 0: + assert q_head_num >= k_head_num, "speculative decode requires q_head_num >= k_head_num" + return self._decode_spec_decode_gqa_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + elif q_head_num == k_head_num and layer.params.sliding_window is None: + return self._decode_mha(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + elif q_head_num >= k_head_num: + return self._decode_gqa(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + raise NotImplementedError("error") + + def _decode_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.basemodel.triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd + + out = alloc_func(q.shape, q.dtype) + token_attention_fwd( + q, + k, + v, + out, + layer.params.alibi, + context.batch.req_to_token_indexs, + context.batch.b_req_idx, + context.batch.b_kv_start_loc, + context.batch.b_seq_len, + context.batch.max_kv_seq_len, + context.batch.total_token_num, + alloc_tensor_func=alloc_func, + ) + return out + + def _decode_mha(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + out = alloc_func(q.shape, q.dtype) + token_decode_attention_flash_decoding( + q=q, batch=context.batch, cache_k=k, cache_v=v, out=out, alloc_tensor_func=alloc_func + ) + return out + + def _decode_gqa(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding import ( + gqa_token_decode_attention_flash_decoding, + ) + + if layer.params.sliding_window is not None: + sliding_window = layer.params.sliding_window + else: + sliding_window = (-1, -1) + out = alloc_func(q.shape, q.dtype) + gqa_token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_v=v, + max_len_in_batch=context.decode_max_kv_seq_len, + out=out, + alloc_tensor_func=alloc_func, + sliding_window=sliding_window, + ) + return out + + def _decode_spec_decode_gqa_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + token_decode_attention_mtp_diverse_single_token, + ) + + out = token_decode_attention_mtp_diverse_single_token( + q=q, + k=k, + v=v, + Req_to_tokens=context.batch.req_to_token_indexs, + B_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_shared_group=context.b_mark_mtp_shared_group, + max_kv_len=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) + return out diff --git a/lightllm/common/attention/backends/dense/triton/int4.py b/lightllm/common/attention/backends/dense/triton/int4.py new file mode 100644 index 0000000000..8f01bcfbed --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int4.py @@ -0,0 +1,146 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.utils.envs_utils import get_env_start_args + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonDecodeContext(AttentionExecutionContext): + decode_max_kv_seq_len: int = None + + +class Int4kvTritonAttBackend(BaseAttentionBackend): + provider = "triton" + supports_alibi = False + supports_sliding_window = False + cache_formats = frozenset({"int4kv"}) + + def __init__(self, model): + super().__init__(model) + self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Int4kvTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Int4kvTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + o = self._prefill_groupsize_quant_prefill_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + return o + + def _prefill_groupsize_quant_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + head_dim = k.shape[2] * 2 + k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) + v_dequant = alloc_func((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv + + dequantize_int4kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + return self._decode_ppl_int4kv_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + + def _decode_ppl_int4kv_decode_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.ppl_int4kv_flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_kv_seq_len=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) diff --git a/lightllm/common/attention/backends/dense/triton/int8.py b/lightllm/common/attention/backends/dense/triton/int8.py new file mode 100644 index 0000000000..d9533e4f28 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int8.py @@ -0,0 +1,184 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers +from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel, get_env_start_args + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonDecodeContext(AttentionExecutionContext): + b_shared_seq_len: torch.Tensor = None + b_mark_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class Int8kvTritonAttBackend(BaseAttentionBackend): + provider = "triton" + supports_alibi = False + supports_sliding_window = False + cache_formats = frozenset({"int8kv"}) + + def __init__(self, model): + super().__init__(model) + self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Int8kvTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Int8kvTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + o = self._prefill_groupsize_quant_prefill_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + return o + + def _prefill_groupsize_quant_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) + v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv + + dequantize_int8kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + if enable_diverse_mode_gqa_decode_fast_kernel(): + context.b_mark_shared_group = build_diverse_shared_group_markers( + b_shared_radix_node_id=context.batch.b_shared_radix_node_id + ) + context.b_shared_seq_len = torch.where(context.b_mark_shared_group == 1, 0, context.batch.b_shared_seq_len) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + if enable_diverse_mode_gqa_decode_fast_kernel(): + return self._decode_diverse_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + else: + return self._decode_normal_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + + def _decode_diverse_decode_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + from .kernels.decode.int8kv.int8kv_flash_decoding_diverse import token_decode_attention_flash_decoding + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + alloc_tensor_func=alloc_func, + b_shared_seq_len=context.b_shared_seq_len, + b_mark_shared_group=context.b_mark_shared_group, + shared_streams_dict=context.workspace.get_object("int8_decode_streams", dict), + ) + + def _decode_normal_decode_att( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_len_in_batch=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py new file mode 100644 index 0000000000..5dd6a683fe --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py @@ -0,0 +1,59 @@ +import torch + + +def gqa_token_decode_attention_flash_decoding( + q: torch.Tensor, + batch, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + max_len_in_batch: int, + out=None, + alloc_tensor_func=torch.empty, + sliding_window=(-1, -1), +): + batch_size = batch.batch_size + q_head_num, head_dim = q.shape[1], q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + from .gqa_flash_decoding_stage1 import flash_decode_stage1 + from .gqa_flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # Because we need to allocate some intermediate tensors, considering parallelism and + # intermediate memory consumption, when batch_size is small, block_num is larger, + # and when batch_size is large, block_num is smaller. This achieves a better balance + # of memory consumption and performance. + BLOCK_SEQ = 256 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + v=cache_v, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + max_len_in_batch=max_len_in_batch, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + sliding_window=sliding_window, + ) + flash_decode_stage2( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + out=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + sliding_window=sliding_window, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py index dc98bb3544..8b2dbc885a 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py @@ -1,13 +1,10 @@ import torch -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( - mtp_diverse_stage1_single_token, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( - mtp_diverse_stage2_single_token, -) from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size +from .stage1_single_token import mtp_diverse_stage1_single_token +from .stage2_single_token import mtp_diverse_stage2_single_token + @torch.no_grad() def token_decode_attention_mtp_diverse_single_token( diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py new file mode 100644 index 0000000000..9e76348baf --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py @@ -0,0 +1,58 @@ +from typing import Optional + +import torch + + +def token_decode_attention_flash_decoding( + q: torch.Tensor, + batch, + cache_k: torch.Tensor, + cache_k_scale: torch.Tensor, + cache_v: torch.Tensor, + cache_v_scale: torch.Tensor, + max_kv_seq_len: int, + out: Optional[torch.Tensor] = None, + alloc_tensor_func=torch.empty, +): + BLOCK_SEQ = 256 + batch_size = batch.batch_size + q_head_num = q.shape[1] + head_dim = q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 + # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 + # 的显存消耗和性能的平衡。 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=q.dtype, device="cuda") + + from .int4kv_flash_decoding_stage1 import int4kv_flash_decode_stage1 + + int4kv_flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + max_kv_seq_len=max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + from ..int8kv.normal.int8kv_flash_decoding_stage2 import flash_decode_stage2 + + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py new file mode 100644 index 0000000000..c5e3af3f36 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py @@ -0,0 +1,102 @@ +# 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 +import torch + +from lightllm.common.attention.api.batch import AttentionBatch +from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size + +from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1 +from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2 +from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3 + + +def token_decode_attention_flash_decoding( + q, + batch: AttentionBatch, + cache_k, + cache_k_scale, + cache_v, + cache_v_scale, + out=None, + alloc_tensor_func=torch.empty, + b_shared_seq_len=None, + b_mark_shared_group=None, + shared_streams_dict=None, +): + if b_shared_seq_len is None or b_mark_shared_group is None or shared_streams_dict is None: + raise ValueError("diverse attention requires prepared group metadata and slot-owned streams") + + if "stream1" not in shared_streams_dict: + shared_streams_dict["stream1"] = torch.cuda.Stream() + if "stream2" not in shared_streams_dict: + shared_streams_dict["stream2"] = torch.cuda.Stream() + + stream1 = shared_streams_dict["stream1"] + stream2 = shared_streams_dict["stream2"] + + q_head_num = q.shape[1] + head_dim = q.shape[2] + + BLOCK_SEQ = 256 + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + mid_o = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2, head_dim], dtype=torch.float32, device="cuda" + ) + mid_o_logexpsum = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2], dtype=torch.float32, device="cuda" + ) + + current_stream = torch.cuda.current_stream() + + stream1.wait_stream(current_stream) + with torch.cuda.stream(stream1): + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + b_shared_seq_len=b_shared_seq_len, + b_mark_shared_group=b_mark_shared_group, + max_len_in_batch=batch.max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + max_batch_group_size=get_diverse_max_batch_shared_group_size(), + ) + stream2.wait_stream(current_stream) + with torch.cuda.stream(stream2): + flash_decode_stage2( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + b_shared_seq_len=b_shared_seq_len, + max_len_in_batch=batch.max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + current_stream.wait_stream(stream1) + current_stream.wait_stream(stream2) + + flash_diverse_decode_stage3( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + b_shared_seq_len=b_shared_seq_len, + O=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py index e69de29bb2..8cc4aa919b 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py @@ -0,0 +1 @@ +from .int8kv_flash_decoding import token_decode_attention_flash_decoding diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py new file mode 100644 index 0000000000..cb4cad2e8e --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py @@ -0,0 +1,67 @@ +from typing import Optional + +import torch + +from lightllm.common.attention.api.batch import AttentionBatch + +from .int8kv_flash_decoding_stage1 import flash_decode_stage1 +from .int8kv_flash_decoding_stage2 import flash_decode_stage2 + + +@torch.no_grad() +def token_decode_attention_flash_decoding( + q: torch.Tensor, + batch: AttentionBatch, + cache_k: torch.Tensor, + cache_k_scale: torch.Tensor, + cache_v: torch.Tensor, + cache_v_scale: torch.Tensor, + max_len_in_batch: int, + out: Optional[torch.Tensor] = None, + alloc_tensor_func=torch.empty, +): + q_head_num = q.shape[1] + head_dim = q.shape[2] + + BLOCK_SEQ = 256 + batch_size = batch.batch_size + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 + # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 + # 的显存消耗和性能的平衡。 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_seq_len=batch.b_seq_len, + max_len_in_batch=max_len_in_batch, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + flash_decode_stage2( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + O=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py new file mode 100644 index 0000000000..94489f2be9 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py @@ -0,0 +1,36 @@ +import torch + + +def token_decode_attention_flash_decoding(q, batch, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty): + BLOCK_SEQ = 256 + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + q_head_num, head_dim = q.shape[1], q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + from .flash_decoding_stage1 import flash_decode_stage1 + from .flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + mid_o = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda" + ) + mid_o_logexpsum = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda" + ) + + flash_decode_stage1( + q.view(calcu_shape1), + cache_k, + cache_v, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + batch.max_kv_seq_len, + mid_o, + mid_o_logexpsum, + BLOCK_SEQ, + ) + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py new file mode 100644 index 0000000000..3ca3f2dba2 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py @@ -0,0 +1,147 @@ +import torch +import triton +import triton.language as tl +from packaging.version import Version + +if Version(triton.__version__) >= Version("2.1.0"): + + @triton.jit + def _fwd_kernel( + Q, + K, + V, + sm_scale, + Out, + q_stride_s, + q_stride_h, + q_stride_d, + k_stride_s, + k_stride_h, + k_stride_d, + v_stride_s, + v_stride_h, + v_stride_d, + o_stride_s, + o_stride_h, + o_stride_d, + head_dim_act, + cu_seqlens, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + cur_batch = tl.program_id(2) + cur_head = tl.program_id(1) + start_m = tl.program_id(0) + + seq_start = tl.load(cu_seqlens + cur_batch).to(tl.int32) + seq_end = tl.load(cu_seqlens + cur_batch + 1).to(tl.int32) + seq_len = seq_end - seq_start + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + + mask_d = offs_d < head_dim_act + off_q = cur_head * q_stride_h + (seq_start + offs_m[:, None]) * q_stride_s + offs_d[None, :] * q_stride_d + q = tl.load(Q + off_q, mask=(offs_m[:, None] < seq_len) & mask_d[None, :], other=0.0) + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + for start_n in range(0, seq_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + off_k = ( + (seq_start + start_n + offs_n[None, :]) * k_stride_s + + cur_head * k_stride_h + + offs_d[:, None] * k_stride_d + ) + k = tl.load(K + off_k, mask=((start_n + offs_n[None, :]) < seq_len) & mask_d[:, None], other=0.0) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) + qk *= sm_scale + qk += tl.where((start_n + offs_n[None, :]) < seq_len, 0, float("-inf")) + + # -- compute m_ij, p, l_ij + m_ij = tl.maximum(tl.max(qk, 1), m_i) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + + acc_scale = tl.exp(m_i - m_ij) + acc = acc * acc_scale[:, None] + + # update acc + off_v = ( + (seq_start + start_n + offs_n[:, None]) * v_stride_s + + cur_head * v_stride_h + + offs_d[None, :] * v_stride_d + ) + v = tl.load(V + off_v, mask=((start_n + offs_n[:, None]) < seq_len) & mask_d[None, :], other=0.0) + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + m_i = m_ij + l_i_new = tl.exp(l_i - m_ij) + l_ij + l_i = m_ij + tl.log(l_i_new) + + o_scale = tl.exp(m_i - l_i) + acc = acc * o_scale[:, None] + # initialize pointers to output + off_o = (seq_start + offs_m[:, None]) * o_stride_s + cur_head * o_stride_h + offs_d[None, :] * o_stride_d + out_ptrs = Out + off_o + tl.store(out_ptrs, acc, mask=(offs_m[:, None] < seq_len) & mask_d[None, :]) + return + + @torch.no_grad() + def _flash_attention_triton_fwd( + q, + k, + v, + o, + cu_seqlens, # q k v cu_seqlens, + max_seqlen, + ): + BLOCK = 64 + # shape constraints + assert q.ndim == k.ndim == v.ndim == o.ndim == 3, "q, k, v, o must be 3D tensors" + _, head_num, head_dim = q.shape + batch_size = cu_seqlens.numel() - 1 + + sm_scale = 1.0 / (head_dim**0.5) # 计算scale系数 + d_pad = triton.next_power_of_2(head_dim) + grid = (triton.cdiv(max_seqlen, BLOCK), head_num, batch_size) # batch, head, + num_warps = 4 + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + o, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + head_dim, + cu_seqlens, + BLOCK_M=BLOCK, + BLOCK_DMODEL=d_pad, + BLOCK_N=BLOCK, + num_warps=num_warps, + num_stages=2, + ) + return + +else: + raise RuntimeError(f"varlen attention requires Triton >= 2.1.0, got {triton.__version__}") diff --git a/lightllm/common/attention/backends/dense/triton/neo.py b/lightllm/common/attention/backends/dense/triton/neo.py new file mode 100644 index 0000000000..60dd33da50 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/neo.py @@ -0,0 +1,48 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend, TritonPrefillContext + + +@dataclasses.dataclass(slots=True) +class NeoTritonPrefillContext(TritonPrefillContext): + pass + + +class NeoTritonAttBackend(TritonAttBackend): + provider = "triton" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NeoTritonPrefillContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) + + out = alloc_func(q.shape, q.dtype) + context_attention_fwd_neo( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + context.batch.b_image_token_end, + ) + return out diff --git a/lightllm/common/attention/backends/dense/triton/varlen.py b/lightllm/common/attention/backends/dense/triton/varlen.py new file mode 100644 index 0000000000..e0e0f6dc9a --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/varlen.py @@ -0,0 +1,23 @@ +from lightllm.common.attention.backends.dense.triton.kernels.prefill.varlen import _flash_attention_triton_fwd +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class TritonVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "triton" + dependency = "triton" + supports_graph = True + + def forward(self, inputs, *, out): + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + if q.shape[0] == 0: + return out + _flash_attention_triton_fwd( + q, + k, + v, + out, + cu_seqlens, # q k v cu_seqlens, + max_seqlen, + ) + return out diff --git a/lightllm/common/attention/backends/dense/varlen.py b/lightllm/common/attention/backends/dense/varlen.py new file mode 100644 index 0000000000..5f6766f115 --- /dev/null +++ b/lightllm/common/attention/backends/dense/varlen.py @@ -0,0 +1,41 @@ +import torch + +from lightllm.common.attention._environment import check_attention_environment +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionOp +from lightllm.platform import Platform + + +class BaseVarlenAttentionOp(VarlenAttentionOp): + provider = "" + + supported_platforms = frozenset({Platform.CUDA}) + dependency = "" + min_device_capability = None + activation_dtypes = frozenset({torch.float16, torch.bfloat16, torch.float32}) + max_head_size = None + head_alignment = 1 + + @classmethod + def check_environment(cls, caps) -> None: + check_attention_environment( + caps, + provider=cls.provider, + supported_platforms=cls.supported_platforms, + min_device_capability=cls.min_device_capability, + dependency=cls.dependency, + ) + + @classmethod + def check_support(cls, inputs, caps=None) -> None: + inputs.validate() + q = inputs.query + if min(q.shape[1:]) <= 0: + raise UnsupportedAttention("head count and dimension must be positive") + if q.dtype not in cls.activation_dtypes: + raise UnsupportedAttention(f"unsupported activation dtype {q.dtype}") + if caps is not None and caps.platform not in cls.supported_platforms: + raise UnsupportedAttention(f"{cls.provider} does not support platform {caps.platform.value}") + if cls.max_head_size is not None and q.shape[-1] > cls.max_head_size: + raise UnsupportedAttention(f"head dimension exceeds {cls.max_head_size}") + if q.shape[-1] % cls.head_alignment: + raise UnsupportedAttention(f"head dimension must be divisible by {cls.head_alignment}") diff --git a/lightllm/common/attention/backends/dense/xformers/varlen.py b/lightllm/common/attention/backends/dense/xformers/varlen.py new file mode 100644 index 0000000000..ec1eb77624 --- /dev/null +++ b/lightllm/common/attention/backends/dense/xformers/varlen.py @@ -0,0 +1,37 @@ +import torch + +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class XformersVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "xformers" + dependency = "xformers.ops" + + @torch.no_grad() + def forward(self, inputs, *, out): + inputs.validate() + from xformers import ops as xformers_ops + from xformers.ops import fmha + + q, k, v, cu_seqlens, max_seqlen = inputs + o = out + assert q.ndim == k.ndim == v.ndim == o.ndim == 3 + assert cu_seqlens is not None and cu_seqlens.ndim == 1 + assert q.shape == k.shape == v.shape == o.shape + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int64).tolist() + seqlens = [int(L) for L in seqlens if int(L) > 0] + + if len(seqlens) == 0: + return o + if max_seqlen: + assert max(seqlens) <= max_seqlen + + attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens, device=q.device) + + q_ = q.unsqueeze(0) # [1, T, H, D] + k_ = k.unsqueeze(0) # [1, T, H, D] + v_ = v.unsqueeze(0) # [1, T, H, D] + + out = xformers_ops.memory_efficient_attention(q_, k_, v_, attn_bias=attn_bias, p=0.0) + o.copy_(out.squeeze(0)) # [T, H, D] + return o diff --git a/lightllm/common/attention/backends/gdn/flashqla/backend.py b/lightllm/common/attention/backends/gdn/flashqla/backend.py new file mode 100644 index 0000000000..5aee2aa86a --- /dev/null +++ b/lightllm/common/attention/backends/gdn/flashqla/backend.py @@ -0,0 +1,17 @@ +import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.backends.gdn.triton.backend import TritonLinearAttBackend + + +class FlashQlaLinearAttBackend(TritonLinearAttBackend): + provider = "flashqla" + dependency = "flash_qla" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supported_phases = frozenset({AttentionPhase.PREFILL}) + cache_formats = frozenset({"gdn"}) + + def get_prefill_kernel(self): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule diff --git a/lightllm/common/attention/backends/gdn/replay/backend.py b/lightllm/common/attention/backends/gdn/replay/backend.py new file mode 100644 index 0000000000..dc112936f2 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/replay/backend.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs + +from .decode import replay + + +@dataclasses.dataclass(slots=True) +class ReplayGDNPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayGDNDecodeContext(LinearDecodeContext): + pass + + +class ReplayGDNTritonAttBackend( + GDNMetadata, BaseAttentionBackend[ReplayGDNPrefillContext | ReplayGDNDecodeContext, GDNInputs, GDNOutput] +): + provider = "triton" + family = AttentionFamily.GDN + cache_formats = frozenset({"gdn"}) + + def __init__(self, model): + super().__init__(model=model) + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + self._init_linear_layer_metadata(model.config, model.tp_world_size_) + self.prefill_kernel = chunk_gated_delta_rule + + def _allocate_prefill_context(self, capacity, capture, workspace): + return ReplayGDNPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return ReplayGDNDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None): + backend = self + conv_states, ssm_states = layer.params.conv, layer.params.ssm + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = backend._split_qkvzba(inputs.mixed_qkvzba) + core_attn_out = self._prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + def _prefill(self, layer, context, mixed_qkv, conv_states, ssm_states, a, b): + backend = self + g, beta = fused_gdn_gating(layer.params.a_log, a, b, layer.params.dt_bias) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer.params.conv_weight, + bias=layer.params.conv_bias, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ).transpose(0, 1) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + core_attn_out, last_state = backend.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[context.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=True) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None): + backend = self + mixed_qkv, z, b, a = backend._split_qkvzba(inputs.mixed_qkvzba) + state = layer.params.replay_state + conv_states = layer.params.conv + draft_step = backend.decode_draft_step + if draft_step > 0: + core = self._mtp(layer, context, mixed_qkv, conv_states, a, b, draft_step, state) + else: + core, z = self._decode(layer, context, mixed_qkv, z, conv_states, a, b, state) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return core, z + + def _decode(self, layer, context, mixed_qkv, z, conv_states, a, b, state): + backend = self + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.params.conv_weight, + layer.params.conv_bias, + context.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) + key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) + value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.params.a_log, + layer.params.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output, z + + def _mtp(self, layer, context, mixed_qkv, conv_states, a, b, draft_step, state): + backend = self + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.params.conv_weight, + mtp_step=draft_step, + bias=layer.params.conv_bias, + activation=backend.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.params.a_log, + layer.params.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output + + +class ReplayGDNFlashQlaAttBackend(ReplayGDNTritonAttBackend): + provider = "flashqla" + dependency = "flash_qla" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supported_phases = frozenset({AttentionPhase.PREFILL}) + + def __init__(self, model): + BaseAttentionBackend.__init__(self, model=model) + from flash_qla import chunk_gated_delta_rule + + self._init_linear_layer_metadata(model.config, model.tp_world_size_) + self.prefill_kernel = chunk_gated_delta_rule diff --git a/lightllm/common/attention/backends/gdn/triton/backend.py b/lightllm/common/attention/backends/gdn/triton/backend.py new file mode 100644 index 0000000000..f24ac01028 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/triton/backend.py @@ -0,0 +1,226 @@ +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.batch import AttentionBatch +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +@dataclasses.dataclass(slots=True) +class LinearAttPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class LinearAttDecodeContext(LinearDecodeContext): + pass + + +class TritonLinearAttBackend( + GDNMetadata, BaseAttentionBackend[LinearAttPrefillContext | LinearAttDecodeContext, GDNInputs, GDNOutput] +): + provider = "triton" + cache_formats = frozenset({"gdn"}) + family = AttentionFamily.GDN + + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) + self.prefill_kernel = self.get_prefill_kernel() + + def get_prefill_kernel(self): + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _allocate_prefill_context(self, capacity, capture, workspace): + return LinearAttPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return LinearAttDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + backend: TritonLinearAttBackend = self + mtp_step = backend.mtp_step + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * (mtp_step + 1) + return + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> torch.Tensor: + mixed_qkvzba: torch.Tensor = inputs.mixed_qkvzba + backend: TritonLinearAttBackend = self + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) + core_attn_out = self._prefill_gdn_prefill_kernel( + layer, context, mixed_qkv, conv_states, ssm_states, a, b, context.batch + ) + return GDNOutput(core_attn_out, z) + + def _prefill_gdn_prefill_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + g, beta = fused_gdn_gating(layer.params.a_log, a, b, layer.params.dt_bias) + mixed_qkv = mixed_qkv.transpose(0, 1) + backend: TritonLinearAttBackend = self + out_tensor = causal_conv1d_fn( + mixed_qkv, + layer.params.conv_weight, + bias=layer.params.conv_bias, + query_start_loc=batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ) + mixed_qkv = out_tensor.transpose(0, 1) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + initial_state = ssm_states[context.b_ssm_buffer_idx] + core_attn_out, last_recurrent_state = backend.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=initial_state, + output_final_state=True, + cu_seqlens=batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=False) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> torch.Tensor: + mixed_qkvzba = inputs.mixed_qkvzba + backend: TritonLinearAttBackend = self + mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + draft_step = self.decode_draft_step + if draft_step > 0: + core_attn_out = self._decode_gdn_mtp_kernel( + layer, context, mixed_qkv, conv_states, ssm_states, a, b, context.batch + ) + else: + core_attn_out, z = self._decode_gdn_decode_kernel( + layer, context, mixed_qkv, z, conv_states, ssm_states, a, b, context.batch + ) + return GDNOutput(core_attn_out, z) + + def _decode_gdn_decode_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + z: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + backend: TritonLinearAttBackend = self + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.params.conv_weight, + layer.params.conv_bias, + context.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + core_attn_out, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + use_qk_l2norm_in_kernel=True, + A_log=layer.params.a_log, + dt_bias=layer.params.dt_bias, + a_raw=a, + b_raw=b, + ) + return GDNOutput(core_attn_out, z) + + def _decode_gdn_mtp_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + backend: TritonLinearAttBackend = self + cu_seqlens_q = context.b1_mtp_cu_q_seq_len + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.params.conv_weight, + mtp_step=backend.decode_draft_step, + bias=layer.params.conv_bias, + activation=backend.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=cu_seqlens_q, + ) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) + assert context.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q.to(torch.long), + ssm_state_indices=context.b_ssm_buffer_idx, + ssm_state_write_indices=context.b_ssm_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + A_log=layer.params.a_log, + dt_bias=layer.params.dt_bias, + a_raw=a, + b_raw=b, + ) + return core_attn_out diff --git a/lightllm/common/attention/backends/kda/replay/backend.py b/lightllm/common/attention/backends/kda/replay/backend.py new file mode 100644 index 0000000000..0d928c5136 --- /dev/null +++ b/lightllm/common/attention/backends/kda/replay/backend.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate + +from .decode import replay + + +@dataclasses.dataclass(slots=True) +class ReplayKDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayKDADecodeContext(LinearDecodeContext): + pass + + +class ReplayKDALinearAttBackend( + BaseAttentionBackend[ReplayKDAPrefillContext | ReplayKDADecodeContext, KDAInputs, torch.Tensor] +): + provider = "triton" + family = AttentionFamily.KDA + cache_formats = frozenset({"kda"}) + + def __init__(self, model): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return ReplayKDAPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return ReplayKDADecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + backend: ReplayKDALinearAttBackend = self + conv_states, ssm_states = layer.params.conv, layer.params.ssm + conv_states = conv_states[..., : backend.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + inputs.mixed_qkv.transpose(0, 1), + layer.params.conv_weight, + bias=None, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) + k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) + v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) + raw_gate = inputs.raw_gate.view(1, -1, backend.tp_hidden_size) + raw_beta = inputs.raw_beta.view(1, -1, backend.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.params.a_log, + g_bias=layer.params.dt_bias, + initial_state=ssm_states[context.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=True) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + backend: ReplayKDALinearAttBackend = self + state = layer.params.replay_state + conv_states = layer.params.conv + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.decode_draft_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.decode_draft_step, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + inputs.mixed_qkv, + conv_states, + layer.params.conv_weight, + **conv_kwargs, + ) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(-1, backend.tp_num_heads, backend.head_dim) + k = k.view(-1, backend.tp_num_heads, backend.head_dim) + v = v.view(-1, backend.tp_num_heads, backend.head_dim) + raw_gate = inputs.raw_gate.reshape(-1, backend.tp_hidden_size) + raw_beta = inputs.raw_beta.reshape(-1, backend.tp_num_heads) + storage_layer = state.storage_layer(layer.logical_layer_id) + output = replay( + state_ssm=state.ssm_state[storage_layer], + state_keys=state.keys[storage_layer], + state_deltas=state.deltas[storage_layer], + state_gates=state.gates[storage_layer], + state_raw_keys=state.raw_keys[storage_layer], + state_raw_values=state.raw_values[storage_layer], + state_betas=state.betas[storage_layer], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=raw_gate, + b=raw_beta, + a_log=layer.params.a_log, + bias=layer.params.dt_bias, + reqs=context.b_conv_buffer_idx, + cu_seqlens=context.b_cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + ) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return output diff --git a/lightllm/common/attention/backends/kda/triton/backend.py b/lightllm/common/attention/backends/kda/triton/backend.py new file mode 100644 index 0000000000..40cccc1be2 --- /dev/null +++ b/lightllm/common/attention/backends/kda/triton/backend.py @@ -0,0 +1,144 @@ +"""KDA attention backend for GLM-5-Next.""" +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +@dataclasses.dataclass(slots=True) +class KDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class KDADecodeContext(LinearDecodeContext): + pass + + +class KDALinearAttBackend(BaseAttentionBackend[KDAPrefillContext | KDADecodeContext, KDAInputs, torch.Tensor]): + provider = "triton" + cache_formats = frozenset({"kda"}) + family = AttentionFamily.KDA + + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def _allocate_prefill_context(self, capacity, capture, workspace): + return KDAPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return KDADecodeContext(capacity, capture, workspace) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def _update_prefill(self, context): + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * (self.mtp_step + 1) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + backend: KDALinearAttBackend = self + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + conv_states = conv_states[..., : backend.conv_kernel_size - 1] + q, k, v = causal_conv1d_qkv( + mixed_qkv, + layer.params.conv_weight, + conv_states, + context.batch.b1_cu_q_seq_len, + context.b_conv_buffer_idx, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + ) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + initial_state = ssm_states[context.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.params.a_log, + g_bias=layer.params.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + update_linear_decode(context, self.mtp_step, self.uses_dynamic_spec_verify_layout(), replay=False) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + backend: KDALinearAttBackend = self + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update(mixed_qkv, conv_states, layer.params.conv_weight, **conv_kwargs) + q, k, v = backend.split_qkv(mixed_qkv) + shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) + q = q.view(*shape, backend.tp_num_heads, backend.head_dim) + k = k.view(*shape, backend.tp_num_heads, backend.head_dim) + v = v.view(*shape, backend.tp_num_heads, backend.head_dim) + raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) + raw_beta = raw_beta.view(*shape, backend.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer.params.a_log, + gate_bias=layer.params.dt_bias, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + cu_seqlens=context.b1_mtp_cu_q_seq_len, + num_accepted_tokens=(context.accepted_tokens(layer.params.mtp_state_index) if backend.mtp_step else None), + ) + return output diff --git a/lightllm/common/attention/backends/mla/fa3/backend.py b/lightllm/common/attention/backends/mla/fa3/backend.py new file mode 100644 index 0000000000..ac781b0d95 --- /dev/null +++ b/lightllm/common/attention/backends/mla/fa3/backend.py @@ -0,0 +1,124 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend +from lightllm.utils.sgl_utils import flash_attn_varlen_func, flash_attn_with_kvcache + + +@dataclasses.dataclass(slots=True) +class MlaFa3PrefillContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class MlaFa3DecodeContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class MlaFa3AttBackend(Fa3AttBackend): + provider = "fa3" + max_head_size = None + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaFa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaFa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.causal = self.uses_causal_attention() + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, alloc_func=torch.empty + ) -> torch.Tensor: + k_nope, k_rope = k + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + assert q.ndim == 3 and k.ndim == 3 and (v.ndim == 3) + softmax_scale = layer.params.softmax_scale + o_tensor = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + max_seqlen_k=context.batch.max_kv_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + return_softmax_lse=False, + ) + return o_tensor + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs=None, + alloc_func=torch.empty, + ): + q_nope, q_rope = q + kv = k + qk_rope_head_dim = layer.params.qk_rope_head_dim + kv_lora_rank = kv.shape[-1] - qk_rope_head_dim + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, self.infer_page_size, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, self.infer_page_size, 1, kv_lora_rank) + k_descale, v_descale = (None, None) + softmax_scale = layer.params.softmax_scale + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/mla/flashinfer/backend.py b/lightllm/common/attention/backends/mla/flashinfer/backend.py new file mode 100644 index 0000000000..9d4251b689 --- /dev/null +++ b/lightllm/common/attention/backends/mla/flashinfer/backend.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Tuple + +import torch +import triton + +from lightllm.common.attention.api import BackendResources +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.flashinfer_mla_plan import fill_mla_decode_plan_for_cuda_graph +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.mla import BatchMLAPagedAttentionWrapper + from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: BatchPrefillWithRaggedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferDecodeContext(AttentionExecutionContext): + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + q_indptr: torch.Tensor = None + decode_wrapper: BatchMLAPagedAttentionWrapper | None = None + + +class MlaFlashInferAttBackend( + BaseAttentionBackend[ + MlaFlashInferPrefillContext | MlaFlashInferDecodeContext, MLAPrefillInputs | MLADecodeInputs, torch.Tensor + ] +): + provider = "flashinfer" + dependency = "flashinfer" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_expert_parallel = False + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + workspace_buffer_key = "flashinfer_mla" + workspace_buffer_size = 256 * 1024 * 1024 + + def __init__(self, model): + set_flashinfer_envs() + super().__init__(model=model) + num_heads = model.config["num_attention_heads"] + self.tp_q_head_num = num_heads // get_dp_world_size() + self.qk_nope_head_dim = model.qk_nope_head_dim + self.qk_rope_head_dim = model.qk_rope_head_dim + self.kv_lora_rank = model.kv_lora_rank + self.v_head_dim = model.v_head_dim + self.q_data_type = model.data_type + self.kv_data_type = model.data_type + self.max_seq_length = model.max_seq_length + self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) + self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5) + from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale + + if model.config["rope_scaling"] is not None: + rope_scaling = model.config["rope_scaling"] + mscale_all_dim = rope_scaling.get("mscale_all_dim", 0) + scaling_factor = rope_scaling["factor"] + if mscale_all_dim: + mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + return + + def configure_spec(self, spec): + self.tp_q_head_num = spec.num_q_heads + + def resource_requirements(self, capacity, *, capture): + return BackendResources( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaFlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaFlashInferDecodeContext(capacity, capture, workspace) + + def _prepare_prefill(self, context): + import flashinfer + + q_starts = context.batch.b1_cu_q_seq_len.int() + kv_starts = context.batch.b1_cu_kv_seq_len.int() + if context.prefill_wrapper is None: + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), "NHD" + ) + context.prefill_wrapper.plan( + qo_indptr=q_starts, + kv_indptr=kv_starts, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_q_head_num, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=True, + sm_scale=self.softmax_scale, + ) + return + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + k_nope, k_rope = k + o_tensor = alloc_func((q.shape[0], q.shape[1], v.shape[-1]), q.dtype, device=q.device) + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + context.prefill_wrapper.run(q, k, v, out=o_tensor) + return o_tensor + + def _decode_refresh_cuda_graph_decode_plan(self, context, kv_lens: torch.Tensor, max_kv_len: int): + fill_mla_decode_plan_for_cuda_graph( + decode_wrapper=context.decode_wrapper, + kv_page_indptr=context.kv_starts, + kv_lens=kv_lens, + batch_size=context.batch.batch_size, + num_heads=self.tp_q_head_num, + max_kv_len=max_kv_len, + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device) + context.decode_wrapper.run( + q_nope, + q_rope, + k[:, :, :-qk_rope_head_dim].view(-1, self.infer_page_size, 1, k.shape[-1] - qk_rope_head_dim), + k[:, :, -qk_rope_head_dim:].view(-1, self.infer_page_size, 1, qk_rope_head_dim), + out=o_tensor, + return_lse=False, + ) + return o_tensor + + def _prepare_decode(self, context): + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + if context.kv_starts is None: + context.kv_starts = starts + context.kv_indices = context.workspace.get( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + context.q_indptr = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + else: + context.kv_starts.copy_(starts, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.infer_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + use_cuda_graph=context.capture, + qo_indptr=context.q_indptr, + kv_indices=context.kv_indices, + kv_indptr=context.kv_starts, + kv_len_arr=batch.b_seq_len, + ) + context.decode_wrapper.plan( + context.q_indptr, + context.kv_starts, + context.kv_indices, + batch.b_seq_len, + self.tp_q_head_num, + self.kv_lora_rank, + self.qk_rope_head_dim, + self.infer_page_size, + False, + self.softmax_scale, + self.q_data_type, + self.kv_data_type, + ) + else: + self._decode_refresh_cuda_graph_decode_plan(context, batch.b_seq_len, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/mla/triton/backend.py b/lightllm/common/attention/backends/mla/triton/backend.py new file mode 100644 index 0000000000..3b42d36daf --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/backend.py @@ -0,0 +1,164 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend + + +@dataclasses.dataclass(slots=True) +class MlaTritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass(slots=True) +class MlaTritonDecodeContext(AttentionExecutionContext): + block_seq: torch.Tensor = None + batch_start: torch.Tensor = None + indices_initialized: bool = False + + +class MlaTritonAttBackend( + BaseAttentionBackend[ + MlaTritonPrefillContext | MlaTritonDecodeContext, MLAPrefillInputs | MLADecodeInputs, torch.Tensor + ] +): + provider = "triton" + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + if q.shape[2] == 256: + return self._prefill_contiguous_kv_prefill_att( + layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func + ) + else: + return self._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends.mla.triton.kernels.prefill import context_attention_fwd_with_v + + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = (q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:]) + o_tensor = alloc_func((q_nope.shape[0], q_nope.shape[1], v.shape[-1]), dtype=q_nope.dtype, device=q_nope.device) + k_nope, k_rope = k + softmax_scale = layer.params.softmax_scale + context_attention_fwd_with_v( + q_nope, + q_rope, + k_nope, + k_rope, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + softmax_scale, + ) + return o_tensor + + def _prefill_contiguous_kv_prefill_att( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + k_nope, k_rope = k + q_head_num = q.shape[1] + k_merged = torch.cat([k_nope.expand(-1, q_head_num, -1), k_rope.expand(-1, q_head_num, -1)], dim=-1) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + context_attention_fwd_contiguous_kv( + q, + k_merged, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.max_q_seq_len, + context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.block_seq = torch.empty((1,), dtype=torch.int64, device=context.batch.device) + context.batch_start = torch.empty((context.batch.batch_size,), dtype=torch.int64, device=context.batch.device) + context.indices_initialized = False + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + softmax_scale = layer.params.softmax_scale + from lightllm.common.attention.backends.mla.triton.kernels.decode import ( + gqa_token_decode_attention_flash_decoding, + ) + + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + kv = k + out = gqa_token_decode_attention_flash_decoding( + q_nope=q_nope, + q_rope=q_rope, + kv_nope=kv[:, :, :-qk_rope_head_dim], + kv_rope=kv[:, :, -qk_rope_head_dim:], + batch=context.batch, + softmax_scale=softmax_scale, + alloc_tensor_func=alloc_func, + block_seq=context.block_seq, + batch_start=context.batch_start, + initialize=not context.indices_initialized, + ) + context.indices_initialized = True + return out diff --git a/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py index e69de29bb2..fb0609a401 100644 --- a/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py @@ -0,0 +1 @@ +from .gqa_flash_decoding import gqa_token_decode_attention_flash_decoding diff --git a/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py new file mode 100644 index 0000000000..a618c9cdea --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py @@ -0,0 +1,147 @@ +import os +from typing import List + +import torch +import torch.multiprocessing as mp +import triton +import triton.language as tl + +from lightllm.utils.device_utils import get_device_sm_count +from lightllm.utils.log_utils import init_logger + +from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig + +logger = init_logger(__name__) + + +def gqa_token_decode_attention_flash_decoding( + q_nope, + q_rope, + kv_nope, + kv_rope, + batch, + softmax_scale, + out=None, + alloc_tensor_func=torch.empty, + block_seq=None, + batch_start=None, + initialize=True, + **run_config, +): + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + + q_head_num, kv_lora_rank = q_nope.shape[1], q_nope.shape[2] + q_rope_dim = q_rope.shape[2] + assert q_rope_dim == 64 + + calcu_shape1 = (batch_size, q_head_num, kv_lora_rank) + calcu_shape2 = (batch_size, q_head_num, q_rope_dim) + + if not run_config: + if torch.cuda.is_current_stream_capturing(): + avg_seq_len_in_batch = max_kv_seq_len + else: + avg_seq_len_in_batch = batch.total_token_num // batch_size + + run_config = MlaDecodeAttentionKernelConfig.try_to_get_best_config( + batch_size=batch_size, + avg_seq_len_in_batch=avg_seq_len_in_batch, + q_head_num=q_head_num, + q_head_dim=kv_lora_rank, + q_rope_dim=q_rope_dim, + out_dtype=torch.bfloat16, + ) + + BLOCK_N = run_config["BLOCK_N"] + + from .gqa_flash_decoding_stage1 import flash_decode_stage1 + from .gqa_flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q_nope.shape, q_nope.dtype, q_nope.device) if out is None else out + + fake_decode_att_block_seq = torch.empty([0], dtype=torch.int64, device="cuda") + mid_o = torch.empty([q_head_num, 0, kv_lora_rank], dtype=torch.float32, device="cuda") + mid_o_logexpsum = torch.empty([q_head_num, 0], dtype=torch.float32, device="cuda") + + vsm_count = flash_decode_stage1( + fake_decode_att_block_seq, + q_nope.view(calcu_shape1), + q_rope.view(calcu_shape2), + kv_nope, + kv_rope, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + mid_o, + mid_o_logexpsum, + softmax_scale, + get_sm_count=True, + **run_config, + ) + + if initialize: + _fwd_kernel_calcu_index_and_block_seq[(1,)]( + batch.b_seq_len, + block_seq, + batch_start, + vsm_count, + batch_size, + BLOCK_N=BLOCK_N, + MAX_BATCH_SIZE=triton.next_power_of_2(batch_size), + num_warps=4, + ) + + mid_o = torch.empty([q_head_num, vsm_count * 4 + batch_size, kv_lora_rank], dtype=torch.float32, device="cuda") + mid_o_logexpsum = torch.empty([q_head_num, vsm_count * 4 + batch_size], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + block_seq, + q_nope.view(calcu_shape1), + q_rope.view(calcu_shape2), + kv_nope, + kv_rope, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + mid_o, + mid_o_logexpsum, + softmax_scale, + get_sm_count=False, + **run_config, + ) + + flash_decode_stage2( + block_seq, batch_start, mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), **run_config + ) + return o_tensor + + +@triton.jit +def _fwd_kernel_calcu_index_and_block_seq( + b_seq_len_ptr, + mid_o_decode_att_block_seq_ptr, + mid_o_batch_start_index_ptr, + num_sm, + batch_size, + BLOCK_N: tl.constexpr, + MAX_BATCH_SIZE: tl.constexpr, +): + b_seq_len = tl.load( + b_seq_len_ptr + tl.arange(0, MAX_BATCH_SIZE), mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, other=0 + ) + total_token_num = tl.sum(b_seq_len) + + block_seq = tl.cast(total_token_num / (num_sm * 4), dtype=tl.int32) + 1 + block_seq = tl.cdiv(block_seq, BLOCK_N) * BLOCK_N + + block_seq_len = tl.cdiv(b_seq_len, block_seq) + cumsum_seq_len = tl.cumsum(block_seq_len) + batch_start_index = cumsum_seq_len - block_seq_len + tl.store( + mid_o_batch_start_index_ptr + tl.arange(0, MAX_BATCH_SIZE), + batch_start_index, + mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, + ) + tl.store(mid_o_decode_att_block_seq_ptr, block_seq) + return diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py new file mode 100644 index 0000000000..8da3941a64 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py @@ -0,0 +1,179 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.utils.dist_utils import get_current_device_id + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparsePrefillContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + # Query-row metadata is consumed by model-specific sparse indexers. It is + # kept on the context so the indexer never reaches into InferStateInfo. + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparseDecodeContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + nsa_cache_seqlens: torch.Tensor = None + nsa_cu_seqlens_k_new: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class NsaFlashMlaSparseAttBackend( + BaseAttentionBackend[ + NsaFlashMlaSparsePrefillContext | NsaFlashMlaSparseDecodeContext, + SparseMLAPrefillInputs | SparseMLADecodeInputs, + torch.Tensor, + ] +): + provider = "flashmla_sparse" + dependency = "flash_mla" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + cache_formats = frozenset({"None"}) + family = AttentionFamily.SPARSE_MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NsaFlashMlaSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return NsaFlashMlaSparseDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + batch = context.batch + context.ragged_mem_index = torch.empty(batch.total_token_num, dtype=torch.int32, device=get_current_device_id()) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=batch.b_seq_len, + b_q_seq_len=batch.b_q_seq_len, + b_req_idx=batch.b_req_idx, + req_to_token_index=batch.req_to_token_indexs, + q_token_num=batch.token_num, + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=batch.hold_req_id, + ) + self._update_query_metadata(context) + return + + @staticmethod + def _update_query_metadata(context): + """Materialize per-query request metadata on the execution context. + + Prefill positions are supplied by the model's input preparation and + therefore remain correct for prompt-cache and speculative layouts. + ``output_size`` keeps the operation graph-capturable without a host + read of the GPU-side query lengths. + """ + batch = context.batch + position_ids = batch.position_ids + if position_ids is None: + # Unit-level callers may construct an AttentionBatch directly. + # The runtime always supplies position_ids, but contiguous query + # ranges are an unambiguous fallback for those callers. + history = batch.b_seq_len - batch.b_q_seq_len + starts = torch.repeat_interleave(history, batch.b_q_seq_len, output_size=batch.token_num) + row_starts = torch.repeat_interleave( + batch.b1_cu_q_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + position_ids = starts + torch.arange(batch.token_num, device=batch.device, dtype=starts.dtype) - row_starts + context.q_kv_seq_len = position_ids + 1 + context.q_req_idx = torch.repeat_interleave(batch.b_req_idx, batch.b_q_seq_len, output_size=batch.token_num) + context.q_kv_start_loc = torch.repeat_interleave( + batch.b1_cu_kv_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: SparseMLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.params.kv + return self._prefill_nsa_prefill_att(layer, context, q=q, kv=k, inputs=inputs) + + def _prefill_nsa_prefill_att(self, layer, context, q: torch.Tensor, kv: torch.Tensor, inputs) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.params.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + mla_out, _, _ = flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_mem_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def _update_decode(self, context): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.get( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + context.nsa_cache_seqlens = torch.minimum( + torch.full(size=(context.batch.batch_size,), fill_value=2048, dtype=torch.int32, device="cuda"), + context.batch.b_seq_len, + ) + padded_seq_lens = torch.zeros(size=(context.nsa_cache_seqlens.shape[0] + 1,), dtype=torch.int32, device="cuda") + padded_seq_lens[1:].copy_(context.nsa_cache_seqlens, non_blocking=True) + context.nsa_cu_seqlens_k_new = padded_seq_lens.cumsum(dim=0, dtype=torch.int32) + + def _forward_decode( + self, layer: LayerAttentionBinding, context, inputs: SparseMLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._decode_nsa_decode_att(layer, context, q=q, kv=k, inputs=inputs) + + def _decode_nsa_decode_att( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], kv: torch.Tensor, inputs + ) -> torch.Tensor: + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.params.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=softmax_scale, + causal=True, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py new file mode 100644 index 0000000000..5e17ab6607 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py @@ -0,0 +1,149 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import NsaFlashMlaSparseAttBackend + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparsePrefillContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparseDecodeContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + flashmla_sched_meta: object = None + capture_flashmla_sched_meta: object = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class NsaFlashMlaFp8SparseAttBackend(NsaFlashMlaSparseAttBackend): + cache_formats = frozenset({"fp8kv_dsa"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NsaFlashMlaFp8SparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return NsaFlashMlaFp8SparseDecodeContext(capacity, capture, workspace) + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: SparseMLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.params.kv + return self._prefill_nsa_prefill_att(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _prefill_nsa_prefill_att( + self, layer, context, q: torch.Tensor, packed_kv: torch.Tensor, inputs + ) -> torch.Tensor: + import flash_mla + + topk_indices = inputs.topk_indices + softmax_scale = layer.params.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + topk_mem_indices = inputs.topk_mem_indices + prefill_cache_kv = inputs.prefill_cache_kv + if context.batch.max_cache_len > 0: + from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, + ) + + kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=packed_kv, + topk_mem_indices=topk_mem_indices, + prefill_mem_index=context.batch.mem_index, + prefill_cache_kv=prefill_cache_kv, + prefill_dtype=prefill_cache_kv.dtype, + ) + else: + kv = prefill_cache_kv + if topk_indices.ndim == 2: + topk_indices = topk_indices.unsqueeze(1) + mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def _update_decode(self, context): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.get( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + return + + def _forward_decode( + self, layer: LayerAttentionBinding, context, inputs: SparseMLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._decode_nsa_decode_att(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _decode_nsa_decode_att( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], packed_kv: torch.Tensor, inputs + ) -> torch.Tensor: + import flash_mla + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.params.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, q_rope = q + q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() + kv = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + o_tensor, _ = flash_mla.flash_mla_with_kvcache( + q=q_all, + k_cache=kv, + block_table=None, + cache_seqlens=None, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return o_tensor[:, 0, :, :] + + def _prepare_decode(self, context): + import flash_mla + + if context.flashmla_sched_meta is None: + context.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] diff --git a/lightllm/common/attention/probe.py b/lightllm/common/attention/probe.py new file mode 100644 index 0000000000..af344da994 --- /dev/null +++ b/lightllm/common/attention/probe.py @@ -0,0 +1,90 @@ +"""启动子进程中的小规模数值检查，不接触模型的真实 KV 或请求状态。""" + +from types import SimpleNamespace + +import torch + +from .api import AttentionBatch, AttentionPhase, ContextCapacity, ExecutionRequirements, LayerAttentionBinding +from .api.inputs import DenseInputs, DenseParams +from .backends._workspace import AttentionWorkspace + + +def probe_dense(backend_cls, spec, runtime): + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.dist_utils import get_dp_world_size + + page = spec.cache.page_size + stride = ((16 + page - 1) // page) * page + q_heads, kv_heads, dim = spec.num_q_heads, spec.cache.key_heads, spec.head_size + model = SimpleNamespace( + args=StartArgs(page_size=page, disable_cudagraph=True), + device_runtime=runtime, + attention_execution=ExecutionRequirements(), + config=dict( + num_attention_heads=q_heads * get_dp_world_size(), + num_key_value_heads=kv_heads * get_dp_world_size(), + hidden_size=q_heads * get_dp_world_size() * dim, + head_dim=dim, + ), + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=2), + max_seq_length=stride, + data_type=spec.dtype, + ) + backend = backend_cls.create(model, spec) + device = runtime.device + lengths = (7, 13) + q_lengths = (2, 3) if spec.phase is AttentionPhase.PREFILL else (1, 1) + query = torch.randn((sum(q_lengths), q_heads, dim), dtype=spec.dtype, device=device) + key, value = (torch.randn((3 * stride, kv_heads, dim), dtype=spec.dtype, device=device) for _ in range(2)) + batch = AttentionBatch( + phase=spec.phase, + batch_size=2, + token_num=sum(q_lengths), + total_token_num=sum(lengths), + max_q_seq_len=max(q_lengths), + max_kv_seq_len=max(lengths), + max_cache_len=max(n - q for n, q in zip(lengths, q_lengths)), + b_req_idx=torch.tensor([0, 1], dtype=torch.int32, device=device), + b_seq_len=torch.tensor(lengths, dtype=torch.int32, device=device), + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device=device), + b1_cu_q_seq_len=torch.tensor([0, q_lengths[0], sum(q_lengths)], dtype=torch.int32, device=device), + b1_cu_kv_seq_len=torch.tensor([0, lengths[0], sum(lengths)], dtype=torch.int32, device=device), + req_to_token_indexs=torch.arange(3 * stride, dtype=torch.int32, device=device).view(3, stride), + hold_req_id=2, + b_ready_cache_len=torch.tensor([n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device=device), + ) + capacity = ContextCapacity(spec.phase, 2, sum(q_lengths), stride, max(q_lengths)) + workspace = AttentionWorkspace(runtime=runtime) + workspace.reserve(backend.resource_requirements(capacity, capture=False)) + context = backend.allocate_context(capacity, capture=False, workspace=workspace) + backend.prepare(context, batch) + backend.update(context) + attention = LayerAttentionBinding( + group_id=0, + logical_layer_id=0, + spec=spec, + params=DenseParams(key, value), + backend=backend, + alloc_tensor=lambda shape, dtype, **kwargs: torch.empty(shape, dtype=dtype, device=device), + ) + actual = attention.forward(context, DenseInputs(query), out=torch.empty_like(query)) + expected, start = [], 0 + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + q = query[start : start + q_length].float() + k = key[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + v = value[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + logits = torch.einsum("qhd,khd->hqk", q, k) * dim**-0.5 + if spec.causal: + mask = torch.arange(length, device=device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=device)[:, None] + ) + logits.masked_fill_(~mask, -torch.inf) + expected.append(torch.einsum("hqk,khd->qhd", logits.softmax(-1), v)) + start += q_length + expected = torch.cat(expected).to(spec.dtype) + runtime.synchronize() + if not torch.allclose(actual, expected, atol=0.02, rtol=0.02): + return False, f"dense {spec.phase.value} output mismatch: max diff {(actual - expected).abs().max().item()}" + return True, None diff --git a/lightllm/common/attention/registry.py b/lightllm/common/attention/registry.py new file mode 100644 index 0000000000..f8b4e9a606 --- /dev/null +++ b/lightllm/common/attention/registry.py @@ -0,0 +1,118 @@ +"""Lazy provider declarations and default preference order; importing this does not load kernels.""" + +import importlib + +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.platform import Platform + + +def load_backend_class(reference): + if isinstance(reference, type): + return reference + module, name = reference.split(":") + if not module.startswith("lightllm."): + module = "lightllm.common.attention." + module + return getattr(importlib.import_module(module), name) + + +varlen_attention_ops = { + "fa3": "backends.dense.fa3.varlen:Fa3VarlenAttentionOp", + "xformers": "backends.dense.xformers.varlen:XformersVarlenAttentionOp", + "sdpa": "backends.dense.sdpa.varlen:SdpaVarlenAttentionOp", + "triton": "backends.dense.triton.varlen:TritonVarlenAttentionOp", +} + + +def get_varlen_attention_op_class(name): + if name not in varlen_attention_ops: + raise ValueError(f"unknown varlen attention operator: {name}") + return load_backend_class(varlen_attention_ops[name]) + + +cache_format_to_backend = { + "None": { + "triton": "backends.dense.triton.backend:TritonAttBackend", + "fa3": "backends.dense.fa3.backend:Fa3AttBackend", + "flashinfer": "backends.dense.flashinfer.backend:FlashInferAttBackend", + }, + "int4kv": { + "triton": "backends.dense.triton.int4:Int4kvTritonAttBackend", + }, + "int8kv": { + "triton": "backends.dense.triton.int8:Int8kvTritonAttBackend", + }, + "fp8kv_sph": { + "fa3": "backends.dense.fa3.fp8:Fp8Fa3AttBackend", + }, + "fp8kv_spt": { + "flashinfer": "backends.dense.flashinfer.fp8:Fp8FlashInferAttBackend", + }, +} + +mla_cache_format_to_backend = { + "None": { + "triton": "backends.mla.triton.backend:MlaTritonAttBackend", + "fa3": "backends.mla.fa3.backend:MlaFa3AttBackend", + "flashinfer": "backends.mla.flashinfer.backend:MlaFlashInferAttBackend", + }, +} + +nsa_cache_format_to_backend = { + "None": { + "flashmla_sparse": "backends.sparse_mla.flashmla.backend:NsaFlashMlaSparseAttBackend", + }, + "fp8kv_dsa": { + "flashmla_sparse": "backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", + }, +} + +neo_cache_format_to_backend = { + "None": { + "triton": "backends.dense.triton.neo:NeoTritonAttBackend", + "fa3": "backends.dense.fa3.neo:NeoFa3AttBackend", + } +} + + +def backend_candidates(family, platform, ssm_state_mode): + tables = { + AttentionFamily.DENSE: cache_format_to_backend, + AttentionFamily.MLA: mla_cache_format_to_backend, + AttentionFamily.SPARSE_MLA: nsa_cache_format_to_backend, + AttentionFamily.GDN: { + "None": { + "flashqla": "backends.gdn.flashqla.backend:FlashQlaLinearAttBackend", + "triton": "backends.gdn.triton.backend:TritonLinearAttBackend", + } + }, + AttentionFamily.KDA: {"None": {"triton": "backends.kda.triton.backend:KDALinearAttBackend"}}, + } + if ssm_state_mode == "replay": + tables[AttentionFamily.GDN] = { + "None": { + "flashqla": "backends.gdn.replay.backend:ReplayGDNFlashQlaAttBackend", + "triton": "backends.gdn.replay.backend:ReplayGDNTritonAttBackend", + } + } + tables[AttentionFamily.KDA] = { + "None": { + "triton": "backends.kda.replay.backend:ReplayKDALinearAttBackend", + } + } + if platform is not Platform.CUDA: + tables = {} + return tables.get(family, {}) + + +def backend_priority(family, phase, specs): + if family is AttentionFamily.GDN: + priority = ["flashqla", "triton"] if phase is AttentionPhase.PREFILL else ["triton"] + elif family is AttentionFamily.KDA: + priority = ["triton"] + elif family is AttentionFamily.SPARSE_MLA: + priority = ["flashmla_sparse"] + else: + priority = ["fa3", "flashinfer", "triton"] + if phase is AttentionPhase.DECODE and not any(spec.use_mtp for spec in specs): + priority = ["flashinfer", "fa3", "triton"] + return priority diff --git a/lightllm/common/attention/selection.py b/lightllm/common/attention/selection.py new file mode 100644 index 0000000000..5fe503b2f9 --- /dev/null +++ b/lightllm/common/attention/selection.py @@ -0,0 +1,137 @@ +import torch +import torch.distributed as dist + +from lightllm.common.attention.api.backend import UnsupportedAttention +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.common.attention.registry import ( + backend_candidates, + backend_priority, + cache_format_to_backend, + get_varlen_attention_op_class, + load_backend_class, + mla_cache_format_to_backend, + neo_cache_format_to_backend, + nsa_cache_format_to_backend, + varlen_attention_ops, +) +from lightllm.platform import get_device_runtime +from lightllm.utils import backend_validator +from lightllm.utils.backend_validator import validate +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class BackendSelectionError(RuntimeError): + def __init__(self, family, phase, failures): + self.failures = failures + details = "; ".join(f"{name}: {reason}" for name, reason in failures.items()) + super().__init__(f"No usable {family.value} {phase.value} attention backend: {details}") + + +def select_varlen_attention_op(requested="auto", *, device=None, devices=None, dtype=None, head_dim=80): + if isinstance(dtype, str): + aliases = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} + dtype = getattr(torch, aliases.get(dtype, dtype)) + if devices is None: + # Resolve once before probing; probes receive concrete device ordinals. + devices = (get_device_runtime(device).device,) + else: + devices = tuple(get_device_runtime(d).device for d in devices) + if not devices: + raise ValueError("varlen selection requires at least one device") + ordered = tuple(varlen_attention_ops) if requested == "auto" else (requested,) + failures = {} + for name in ordered: + if name not in varlen_attention_ops: + failures[name] = "unknown varlen attention operator" + continue + # 在子进程中检查实际 Q/K/V 路径，避免父进程初始化 CUDA 或误用文本 attention 探针。 + if not all( + backend_validator._validate(name, varlen=True, device=d, dtype=dtype, head_dim=head_dim) for d in devices + ): + failures[name] = "varlen environment or numerical probe failed on an assigned device" + continue + logger.info("Selected %s varlen attention", name) + return name + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) + + +def select_attention_backend( + *, + family, + phase, + specs, + caps, + execution, + requested="auto", + candidates=None, + priority=None, + encoding="None", + ssm_state_mode="native", +): + specs = tuple(specs) + if candidates is None: + candidates = backend_candidates(family, caps.platform, ssm_state_mode) + encoding = "None" if family in (AttentionFamily.GDN, AttentionFamily.KDA) else encoding + available = candidates.get(encoding, {}) + if priority is None: + priority = backend_priority(family, phase, specs) + ordered = [requested] if requested != "auto" else list(priority) + if not specs: + raise ValueError(f"model has no {family.value} attention layers") + failures = {} + selected_backend = None + selected_name = None + for name in ordered: + if name not in available: + failures[name] = f"no implementation for cache encoding {encoding}" + continue + try: + backend = load_backend_class(available[name]) + backend.check_environment(caps) + for spec in specs: + backend.check_support(spec, caps, execution) + for spec in specs: + if not validate(available[name], device=str(caps.device), dtype=spec.dtype, spec=spec): + failures[name] = "provider numerical probe failed on the assigned device/dtype" + break + if name in failures: + continue + except UnsupportedAttention as exc: + failures[name] = str(exc) + continue + except Exception as exc: + # All ranks must reach consensus, including a rank whose provider check failed. + logger.exception( + "Unexpected %s attention provider failure: family=%s phase=%s device=%s specs=%r", + name, + family.value, + phase.value, + caps.device, + specs, + ) + failures[name] = f"provider check failed: {type(exc).__name__}: {exc}" + continue + logger.info("Selected %s %s attention on %s: %s", family.value, phase.value, caps.device, name) + selected_backend = backend + selected_name = name + break + + error = None if selected_backend is not None else BackendSelectionError(family, phase, failures) + _check_rank_selection(family, phase, selected_name, None if error is None else str(error)) + if error is not None: + raise error + return selected_backend + + +def _check_rank_selection(family, phase, selected_name, failure): + if dist.is_initialized() and dist.get_world_size() > 1: + selections = [None] * dist.get_world_size() + dist.all_gather_object(selections, {"backend": selected_name, "failure": failure}) + selected = {item["backend"] for item in selections if item["backend"] is not None} + rejected = {f"rank {rank}": item["failure"] for rank, item in enumerate(selections) if item["failure"]} + if len(selected) > 1: + rejected["backend selection"] = ", ".join(sorted(selected)) + if rejected: + raise BackendSelectionError(family, phase, rejected) diff --git a/lightllm/common/attention/varlen.py b/lightllm/common/attention/varlen.py new file mode 100644 index 0000000000..aa758efdea --- /dev/null +++ b/lightllm/common/attention/varlen.py @@ -0,0 +1,35 @@ +from lightllm.platform import get_device_runtime +from lightllm.utils.log_utils import init_logger + +from .api import AttentionFamily, AttentionPhase, UnsupportedAttention +from .selection import BackendSelectionError, get_varlen_attention_op_class + +logger = init_logger(__name__) + + +class AutoVarlenAttentionOp: + """音频编码器的自动路径；失败的实现只在当前调用方禁用，并在本次调用中重算输出。""" + + def __init__(self, providers=("fa3", "triton")): + self.providers = providers + self.ops = {} + self.failures = {} + + def forward(self, inputs, *, out): + for name in self.providers: + if name in self.failures: + continue + try: + if name not in self.ops: + op_cls = get_varlen_attention_op_class(name) + caps = get_device_runtime(inputs.query.device).caps + op_cls.check_environment(caps) + op_cls.check_support(inputs, caps) + self.ops[name] = op_cls() + return self.ops[name].forward(inputs, out=out) + except UnsupportedAttention as exc: + self.failures[name] = str(exc) + except Exception as exc: + self.failures[name] = str(exc) + logger.warning("Disabling %s varlen attention for this encoder: %s", name, exc) + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, self.failures) diff --git a/lightllm/common/basemodel/graph.py b/lightllm/common/basemodel/graph.py new file mode 100644 index 0000000000..dd03b44242 --- /dev/null +++ b/lightllm/common/basemodel/graph.py @@ -0,0 +1,72 @@ +"""持有一次捕获的输入、provider 任务和完成事件；bucket 策略仍由模型决定。""" + +from lightllm.platform import get_device_runtime + + +class GraphExecution: + def __init__(self, graph, states=(), outputs=(), *, runtime=None): + self.graph = graph + self.states = tuple(states) + self.outputs = tuple(outputs) + self.runtime = runtime if runtime is not None else get_device_runtime(self.states[0].input_ids.device) + self.bindings = [] + self.workspaces = [] + self.generations = [] + try: + for state in self.states: + manager = state.att_state_manager + self.generations.append((manager, manager.storage_generation)) + for group_id, context in state.attention_contexts.items(): + backend = state.attention_groups[group_id].backend_for(context.capacity.phase) + binding = backend.bind_graph(context) + if binding is not None: + self.bindings.append(binding) + if context.workspace not in self.workspaces: + self.workspaces.append(context.workspace) + manager.register_graph(self) + for workspace in self.workspaces: + workspace.freeze() + except BaseException: + self.close() + raise + + def replay(self, *states): + if self.graph is None: + raise RuntimeError("Graph execution is closed") + if states and len(states) != len(self.states): + raise ValueError("Graph execution lane count changed") + for manager, generation in self.generations: + if manager.storage_generation != generation: + raise RuntimeError("attention storage changed; rebind layers and recapture Graph") + for binding in self.bindings: + binding.wait_for_update() + if not states: + for workspace in self.workspaces: + workspace.begin() + for target, source in zip(self.states, states): + target.copy_for_cuda_graph(source) + for binding in self.bindings: + binding.before_replay() + self.graph.replay() + # 必须先提交辅助更新，再记录完成事件：Graph 本身可能正在等待更新流。 + events = [] + for binding in self.bindings: + events.extend(binding.after_submit()) + for workspace in self.workspaces: + workspace.finish(dependencies=events) + return self.outputs[0] if len(self.outputs) == 1 else self.outputs + + def close(self): + if self.graph is None: + return + for binding in self.bindings: + binding.wait_for_update() + self.runtime.synchronize() + for binding in self.bindings: + binding.close() + self.graph.reset() + self.graph = None + self.bindings.clear() + self.workspaces.clear() + self.states = self.outputs = () + self.generations.clear() diff --git a/lightllm/models/deepseek_v4/attention/__init__.py b/lightllm/models/deepseek_v4/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/deepseek_v4/attention/sparse_mla.py b/lightllm/models/deepseek_v4/attention/sparse_mla.py new file mode 100644 index 0000000000..daca45985c --- /dev/null +++ b/lightllm/models/deepseek_v4/attention/sparse_mla.py @@ -0,0 +1,208 @@ +import inspect +from dataclasses import dataclass + +import torch +import triton + +from lightllm.common.attention.api import AttentionFamily, UnsupportedAttention +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend + + +@dataclass(frozen=True) +class DeepseekV4Params: + swa: torch.Tensor + prefill_swa: torch.Tensor + compressed: torch.Tensor | None + compress_ratio: int + head_dim_v: int + softmax_scale: float + attn_sink: torch.Tensor + + +@dataclass(frozen=True) +class DeepseekV4Inputs: + query: torch.Tensor + swa_indices: torch.Tensor + swa_lengths: torch.Tensor + extra_indices: torch.Tensor | None = None + extra_lengths: torch.Tensor | None = None + + +@dataclass +class SwaPrefillWorkspace: + swa_buffer: torch.Tensor + HOLD_SWA_SLOT: int + prefill_token_ranges: torch.Tensor + swa_prefill_workspace_offsets: torch.Tensor + + +def prepare_swa_workspace(swa_state, buffer, capacity, batch, slot): + from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace + + metadata = SwaPrefillWorkspace( + buffer, + buffer.shape[1] - swa_state.config.page_token_size, + torch.zeros((swa_state.request_capacity, 2), dtype=torch.int32, device=buffer.device), + torch.full((swa_state.request_capacity,), -1, dtype=torch.int32, device=buffer.device), + ) + prepare_swa_prefill_workspace( + swa_state=swa_state, + prefill_state=metadata, + b_req_idx=batch.b_req_idx, + b_ready_cache_len=batch.b_ready_cache_len, + b_seq_len=batch.b_seq_len, + workspace_base=slot * capacity, + ) + return metadata + + +@dataclass(slots=True) +class DeepseekV4Context(AttentionExecutionContext): + sched_meta_by_ratio: dict | None = None + capture_sched_meta_by_ratio: dict | None = None + + +class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttentionBackend): + supports_attention_sink = True + supports_noncausal_decode = True + family = AttentionFamily.SPARSE_MLA + provider = "flashmla_sparse" + dependency = "flash_mla" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + cache_formats = frozenset({"fp8kv_dsa"}) + + @classmethod + def check_environment(cls, caps) -> None: + super().check_environment(caps) + import flash_mla + + parameters = inspect.signature(flash_mla.flash_mla_with_kvcache).parameters + required = {"out", "attn_sink", "indices", "topk_length", "extra_k_cache", "extra_topk_length"} + missing = required.difference(parameters) + if missing: + raise UnsupportedAttention( + "DSV4 requires the FlashMLA version pinned in docker/Dockerfile; missing: " + ", ".join(sorted(missing)) + ) + + def __init__(self, model): + super().__init__(model=model) + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ) + self.tp_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_ + # Sparse FlashMLA uses 64-head tiles; TP can leave fewer local Q heads. + self.padded_q_head_num = ((self.tp_q_head_num + 63) // 64) * 64 + self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) + self.layer_compress_ratios = dict(zip(config.logical_layer_ids, config.compress_ratios)) + self.swa_config = config.create_swa_config() + self.compressed_configs = { + ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio + } + + @classmethod + def get_swa_prefill_workspace(cls, model): + """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" + if hasattr(model, "swa_prefill_workspace"): + return model.swa_prefill_workspace + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ).create_swa_config() + page_token_num = config.page_token_size + max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ + capacity = ( + triton.cdiv( + max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num + ) + * page_token_num + ) + if model.args.run_mode == "decode": + # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， + # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 + capacity = 0 + regions = 1 + int(model.args.enable_prefill_microbatch_overlap) + # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 + token_num = regions * capacity + page_token_num + # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 + slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num + layer_count = len(config.logical_layer_ids) + buffer = torch.empty(layer_count * token_num * slot_bytes, dtype=torch.uint8, device="cuda") + model.swa_prefill_workspace = buffer.view(layer_count, token_num, slot_bytes), capacity + return model.swa_prefill_workspace + + def _allocate_prefill_context(self, capacity, capture, workspace): + return DeepseekV4Context(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return DeepseekV4Context(capacity, capture, workspace) + + def _prepare_prefill(self, context): + self._prepare_decode(context) + + def _prepare_decode(self, context): + import flash_mla as flashmla + + if context.sched_meta_by_ratio is None: + context.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios} + context.capture_sched_meta_by_ratio = { + ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios + } + + def _run(self, layer, context, inputs, swa, out): + import flash_mla as flashmla + + params = layer.params + compressed = None + if params.compressed is not None: + compressed = self.compressed_configs[params.compress_ratio].get_attention_view(packed_kv=params.compressed) + scheduler = ( + context.capture_sched_meta_by_ratio + if torch.cuda.is_current_stream_capturing() + else context.sched_meta_by_ratio + ) + q = inputs.query + if out is None: + out = layer.alloc_tensor((q.shape[0], q.shape[1], params.head_dim_v), dtype=q.dtype, device=q.device) + padded_out, _ = flashmla.flash_mla_with_kvcache( + q=q.unsqueeze(1), + k_cache=self.swa_config.get_attention_view(packed_kv=swa), + block_table=None, + cache_seqlens=None, + head_dim_v=params.head_dim_v, + tile_scheduler_metadata=scheduler[params.compress_ratio], + num_splits=None, + softmax_scale=params.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=inputs.swa_indices, + attn_sink=params.attn_sink, + topk_length=inputs.swa_lengths, + extra_k_cache=compressed, + extra_indices_in_kvcache=inputs.extra_indices, + extra_topk_length=inputs.extra_lengths, + out=out.unsqueeze(1), + ) + return padded_out[:, 0, : self.tp_q_head_num] + + def _forward_prefill(self, layer, context, inputs: DeepseekV4Inputs, *, out=None): + return self._run(layer, context, inputs, layer.params.prefill_swa, out) + + def _forward_decode(self, layer, context, inputs: DeepseekV4Inputs, *, out=None): + return self._run(layer, context, inputs, layer.params.swa, out) + + +DSV4_NSA_BACKENDS = { + "fp8kv_dsa": { + "flashmla_sparse": "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4FlashMlaFp8SparseAttBackend", + } +} diff --git a/lightllm/models/glm5_next/attention/__init__.py b/lightllm/models/glm5_next/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py new file mode 100644 index 0000000000..f91330042c --- /dev/null +++ b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py @@ -0,0 +1,224 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.execution import BackendResources, WorkspaceRequirement +from lightllm.common.attention.api.inputs import SparseMLAPrefillInputs +from lightllm.common.attention.backends.sparse_mla.flashmla.fp8 import ( + NsaFlashMlaFp8SparseAttBackend, + NsaFlashMlaFp8SparseDecodeContext, +) + +from .sparse_mla import Glm5NextSparsePrefillContext + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparsePrefillContext(Glm5NextSparsePrefillContext): + smallq_workspace: object = None + generic_workspace: object = None + generic_binding: object = None + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparseDecodeContext(NsaFlashMlaFp8SparseDecodeContext): + decode_topk_bound: int = None + + +class Fp8Glm5NextSparseAttBackend(NsaFlashMlaFp8SparseAttBackend): + provider = "flashmla_sparse" + cache_formats = frozenset({"fp8kv_dsa"}) + + def __init__(self, model): + super().__init__(model) + self.prefill_graph_enabled = model.args.enable_prefill_cudagraph + self.prefill_token_capacity = model.batch_max_tokens or 0 + self.decode_only = model.args.run_mode == "decode" + + def resource_requirements(self, capacity, *, capture): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + if capacity.phase is not AttentionPhase.PREFILL: + return super().resource_requirements(capacity, capture=capture) + tokens = max(1, self.prefill_token_capacity, capacity.token_num) + kv_tokens = tokens if self.decode_only else capacity.batch_size * capacity.max_kv_seq_len + tokens + gather_tokens = min(KV_GATHER_CHUNK_SIZE, kv_tokens) + topk = (self.network_config["index_topk"] + 3 + 127) // 128 * 128 + return BackendResources( + workspace=( + WorkspaceRequirement("glm_gathered_kv", gather_tokens * 512, torch.bfloat16), + WorkspaceRequirement("glm_chunk_indices", tokens * topk, torch.int32), + WorkspaceRequirement("glm_chunk_lengths", tokens, torch.int32), + WorkspaceRequirement("glm_cu_q", tokens + 1, torch.int32), + ) + ) + + def get_smallq_workspace(self, context, q): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + allocate_direct_smallq_workspace, + allocate_selected_prefill_workspace, + ) + + context.smallq_workspace = context.workspace.get_object( + "glm_smallq", + lambda: { + "direct": allocate_direct_smallq_workspace(q[:1]), + "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), + }, + ) + if q.shape[0] == 1: + return context.smallq_workspace["direct"] + batch = q.shape[0] + selected = context.smallq_workspace["selected"] + slots = selected["compact"].shape[1] + return { + "compact": selected["compact"][:batch], + "lengths": selected["lengths"][:batch], + "page_table": selected["page_table"][:batch], + "cu_q": selected["cu_q"][: batch + 1], + "cu_k": selected["cu_k"][: batch + 1], + "selected_kv": selected["selected_kv"][: batch * slots], + } + + def get_generic_prefill_workspace(self, context, q, packed_kv, topk_indices): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + rows, columns = topk_indices.shape + workspace = context.workspace + tokens = max(self.prefill_token_capacity, rows) + binding = (rows, columns, topk_indices.dtype, packed_kv.shape[0]) + if context.generic_workspace is None or context.generic_binding != binding: + row_table = workspace.get("glm_prefill_rows", packed_kv.shape[0], torch.int32)[: packed_kv.shape[0]] + row_table.fill_(-1) + cu_q = workspace.get("glm_cu_q", tokens + 1, torch.int32)[: rows + 1] + torch.arange(rows + 1, device=q.device, dtype=torch.int32, out=cu_q) + context.generic_workspace = { + "prefill_row_table": row_table, + "chunk_indices": workspace.get("glm_chunk_indices", tokens * columns, topk_indices.dtype)[ + : rows * columns + ].view(rows, columns), + "chunk_lengths": workspace.get("glm_chunk_lengths", tokens, torch.int32)[:rows], + "cu_seqlens_q": cu_q, + "row_table_built": False, + } + context.generic_binding = binding + gather_tokens = min(KV_GATHER_CHUNK_SIZE, context.ragged_mem_index.numel()) + if self.decode_only: + gather_tokens = min(gather_tokens, max(1, self.prefill_token_capacity, rows)) + context.generic_workspace["gathered_kv"] = workspace.get( + "glm_gathered_kv", max(1, gather_tokens) * 512, q.dtype + )[: gather_tokens * 512].view(gather_tokens, 1, 512) + context.generic_workspace["prefill_row_table"].fill_(-1) + context.generic_workspace["row_table_built"] = False + return context.generic_workspace + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8Glm5NextSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8Glm5NextSparseDecodeContext( + capacity, + capture, + workspace, + decode_topk_bound=2051 + if self.network_config["index_topk"] == 2048 and self.network_config["index_kpool"] == 4 + else None, + ) + + def _prefill_nsa_prefill_att(self, layer, context, q, kv, inputs): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8 + + if q.shape[1] != 64: + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + layer.params.softmax_scale, + generic_workspace=self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices), + ) + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + can_use_direct_smallq, + can_use_selected_prefill, + ) + + direct_workspace = ( + self.get_smallq_workspace(context, q) + if can_use_direct_smallq( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + else None + ) + selected_eligible = direct_workspace is None and can_use_selected_prefill( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + workspace = self.get_smallq_workspace(context, q) if selected_eligible else None + generic_workspace = None + if direct_workspace is None and (not selected_eligible): + generic_workspace = self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices) + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + layer.params.softmax_scale, + workspace, + direct_workspace, + generic_workspace, + ) + + def _decode_nsa_decode_att(self, layer, context, q, packed_kv, inputs): + import flash_mla + + from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode + from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads + + topk_mem_indices = inputs.topk_mem_indices + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, _ = q + num_heads = q_nope.shape[1] + if fp8_sparse_decode.can_use_fp8_sparse_direct_decode( + q_nope, + packed_kv, + topk_mem_indices, + layer.params.kv_lora_rank, + context.decode_topk_bound, + is_cuda_graph_capture=context.capture, + ): + return fp8_sparse_decode.fp8_sparse_direct_decode( + q_nope, packed_kv, topk_mem_indices, layer.params.softmax_scale + ) + padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=context.capture) + kv_cache = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + output, _ = flash_mla.flash_mla_with_kvcache( + q=padded_q, + k_cache=kv_cache, + block_table=None, + cache_seqlens=None, + head_dim_v=layer.params.kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=layer.params.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return output[:, 0, :num_heads, :] + + def _forward_prefill(self, layer, context, inputs: SparseMLAPrefillInputs, *, out=None): + return self._prefill_nsa_prefill_att(layer, context, inputs.query, layer.params.kv, inputs) diff --git a/lightllm/models/glm5_next/attention/sparse_mla.py b/lightllm/models/glm5_next/attention/sparse_mla.py new file mode 100644 index 0000000000..40c1b8e871 --- /dev/null +++ b/lightllm/models/glm5_next/attention/sparse_mla.py @@ -0,0 +1,94 @@ +import dataclasses +import inspect + +import torch + +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import ( + NsaFlashMlaSparseAttBackend, + NsaFlashMlaSparseDecodeContext, + NsaFlashMlaSparsePrefillContext, +) + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparsePrefillContext(NsaFlashMlaSparsePrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparseDecodeContext(NsaFlashMlaSparseDecodeContext): + pass + + +class Glm5NextSparseAttBackend(NsaFlashMlaSparseAttBackend): + provider = "flashmla_sparse" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Glm5NextSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Glm5NextSparseDecodeContext(capacity, capture, workspace) + + def _prefill_nsa_prefill_att(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + tokens, heads, dim = q.shape + padded_heads = (heads + 63) // 64 * 64 + padded_q = q + if padded_heads != heads: + padded_q = q.new_zeros((tokens, padded_heads, dim)) + padded_q[:, :heads] = q + out, _, _ = flash_mla_sparse_fwd( + q=padded_q, + kv=kv, + indices=inputs.topk_mem_indices.unsqueeze(1), + sm_scale=layer.params.softmax_scale, + d_v=layer.params.kv_lora_rank, + ) + return out[:, :heads] + + def _update_decode(self, context): + super()._update_decode(context) + pool = self.network_config["index_kpool"] + topk = self.network_config["index_topk"] + context.nsa_cache_seqlens = ( + torch.minimum(context.lengths // pool * pool, torch.full_like(context.lengths, topk)) + + context.lengths % pool + ) + context.nsa_cu_seqlens_k_new = torch.nn.functional.pad( + context.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0) + ) + + def _decode_nsa_decode_att(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + q_nope, _ = q + kv_nope = kv.view(-1, 1, 1, layer.params.kv_lora_rank) + common_kwargs = dict( + page_table=inputs.topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=layer.params.softmax_scale, + causal=False, + ) + if "only_qv" in inspect.signature(flash_attn_with_kvcache).parameters: + return flash_attn_with_kvcache( + q=None, qv=q_nope, k_cache=None, v_cache=kv_nope, only_qv=True, **common_kwargs + ) + if q_nope.shape[-1] > 256: + page_indices = inputs.topk_mem_indices.clamp_min(0) + key = kv_nope[page_indices, 0, 0] + score = torch.einsum("bhd,bkd->bhk", q_nope.float(), key.float()) + score = score * layer.params.softmax_scale + valid = ( + torch.arange(page_indices.shape[1], device=page_indices.device)[None, :] + < context.nsa_cache_seqlens[:, None] + ) + score = score.masked_fill(~valid[:, None, :], float("-inf")) + weights = score.softmax(dim=-1) + weights = torch.where(valid[:, None, :], weights, torch.zeros_like(weights)) + return torch.einsum("bhk,bkd->bhd", weights, key.float()).to(q_nope.dtype) + return flash_attn_with_kvcache(q=q_nope, k_cache=kv_nope, v_cache=kv_nope, **common_kwargs) diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 3458f10131..90c1872263 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -2,10 +2,14 @@ import multiprocessing as mp import os +import time +import traceback +from functools import lru_cache + import torch -from lightllm.utils.log_utils import init_logger + from lightllm.utils.dist_utils import get_global_rank -from functools import lru_cache +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -18,20 +22,19 @@ def _compute_ground_truth(q, k, v, is_causal=True): return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal) -def _validate_fa3(): +def _validate_fa3(dtype=torch.bfloat16, spec=None): """Validate FA3 with ground truth.""" - from lightllm.utils.device_utils import is_hopper from lightllm.utils.sgl_utils import flash_attn_varlen_func - if not is_hopper(): - return False, "Not a Hopper GPU" + if torch.cuda.get_device_capability()[0] < 9: + return False, "FA3 requires SM90 or newer" if flash_attn_varlen_func is None: return False, "flash_attn_varlen_func is None" - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -48,18 +51,21 @@ def _validate_fa3(): cu_seqlens_k=cu_seqlens, max_seqlen_q=seq, max_seqlen_k=seq, - softmax_scale=1.0 / (dim ** 0.5), + softmax_scale=1.0 / (dim**0.5), causal=True, ) out = out.reshape(batch, seq, heads, dim).transpose(1, 2) torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashinfer(): +def _validate_flashinfer(dtype=torch.float16, spec=None): """Validate FlashInfer with ground truth.""" capability = torch.cuda.get_device_capability() arch = f"{capability[0]}.{capability[1]}" @@ -67,10 +73,10 @@ def _validate_flashinfer(): import flashinfer - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -88,15 +94,18 @@ def _validate_flashinfer(): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashqla(): +def _validate_flashqla(dtype=None, spec=None): """Validate FlashQLA against LightLLM's vendored FLA kernel.""" from flash_qla import chunk_gated_delta_rule as flashqla_chunk_gated_delta_rule - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig @@ -107,7 +116,11 @@ def _validate_flashqla(): num_v_heads = max(1, linear_config.global_linear_v_heads // linear_config.tp_world_size) head_k_dim = linear_config.head_linear_k_dim head_v_dim = linear_config.head_linear_v_dim - qkv_dtype = linear_config.conv_state_dtype + qkv_dtype = dtype or linear_config.conv_state_dtype + if spec is not None: + num_k_heads = spec.num_q_heads + num_v_heads = max(num_k_heads, num_v_heads) + head_k_dim, head_v_dim = spec.head_size, spec.value_head_size state_dtype = linear_config.ssm_state_dtype batch, seq = 1, 64 @@ -140,7 +153,10 @@ def _validate_flashqla(): ("final state", final_state, expected_state), ): if not torch.allclose(actual, expected, rtol=1e-2, atol=1e-2): - return False, f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}" + return ( + False, + f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}", + ) return True, None @@ -169,7 +185,10 @@ def _softmax_kernel(input_ptr, output_ptr, n_cols, BLOCK: tl.constexpr): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-3, atol=1e-3): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -208,7 +227,10 @@ def _validate_xformers(): out = out.transpose(1, 2).contiguous() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -234,12 +256,15 @@ def _validate_sdpa(): out = scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashmla_sparse(): +def _validate_flashmla_sparse(dtype=torch.bfloat16, spec=None): """Validate flashmla_sparse (NSA) with ground truth.""" try: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -252,7 +277,6 @@ def _validate_flashmla_sparse(): kv_lora_rank = 512 qk_rope_head_dim = 64 qk_dim = kv_lora_rank + qk_rope_head_dim - dtype = torch.bfloat16 device = "cuda" q = torch.randn(batch * seq, heads, qk_dim, dtype=dtype, device=device) @@ -265,7 +289,7 @@ def _validate_flashmla_sparse(): topk_indices = topk_indices.view(batch * seq, 1, index_topk) - softmax_scale = 1.0 / (qk_dim ** 0.5) + softmax_scale = 1.0 / (qk_dim**0.5) try: mla_out, _, _ = flash_mla_sparse_fwd( @@ -282,7 +306,40 @@ def _validate_flashmla_sparse(): return True, None -def _run_in_subprocess(backend_name, pipe): +def _validate_varlen(backend_name, dtype=None, device=None, head_dim=80): + from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs + from lightllm.common.attention.selection import get_varlen_attention_op_class + from lightllm.platform import get_device_runtime + + runtime = get_device_runtime(device or ("cuda" if torch.cuda.is_available() else "cpu")) + op_cls = get_varlen_attention_op_class(backend_name) + try: + op_cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + device = runtime.device + dtype = dtype or (torch.float32 if device.type == "cpu" else torch.bfloat16) + q, k, v = [torch.randn(8, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device=device) + out = torch.empty_like(q) + inputs = VarlenAttentionInputs(q, k, v, cu, 5) + try: + op_cls.check_support(inputs, runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + op_cls().forward(inputs, out=out) + expected = torch.empty_like(q) + for start, end in ((0, 3), (3, 8)): + result = _compute_ground_truth( + q[start:end].transpose(0, 1), k[start:end].transpose(0, 1), v[start:end].transpose(0, 1), is_causal=False + ) + expected[start:end].copy_(result.transpose(0, 1)) + if not torch.allclose(out, expected, rtol=2e-2, atol=2e-2): + return False, f"varlen output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return True, None + + +def _run_in_subprocess(backend_name, pipe, varlen=False, device=None, dtype=None, spec=None, head_dim=80): """Run validation in subprocess with suppressed output.""" import sys @@ -292,50 +349,81 @@ def _run_in_subprocess(backend_name, pipe): sys.stderr = devnull try: - if backend_name == "fa3": - success, err = _validate_fa3() + if device is not None or spec is not None: + from lightllm.platform import get_device_runtime + + runtime = get_device_runtime(device) + if runtime.api is not None: + runtime.api.set_device(runtime.device) + if spec is not None: + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.attention.selection import load_backend_class + + cls = load_backend_class(backend_name) + if ( + spec.family is AttentionFamily.DENSE + and spec.cache.format == "None" + and not spec.use_mtp + and not spec.use_alibi + and not spec.use_sink + and spec.sliding_window is None + and not spec.variant + ): + from lightllm.common.attention.probe import probe_dense + + success, err = probe_dense(cls, spec, runtime) + pipe.send((success, err)) + return + backend_name = cls.provider + if varlen: + success, err = _validate_varlen(backend_name, dtype, device, head_dim) + elif backend_name == "fa3": + success, err = _validate_fa3(dtype or torch.bfloat16, spec) elif backend_name == "xformers": success, err = _validate_xformers() elif backend_name == "sdpa": success, err = _validate_sdpa() elif backend_name == "flashinfer": - success, err = _validate_flashinfer() + success, err = _validate_flashinfer(dtype or torch.float16, spec) elif backend_name == "flashqla": - success, err = _validate_flashqla() + success, err = _validate_flashqla(dtype, spec) elif backend_name == "triton": success, err = _validate_triton() elif backend_name == "flashmla_sparse": - success, err = _validate_flashmla_sparse() + success, err = _validate_flashmla_sparse(dtype or torch.bfloat16, spec) else: success, err = False, f"Unknown backend: {backend_name}" pipe.send((success, err)) - except Exception as e: - pipe.send((False, str(e))) + except Exception: + pipe.send((False, traceback.format_exc())) finally: devnull.close() @lru_cache(maxsize=None) -def validate(backend_name: str) -> bool: - if get_global_rank() == 0: - validate_ok = _validate(backend_name) - torch.distributed.broadcast_object_list([validate_ok], src=0) - else: - validate_ok = [None] - torch.distributed.broadcast_object_list(validate_ok, src=0) - validate_ok = validate_ok[0] - return validate_ok - - -def _validate(backend_name: str) -> bool: +def validate(backend_name: str, *, device=None, dtype=None, spec=None) -> bool: + """每个 rank 在自身设备验证；dtype/完整 spec 不同不能复用另一配置的结果。""" + return _validate(backend_name, device=device, dtype=dtype, spec=spec) + + +def _validate(backend_name: str, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80) -> bool: """Validate backend in subprocess with ground truth check.""" try: ctx = mp.get_context("spawn") parent, child = ctx.Pipe(duplex=False) logger.info(f"Validating {backend_name} backend start, please wait ...") - proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child)) + proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child, varlen, device, dtype, spec, head_dim)) proc.start() - proc.join(timeout=_VALIDATION_TIMEOUT) + child.close() + deadline = time.monotonic() + _VALIDATION_TIMEOUT + # Drain the pipe before joining: a full traceback can exceed its buffer. + result = None + if parent.poll(_VALIDATION_TIMEOUT): + try: + result = parent.recv() + except EOFError: + pass + proc.join(timeout=max(0, deadline - time.monotonic())) if proc.is_alive(): proc.kill() @@ -347,14 +435,21 @@ def _validate(backend_name: str) -> bool: logger.warning(f"{backend_name} validation crashed (exit {proc.exitcode})") return False - if parent.poll(): - success, err = parent.recv() + if result is not None: + success, err = result if success: logger.info(f"{backend_name} validated") return True - logger.warning(f"{backend_name} validation failed: {str(err)[0:88]}") + logger.warning( + "%s validation failed on device=%s dtype=%s spec=%r:\n%s", backend_name, device, dtype, spec, err + ) return False - except Exception as e: - logger.warning(f"{backend_name} validation error: {e}") + except Exception: + logger.exception("%s validation error on device=%s dtype=%s spec=%r", backend_name, device, dtype, spec) return False + finally: + if "parent" in locals(): + parent.close() + if "child" in locals(): + child.close() diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py index e52a18cacc..9c31f725ee 100644 --- a/unit_tests/common/attention/test_contracts.py +++ b/unit_tests/common/attention/test_contracts.py @@ -4,11 +4,23 @@ import pytest import torch +from lightllm.common.attention import selection from lightllm.common.attention.api import AttentionFamily, AttentionPhase, UnsupportedAttention from lightllm.common.attention.backends.base import BaseAttentionBackend from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub +def _select(*, model, specs=None, **kwargs): + return selection.select_attention_backend( + specs=specs or (make_spec(kwargs["family"], kwargs["phase"]),), + caps=model.device_runtime.caps, + execution=model.attention_execution, + encoding=model.args.llm_kv_type, + ssm_state_mode=model.args.ssm_state_mode, + **kwargs, + ) + + def test_graph_copy_validates_all_inputs_before_writing(): captured = make_batch() incoming = make_batch((9, 7)) @@ -55,6 +67,174 @@ def test_layer_attention_rejects_wrong_phase(): attention.forward(context, inputs=None) +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_selection_never_falls_back_to_a_rejected_backend(monkeypatch, requested): + spec = make_spec(cache=replace(make_spec().cache, format="fp8kv_sph")) + monkeypatch.setattr( + BaseAttentionBackend, + "check_environment", + classmethod(lambda cls, caps: None), + ) + with pytest.raises(selection.BackendSelectionError, match="encoding fp8kv_sph"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_environment_checks_apply_to_explicit_and_auto(monkeypatch, requested): + def reject(cls, caps): + raise UnsupportedAttention("missing extension") + + monkeypatch.setattr(BaseAttentionBackend, "check_environment", classmethod(reject)) + with pytest.raises(selection.BackendSelectionError, match="missing extension"): + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + + +def _backend_check_exception_candidates(monkeypatch, error_type=RuntimeError): + class BrokenBackend(BaseAttentionBackend): + provider = "broken" + + @classmethod + def check_environment(cls, caps): + raise error_type("provider ABI mismatch") + + class WorkingBackend(BaseAttentionBackend): + provider = "working" + + @classmethod + def check_environment(cls, caps): + return None + + monkeypatch.setattr("lightllm.common.attention.selection.validate", lambda *args, **kwargs: True) + return {"None": {"broken": BrokenBackend, "working": WorkingBackend}} + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_allows_auto_fallback(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + selected = _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="auto", + candidates=candidates, + priority=["broken", "working"], + ) + assert selected.provider == "working" + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_is_reported_for_explicit_selection(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + with pytest.raises(selection.BackendSelectionError) as error: + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="broken", + candidates=candidates, + priority=["broken", "working"], + ) + + reason = "provider ABI mismatch" + if error_type is not UnsupportedAttention: + reason = f"provider check failed: {error_type.__name__}: {reason}" + assert error.value.failures == {"broken": reason} + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_rejection_reaches_rank_consensus(monkeypatch, requested): + spec = make_spec(dtype=torch.float64) + ranks = [] + monkeypatch.setattr(BaseAttentionBackend, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr(selection, "validate", lambda *args, **kwargs: pytest.fail("probed rejected inputs")) + monkeypatch.setattr(selection, "_check_rank_selection", lambda *args: ranks.append(args)) + with pytest.raises(selection.BackendSelectionError, match="activation dtype"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + assert len(ranks) == 1 + assert ranks[0][:3] == (spec.family, spec.phase, None) + assert "activation dtype" in ranks[0][3] + + +def test_rank_selection_checks_all_ranks_after_local_selection(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + + def gather(output, value): + output[:] = [value, {"backend": "triton", "failure": None}] + + monkeypatch.setattr(dist, "all_gather_object", gather) + with pytest.raises(selection.BackendSelectionError, match="backend selection"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +def test_rank_selection_propagates_remote_failure(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr( + dist, + "all_gather_object", + lambda output, value: output.__setitem__(slice(None), [value, {"backend": None, "failure": "probe failed"}]), + ) + with pytest.raises(selection.BackendSelectionError, match="rank 1: probe failed"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +@pytest.mark.parametrize("format,provider", [("fp8kv_sph", "flashinfer"), ("fp8kv_spt", "fa3")]) +def test_fp8_scale_layouts_are_not_interchangeable(format, provider): + candidates = selection.cache_format_to_backend[format] + assert provider not in candidates + + +def test_layer_shape_configures_group_planner(): + from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferAttBackend + + model = model_stub(page_size=16) + model.config = {"num_attention_heads": 16, "num_key_value_heads": 4, "hidden_size": 2048} + model.max_seq_length = 128 + model.data_type = torch.bfloat16 + small = FlashInferAttBackend.create(model, make_spec(num_q_heads=8, head_size=64)) + large = FlashInferAttBackend.create(model, make_spec(num_q_heads=16, head_size=128)) + assert (small.tp_q_head_num, small.head_dim) == (8, 64) + assert (large.tp_q_head_num, large.head_dim) == (16, 128) + + +def test_replay_flashqla_rejects_decode_and_fp32(): + from lightllm.common.attention.backends.gdn.replay.backend import ReplayGDNFlashQlaAttBackend + + spec = make_spec(family=AttentionFamily.GDN) + with pytest.raises(UnsupportedAttention, match="decode"): + ReplayGDNFlashQlaAttBackend.check_support(spec) + with pytest.raises(UnsupportedAttention, match="dtype"): + ReplayGDNFlashQlaAttBackend.check_support(replace(spec, phase=AttentionPhase.PREFILL, dtype=torch.float32)) + + class _OutputBackend(BaseAttentionBackend): def _forward_decode(self, layer, context, inputs, *, out=None): return inputs diff --git a/unit_tests/common/attention/test_dense_numerics.py b/unit_tests/common/attention/test_dense_numerics.py new file mode 100644 index 0000000000..e84ad83c9e --- /dev/null +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -0,0 +1,54 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionPhase, CacheLayout, ContextCapacity, UnsupportedAttention +from lightllm.common.attention.api.inputs import DenseInputs, DenseParams +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.common.attention.selection import cache_format_to_backend, load_backend_class +from lightllm.common.attention_state import AttentionStateManager +from lightllm.platform import get_device_runtime +from unit_tests.common.attention.helpers import make_batch, make_layer_attention, make_spec, model_stub + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(query, key, value, lengths, q_lengths, window=None): + rows, start = [], 0 + groups = query.shape[1] // key.shape[1] + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + k = key[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + v = value[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + q = query[start : start + q_length].float() + logits = torch.einsum("qhd,khd->hqk", q, k) * query.shape[-1] ** -0.5 + causal = torch.arange(length, device=q.device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] + ) + if window is not None: + causal &= torch.arange(length, device=q.device)[None, :] >= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] - window[0] + ) + probs = logits.masked_fill(~causal, -torch.inf).softmax(-1) + rows.append(torch.einsum("hqk,khd->qhd", probs, v)) + start += q_length + return torch.cat(rows).to(query.dtype) + + +@pytest.mark.parametrize("provider", ["triton", "fa3", "flashinfer"]) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +def test_startup_probe_runs_the_selected_dense_implementation(provider, phase): + from lightllm.common.attention.probe import probe_dense + + cls = load_backend_class(cache_format_to_backend["None"][provider]) + runtime = get_device_runtime() + try: + cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + spec = make_spec( + phase=phase, head_size=64, value_head_size=64, cache=CacheLayout("None", 16, torch.bfloat16, 2, 64, 64) + ) + success, error = probe_dense(cls, spec, runtime) + assert success, error diff --git a/unit_tests/common/attention/test_linear_layout.py b/unit_tests/common/attention/test_linear_layout.py new file mode 100644 index 0000000000..59f2ef8bf6 --- /dev/null +++ b/unit_tests/common/attention/test_linear_layout.py @@ -0,0 +1,72 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import BaseAttentionBackend +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, model_stub + + +class _LinearBackend(BaseAttentionBackend): + def _allocate_decode_context(self, capacity, capture, workspace): + return LinearDecodeContext(capacity, capture, workspace) + + def _update_decode(self, context): + update_linear_decode(context, 2, self.dynamic_verify, replay=self.replay) + + def _forward_decode(self, layer, context, inputs, *, out=None): + return context.accepted_tokens(layer.params) + + +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("replay", [False, True]) +def test_shared_backend_selects_each_states_cursor_through_graph_replay(dynamic, replay): + backend = _LinearBackend(model_stub()) + backend.dynamic_verify, backend.replay = dynamic, replay + requests = [0, 0, 2, 7, 7, 7] if dynamic else [0, 0, 0, 2, 2, 2] + positions = [0, 1, 0, 0, 0, 0] if dynamic else [0, 1, 2, 0, 1, 2] + batch = make_batch( + (5,) * 6, + device="cuda", + b_req_idx=torch.tensor(requests, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor(positions, dtype=torch.int32, device="cuda"), + ) + context = make_context(backend, batch, capture=True, prepare=True) + first = torch.tensor([0, 0, 2, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + second = torch.tensor([2, 0, 1, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + layer_a = make_layer_attention(backend, context, params=first) + layer_b = replace(layer_a, logical_layer_id=1, params=second) + + def step(): + backend.update(context) + a, b = layer_a.forward(context, None), layer_b.forward(context, None) + assert layer_a.forward(context, None) is a + return a, b + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + a, b = step() + graph.replay() + tail = [1] * 4 if dynamic else [] + assert a.tolist() == [1, 3] + tail + assert b.tolist() == [3, 2] + tail + first[0], first[2] = 1, 0 + second[0], second[2] = 0, 2 + graph.replay() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + # A new eager batch must not reuse the previous batch's gathered counts. + backend.prepare(context, batch) + a, b = step() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + expected_requests = torch.tensor([0, 2] + ([7] * 4 if dynamic else []), device="cuda") + torch.testing.assert_close(context.b_conv_buffer_idx.long(), expected_requests) + expected_ssm = expected_requests if replay else expected_requests[:, None] * 3 + torch.arange(3, device="cuda") + torch.testing.assert_close(context.b_ssm_buffer_idx.long(), expected_ssm) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 344bfea57f..3078cb4c4e 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -1,22 +1,22 @@ -import weakref import inspect import sys +import weakref from types import ModuleType, SimpleNamespace import pytest import torch -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.linear.gdn as gdn import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils +import lightllm.common.basemodel.attention.linear.gdn as gdn import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops +import lightllm.utils.backend_validator as backend_validator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.api_cli import make_argument_parser -import lightllm.utils.backend_validator as backend_validator @pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) @@ -251,7 +251,9 @@ def kernel(**kwargs): flashqla = ModuleType("flash_qla") flashqla.chunk_gated_delta_rule = kernel monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr(fla_ops, "chunk_gated_delta_rule", kernel) + from lightllm.common.attention.backends.gdn.triton.kernels import chunk + + monkeypatch.setattr(chunk, "chunk_gated_delta_rule", kernel) original_randn = torch.randn original_rand = torch.rand