clean model_executor/layers/attention/append_attn_backend.py

Author: zhoutianzi666

fastdeploy/model_executor/layers/attention/append_attn_backend.py

@@ -54,9 +54,6 @@ class AppendAttentionMetadata(AttentionMetadata):
     _dtype: paddle.dtype = paddle.bfloat16
     encoder_max_partition_size: int = 32768
     max_partition_size: int = 32768
-    block_tables: Optional[paddle.Tensor] = None
-    rotary_embs: Optional[paddle.Tensor] = None
-    attn_mask: Optional[paddle.Tensor] = None
     _fuse_kernel_compute_dtype: str = "bf16"
 
     # pd_disaggregation
@@ -75,31 +72,34 @@ def allocate_launch_related_buffer(
     block_size,
 ):
     # Initialize AttentionBackend buffers
-    group_size = np.ceil(num_heads / kv_num_heads)
+    assert num_heads % kv_num_heads == 0
+    assert max_model_len % block_size == 0
+    assert max_model_len % encoder_block_shape_q == 0
+    group_size = num_heads // kv_num_heads
 
     # NOTE: (changwenbin) When using auto_chunk,
     # decode_max_tile_size must take into account the maximum case, where *1024 can cover 128K.
     decode_max_tile_size = (
-        1024 * max_batch_size * np.ceil((decoder_step_token_num * group_size) / decoder_block_shape_q)
+        1024 * max_batch_size * (int)(np.ceil(decoder_step_token_num * group_size / decoder_block_shape_q))
     )
-    encode_max_tile_size = max_batch_size * np.ceil((max_model_len * group_size) / encoder_block_shape_q)
-    kv_max_tile_size = max_batch_size * np.ceil(max_model_len / block_size)
+    encode_max_tile_size = max_batch_size * (max_model_len * group_size // encoder_block_shape_q)
+    kv_max_tile_size = max_batch_size * (max_model_len // block_size)
     res = {}
-    res["decoder_batch_ids"] = paddle.full([int(decode_max_tile_size)], 0, dtype="int32")
-    res["decoder_tile_ids_per_batch"] = paddle.full([int(decode_max_tile_size)], 0, dtype="int32")
+    res["decoder_batch_ids"] = paddle.full([decode_max_tile_size], 0, dtype="int32")
+    res["decoder_tile_ids_per_batch"] = paddle.full([decode_max_tile_size], 0, dtype="int32")
     res["decoder_num_blocks_cpu"] = paddle.full([1], 0, dtype="int32").pin_memory()
     # NOTE: (changwenbin) MLA kernel only needs decoder_num_blocks_device in place of GPU tensor,
     # adapted to cudagraph.
     res["decoder_num_blocks_device"] = paddle.full([1], 0, dtype="int32")
     res["decoder_chunk_size_device"] = paddle.full([1], 64, dtype="int32")
     res["max_len_tensor_cpu"] = paddle.full([9], 0, dtype="int32").cpu()
 
-    res["encoder_batch_ids"] = paddle.full([int(encode_max_tile_size)], 0, dtype="int32")
-    res["encoder_tile_ids_per_batch"] = paddle.full([int(encode_max_tile_size)], 0, dtype="int32")
+    res["encoder_batch_ids"] = paddle.full([encode_max_tile_size], 0, dtype="int32")
+    res["encoder_tile_ids_per_batch"] = paddle.full([encode_max_tile_size], 0, dtype="int32")
     res["encoder_num_blocks_x_cpu"] = paddle.full([1], 0, dtype="int32").cpu()
 
-    res["kv_batch_ids"] = paddle.full([int(kv_max_tile_size)], 0, dtype="int32")
-    res["kv_tile_ids_per_batch"] = paddle.full([int(kv_max_tile_size)], 0, dtype="int32")
+    res["kv_batch_ids"] = paddle.full([kv_max_tile_size], 0, dtype="int32")
+    res["kv_tile_ids_per_batch"] = paddle.full([kv_max_tile_size], 0, dtype="int32")
     res["kv_num_blocks_x_cpu"] = paddle.full([1], 0, dtype="int32").cpu()
 
     return res
@@ -175,10 +175,6 @@ def init_attention_metadata(self, forward_meta: ForwardMeta):
             metadata._fuse_kernel_compute_dtype = "fp16"
         elif metadata._dtype == "float32":
             metadata._fuse_kernel_compute_dtype = "fp32"
-        metadata.block_tables = forward_meta.block_tables
-        metadata.rotary_embs = forward_meta.rotary_embs
-        metadata.attn_mask = forward_meta.attn_mask
-        metadata.pre_caches_length = forward_meta.pre_caches_length
 
         # pd_disaggregation
         metadata.kv_signal_data_list = [None] * self.num_layers
@@ -330,7 +326,7 @@ def forward_mixed(
                 forward_meta.seq_lens_this_time,
                 forward_meta.batch_id_per_token,
                 forward_meta.cu_seqlens_q,
-                metadata.block_tables,
+                forward_meta.block_tables,
                 forward_meta.encoder_batch_ids,
                 forward_meta.encoder_tile_ids_per_batch,
                 forward_meta.encoder_num_blocks_x_cpu,
@@ -342,8 +338,8 @@ def forward_mixed(
                 forward_meta.decoder_num_blocks_cpu,
                 forward_meta.max_len_tensor_cpu,
                 res,
-                metadata.rotary_embs,
-                metadata.attn_mask,
+                forward_meta.rotary_embs,
+                forward_meta.attn_mask,
                 layer.qkv_bias,
                 layer.qkv_scale,
                 cache_k_scales,
@@ -387,7 +383,7 @@ def forward_mixed(
                 forward_meta.seq_lens_this_time,
                 forward_meta.batch_id_per_token,
                 forward_meta.cu_seqlens_q,
-                metadata.block_tables,
+                forward_meta.block_tables,
                 forward_meta.encoder_batch_ids,
                 forward_meta.encoder_tile_ids_per_batch,
                 forward_meta.encoder_num_blocks_x_cpu,
@@ -398,8 +394,8 @@ def forward_mixed(
                 forward_meta.decoder_tile_ids_per_batch,
                 forward_meta.decoder_num_blocks_cpu,
                 forward_meta.max_len_tensor_cpu,
-                metadata.rotary_embs,
-                metadata.attn_mask,
+                forward_meta.rotary_embs,
+                forward_meta.attn_mask,
                 layer.qkv_bias,
                 layer.qkv_scale,
                 cache_k_scales,