[Executor] Move forward_meta.py to fastdeploy/model_executor (#2774)

* Use PEP 563 in attention.py and fix conflict

* merge commit

* Change what was left out last time
This commit is contained in:
littledgg
2025-07-10 20:36:51 +08:00
committed by GitHub
parent 8c660a0dfb
commit 59071268b6
27 changed files with 53 additions and 55 deletions

View File

@@ -0,0 +1,153 @@
"""
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
import logging
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import TYPE_CHECKING, Optional
from fastdeploy.model_executor.layers.attention import AttentionBackend
import paddle
logger = logging.getLogger(__name__)
class ForwardMode(IntEnum):
"""
Forward mode used during attention.
"""
# Prefill and Extend mode
EXTEND = auto()
# Decode mode
DECODE = auto()
# Mixed mode
MIXED = auto()
def is_prefill(self):
""" Is Extend mode """
return self == ForwardMode.EXTEND
def is_decode(self):
""" Is Decode mode """
return self == ForwardMode.DECODE
def is_mixed(self):
""" Is Mixed mode """
return self == ForwardMode.MIXED
@dataclass
class ForwardMeta():
"""
ForwardMeta is used to store the global meta information of the model forward.
"""
# Input tokens IDs
input_ids: paddle.Tensor
# Input tokens IDs of removed padding
ids_remove_padding: paddle.Tensor
# Rotation position embedding
rotary_embs: Optional[paddle.Tensor] = None
# Use cuda graph in this step or not. Used to avoid run cuda graph when in dummy run or prefill stage.
step_use_cudagraph: bool = False
# Batch type flag
is_decode_batch: bool = False
# Attention backend object
attn_backend: AttentionBackend = None
# Forward mode used during attention
forward_mode: ForwardMode = ForwardMode.MIXED
# Attention mask
attn_mask: Optional[paddle.Tensor] = None
# Decoder batch id. Used by attention backend.
decoder_batch_ids: Optional[paddle.Tensor] = None
# Tile ID for each batch of the decoder. Used by attention backend.
decoder_tile_ids_per_batch: Optional[paddle.Tensor] = None
# Sequence length of encoder for ever batch
seq_lens_encoder: Optional[paddle.Tensor] = None
# Sequence length of Encoder for ever batch
seq_lens_decoder: Optional[paddle.Tensor] = None
# The sequence length processed in the current step
seq_lens_this_time: Optional[paddle.Tensor] = None
# Accumulated offset
cum_offsets: Optional[paddle.Tensor] = None
# Offset tensor, used to restore the position of ids_remove_madding after padding removal to the original input_ids
padding_offset: Optional[paddle.Tensor] = None
# Accumulated sequence length of query
cu_seqlens_q: Optional[paddle.Tensor] = None
# Accumulated sequence length of key
cu_seqlens_k: Optional[paddle.Tensor] = None
# Pre-cache length
pre_caches_length: int = 0
# Block tables
block_tables: Optional[paddle.Tensor] = None
# KV caches
caches: Optional[list[paddle.Tensor]] = None
def clear_caches(self):
""" Safely clean up the caches """
if self.caches:
del self.caches
@dataclass
class XPUForwardMeta(ForwardMeta):
"""
XPUForwardMeta is used to store the global meta information of the forward, and some XPU specific meta info.
"""
# TODO(wanghaitao): Supplementary notes
#
encoder_batch_map: Optional[paddle.Tensor] = None
#
decoder_batch_map: Optional[paddle.Tensor] = None
#
encoder_batch_idx: Optional[paddle.Tensor] = None
#
decoder_batch_idx: Optional[paddle.Tensor] = None
#
encoder_seq_lod: Optional[paddle.Tensor] = None
#
decoder_context_len: Optional[paddle.Tensor] = None
#
decoder_context_len_cache: Optional[paddle.Tensor] = None
#
encoder_batch_map_cpu: Optional[paddle.Tensor] = None
#
decoder_batch_map_cpu: Optional[paddle.Tensor] = None
#
encoder_batch_idx_cpu: Optional[paddle.Tensor] = None
#
decoder_batch_idx_cpu: Optional[paddle.Tensor] = None
#
encoder_seq_lod_cpu: Optional[paddle.Tensor] = None
#
decoder_context_len_cpu: Optional[paddle.Tensor] = None
#
decoder_context_len_cache_cpu: Optional[paddle.Tensor] = None
#
batch_tensor: Optional[paddle.Tensor] = None
#
enc_batch: Optional[paddle.Tensor] = None
#
dec_batch: Optional[paddle.Tensor] = None
#
total_enc_len: Optional[paddle.Tensor] = None

View File

@@ -27,14 +27,13 @@ from fastdeploy.model_executor.layers.attention.ops import (
init_signal_layerwise, open_shm_and_get_meta_signal)
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.model_executor.forward_meta import ForwardMeta
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.model_executor.layers.attention.utils import init_rank_and_device_id
from fastdeploy.worker.forward_meta import ForwardMeta
@dataclass
@@ -54,7 +53,7 @@ class AppendAttentionMetadata(AttentionMetadata):
decoder_tile_ids_per_batch: paddle.Tensor = None
decoder_num_blocks: paddle.Tensor = None
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
encoder_max_partition_size: int = 32768
max_partition_size: int = 32768
block_tables: Optional[paddle.Tensor] = None

View File

@@ -14,7 +14,9 @@
# limitations under the License.
"""
from typing import Dict, Optional
from __future__ import annotations
from typing import TYPE_CHECKING, Dict, Optional
import numpy as np
import paddle
@@ -24,7 +26,8 @@ from paddleformers.utils.log import logger
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.quantization.quant_base import \
QuantMethodBase
from fastdeploy.worker.forward_meta import ForwardMeta
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Attention(nn.Layer):

View File

@@ -21,10 +21,11 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING
import paddle
from fastdeploy.worker.forward_meta import ForwardMeta
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
@dataclass

View File

@@ -23,13 +23,13 @@ from typing import TYPE_CHECKING, List, Optional
import paddle
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.model_executor.forward_meta import ForwardMeta
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.worker.forward_meta import ForwardMeta
@dataclass
class BlockAttentionMetadata(AttentionMetadata):
@@ -48,7 +48,7 @@ class BlockAttentionMetadata(AttentionMetadata):
decoder_tile_ids_per_batch: paddle.Tensor = None
decoder_num_blocks: paddle.Tensor = None
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
encoder_max_partition_size: int = 32768
max_partition_size: int = 32768
block_tables: Optional[paddle.Tensor] = None

View File

@@ -18,7 +18,7 @@ from __future__ import annotations
import os
from dataclasses import dataclass, field
from typing import List, Optional
from typing import List, Optional, TYPE_CHECKING
import paddle
@@ -35,7 +35,8 @@ from fastdeploy.model_executor.layers.attention.ops import (
get_block_shape_and_split_kv_block, gqa_rope_write_cache,
init_signal_layerwise, open_shm_and_get_meta_signal, pre_cache_len_concat)
from fastdeploy.model_executor.layers.attention.utils import init_rank_and_device_id
from fastdeploy.worker.forward_meta import ForwardMeta
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
@dataclass

View File

@@ -20,7 +20,7 @@ import os
import paddle
from dataclasses import dataclass
from typing import Optional
from typing import Optional, TYPE_CHECKING
from math import sqrt
from paddle.nn.functional.flash_attention import flash_attn_unpadded
@@ -30,7 +30,8 @@ from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.worker.forward_meta import ForwardMeta
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
@dataclass

View File

@@ -35,15 +35,13 @@ if current_platform.is_cuda() and not current_platform.is_dcu():
prefill_mla_write_cache)
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.model_executor.forward_meta import ForwardMeta
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.model_executor.layers.attention.utils import \
init_rank_and_device_id
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.layers.attention.utils import init_rank_and_device_id
def yarn_get_mscale(scale=1, mscale=1):
@@ -71,7 +69,7 @@ class MLAAttentionMetadata(AttentionMetadata):
decoder_tile_ids_per_batch: paddle.Tensor = None
decoder_num_blocks: paddle.Tensor = None
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
encoder_max_partition_size: int = 32768
max_partition_size: int = 32768
block_tables: Optional[paddle.Tensor] = None

View File

@@ -17,12 +17,14 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import paddle
from paddle.nn.functional import scaled_dot_product_attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import \
AttentionBackend
from fastdeploy.worker.forward_meta import ForwardMeta
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta
class PaddleNativeAttnBackend(AttentionBackend):

View File

@@ -26,13 +26,12 @@ from fastdeploy.model_executor.layers.attention.ops import (
init_signal_layerwise, open_shm_and_get_meta_signal)
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.model_executor.forward_meta import ForwardMeta
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.worker.forward_meta import ForwardMeta
@dataclass
@@ -52,7 +51,7 @@ class XPUAttentionMetadata(AttentionMetadata):
decoder_tile_ids_per_batch: paddle.Tensor = None
decoder_num_blocks: paddle.Tensor = None
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
encoder_max_partition_size: int = 32768
max_partition_size: int = 32768
block_tables: Optional[paddle.Tensor] = None

View File

@@ -24,15 +24,12 @@ import paddle
import numpy as np
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.worker.forward_meta import ForwardMeta, ForwardMode
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta, ForwardMode
from fastdeploy.model_executor.ops.gcu import (fused_rotary_embedding,
mem_efficient_attention,
@@ -47,7 +44,7 @@ class GCUFlashAttnMetadata(AttentionMetadata):
"""
forward_mode: ForwardMode = ForwardMode.MIXED
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
seq_lens_encoder: Optional[paddle.Tensor] = None
seq_lens_decoder: Optional[paddle.Tensor] = None

View File

@@ -25,28 +25,26 @@ import paddle
import numpy as np
import math
if TYPE_CHECKING:
from paddle._typing.dtype_like import _DTypeLiteral
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.layers.attention.attention import Attention
from fastdeploy.model_executor.layers.attention.base_attention_backend import (
AttentionBackend, AttentionMetadata)
from fastdeploy.worker.forward_meta import ForwardMeta, ForwardMode
from fastdeploy.model_executor.ops.gcu import (fused_rotary_embedding,
mem_efficient_attention,
flash_attn_var_len)
from paddleformers.utils.log import logger
if TYPE_CHECKING:
from fastdeploy.model_executor.forward_meta import ForwardMeta, ForwardMode
@dataclass
class GCUMemEfficientAttnMetadata(AttentionMetadata):
"""
GCUMemEfficientAttnMetadata
"""
forward_mode: ForwardMode = ForwardMode.MIXED
_dtype: _DTypeLiteral = paddle.bfloat16
_dtype: paddle.dtype = paddle.bfloat16
seq_lens_encoder: Optional[paddle.Tensor] = None
seq_lens_decoder: Optional[paddle.Tensor] = None

View File

@@ -40,7 +40,7 @@ from fastdeploy.model_executor.layers.rotary_embedding import \
DeepseekScalingRotaryEmbedding
from fastdeploy.model_executor.models.model_base import ModelForCasualLM
from fastdeploy.platforms import current_platform
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
if current_platform.is_cuda():
from fastdeploy.model_executor.ops.gpu import \

View File

@@ -41,7 +41,7 @@ from fastdeploy.model_executor.models.tp_utils import TensorSplitMode as tsm
from fastdeploy.model_executor.models.utils import \
LayerIdPlaceholder as layerid
from fastdeploy.model_executor.models.utils import WeightMeta
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Ernie4_5_MLP(nn.Layer):

View File

@@ -30,7 +30,7 @@ from fastdeploy.model_executor.layers.mtp_linear import ParallelEHProjection
from fastdeploy.model_executor.layers.normalization import RMSNorm
from fastdeploy.model_executor.models.ernie4_5_moe import Ernie4_5_DecoderLayer
from fastdeploy.model_executor.models.model_base import ModelForCasualLM
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Ernie4_5_MTPPretrainedModel(PretrainedModel):

View File

@@ -44,7 +44,7 @@ if current_platform.is_cuda() and not current_platform.is_dcu():
text_image_gather_scatter,
text_image_index_out)
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Ernie4_5_VLMLP(Ernie4_5_MLP):

View File

@@ -34,7 +34,7 @@ from fastdeploy.model_executor.layers.linear import (
from fastdeploy.model_executor.layers.lm_head import ParallelLMHead
from fastdeploy.model_executor.layers.normalization import RMSNorm
from fastdeploy.model_executor.models.model_base import ModelForCasualLM
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Qwen2MLP(nn.Layer):

View File

@@ -34,7 +34,7 @@ from fastdeploy.model_executor.layers.lm_head import ParallelLMHead
from fastdeploy.model_executor.layers.normalization import RMSNorm
from fastdeploy.model_executor.models.model_base import ModelForCasualLM
from fastdeploy.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Qwen3MLP(Qwen2MLP):

View File

@@ -35,7 +35,7 @@ from fastdeploy.model_executor.layers.lm_head import ParallelLMHead
from fastdeploy.model_executor.layers.moe.moe import FusedMoE
from fastdeploy.model_executor.layers.normalization import RMSNorm
from fastdeploy.model_executor.models.model_base import ModelForCasualLM
from fastdeploy.worker.forward_meta import ForwardMeta
from fastdeploy.model_executor.forward_meta import ForwardMeta
class Qwen3MLP(nn.Layer):