From d063983db8aa3353f4c46c2da4eb453d995ba384 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 29 Oct 2026 11:20:16 +0000 Subject: [PATCH 01/56] refactor(att_resources): unify attention state and cache transfers Introduce full-attention, bounded-attention and attention-tail states with shared resource allocation and staged model initialization. Integrate CPU cache, DP sharing and PD transfers; add page transfer kernels, regression tests and architecture documentation. --- docs/CN/source/models/add_new_model.md | 25 +- docs/EN/source/models/add_new_model.md | 25 +- docs/cache_manager_design.md | 1702 +++++++++++++++++ lightllm/common/att_resources/__init__.py | 5 + .../att_resources/allocators/__init__.py | 4 + .../allocators/full_att_kv_allocator.py | 100 + .../allocators/state_page_allocator.py | 39 + .../att_resources/att_resource_manager.py | 711 +++++++ .../common/att_resources/att_tail/__init__.py | 5 + .../att_resources/att_tail/base_state.py | 67 + .../att_tail/create_kpool_state.py | 18 + .../att_resources/att_tail/kpool_state.py | 194 ++ .../att_resources/att_tail/transfer_utils.py | 50 + .../att_resources/bounded_att/__init__.py | 19 + .../att_resources/bounded_att/base_state.py | 167 ++ .../bounded_att/conv_ssm_types.py | 53 + .../bounded_att/create_gdn_state.py | 19 + .../bounded_att/create_kda_state.py | 19 + .../bounded_att/gdn_pages_pool.py | 45 + .../att_resources/bounded_att/gdn_state.py | 407 ++++ .../bounded_att/kda_pages_pool.py | 49 + .../att_resources/bounded_att/kda_state.py | 407 ++++ .../bounded_att/transfer_utils.py | 73 + .../common/att_resources/bounded_att/utils.py | 114 ++ .../common/att_resources/full_att/__init__.py | 31 + .../att_resources/full_att/base_state.py | 201 ++ .../full_att/create_gqa_state.py | 67 + .../full_att/create_mla_state.py | 55 + .../att_resources/full_att/gqa_state.py | 339 ++++ .../att_resources/full_att/gqa_transfer.py | 238 +++ .../full_att/group_quantized_state.py | 408 ++++ .../full_att/group_quantized_transfer.py | 248 +++ .../att_resources/full_att/mla_state.py | 307 +++ .../att_resources/full_att/mla_transfer.py | 233 +++ .../full_att/packed_mla_state.py | 413 ++++ .../full_att/packed_mla_transfer.py | 257 +++ .../full_att/static_fp8_state.py | 381 ++++ .../full_att/static_fp8_transfer.py | 240 +++ lightllm/common/att_resources/types.py | 216 +++ lightllm/common/att_resources/utils.py | 16 + .../common/basemodel/attention/base_att.py | 16 +- .../common/basemodel/attention/fa3/fp8.py | 34 +- .../basemodel/attention/flashinfer/fp8.py | 20 +- .../common/basemodel/attention/linear/gdn.py | 10 +- .../common/basemodel/attention/linear/kda.py | 10 +- .../attention/nsa/fp8_flashmla_sparse.py | 9 +- lightllm/common/basemodel/basemodel.py | 126 +- lightllm/common/basemodel/basemodel_utils.py | 153 ++ lightllm/common/basemodel/cuda_graph.py | 4 +- lightllm/common/basemodel/infer_struct.py | 5 +- .../transformer_layer_infer_template.py | 8 +- .../common/basemodel/prefill_cuda_graph.py | 4 +- .../triton_kernel/destindex_copy_kv_fp8.py | 23 +- .../destindex_copy_kv_flashmla_fp8.py | 142 +- .../prefill_compact_kv_flashmla_fp8.py | 186 +- .../common/basemodel/triton_kernel/kv_move.py | 8 +- .../linear_att_cpu_cache_copy.py | 531 ----- .../common/kv_cache_mem_manager/__init__.py | 30 - .../common/kv_cache_mem_manager/allocator.py | 103 - .../deepseek2_mem_manager.py | 83 - .../deepseek3_2mem_manager.py | 24 - ...oken_group_quant_deepseek3_2mem_manager.py | 66 - .../fp8_static_per_head_quant_mem_manager.py | 98 - ...fp8_static_per_tensor_quant_mem_manager.py | 70 - .../glm5_next_mem_manager.py | 122 -- .../kv_cache_mem_manager/mem_manager.py | 308 --- .../common/kv_cache_mem_manager/mem_utils.py | 76 - .../kv_cache_mem_manager/operator/__init__.py | 14 - .../kv_cache_mem_manager/operator/base.py | 48 - .../kv_cache_mem_manager/operator/deepseek.py | 207 -- .../operator/fp8_quant.py | 48 - .../operator/linear_att.py | 210 -- .../kv_cache_mem_manager/operator/normal.py | 115 -- .../kv_cache_mem_manager/operator/quant.py | 129 -- .../ppl_int4kv_mem_manager.py | 50 - .../ppl_int8kv_mem_manager.py | 49 - .../qwen3next_mem_manager.py | 464 ----- lightllm/common/req_manager/__init__.py | 8 - lightllm/common/req_manager/base.py | 37 +- lightllm/common/req_manager/glm5_next.py | 35 - lightllm/common/req_manager/hybrid_base.py | 69 - lightllm/common/req_manager/linear_att.py | 119 -- .../common/state_cache_manager/__init__.py | 22 - lightllm/common/state_cache_manager/base.py | 57 - .../common/state_cache_manager/glm5_next.py | 64 - .../common/state_cache_manager/layer_cache.py | 37 - .../common/state_cache_manager/linear_att.py | 198 -- .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/bloom/model.py | 8 + .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/deepseek2/model.py | 20 +- .../triton_kernel/gqa_flash_decoding_fp8.py | 2 +- .../layer_infer/transformer_layer_infer.py | 11 +- lightllm/models/deepseek3_2/model.py | 6 + lightllm/models/deepseek_mtp/model.py | 4 - lightllm/models/gemma3/model.py | 19 +- .../layer_infer/transformer_layer_infer.py | 5 +- lightllm/models/gemma4/model.py | 33 +- lightllm/models/gemma_2b/model.py | 19 +- lightllm/models/glm4_moe_lite_mtp/model.py | 12 +- .../models/glm5_next/att_resource_config.py | 65 + .../models/glm5_next/att_resource_manager.py | 160 ++ lightllm/models/glm5_next/indexer.py | 11 +- .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/glm5_next/model.py | 86 +- lightllm/models/glm5_next_mtp/model.py | 20 +- .../layer_infer/transformer_layer_infer.py | 6 +- .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/llama/model.py | 23 +- lightllm/models/mistral/model.py | 24 +- lightllm/models/mistral_mtp/model.py | 4 - lightllm/models/mixtral/model.py | 21 +- lightllm/models/qwen2/model.py | 17 - lightllm/models/qwen3_5_dflash/model.py | 22 +- lightllm/models/qwen3_5_dspark/model.py | 22 +- lightllm/models/qwen3_5_mtp/model.py | 14 +- lightllm/models/qwen3_dflash/model.py | 7 +- .../layer_infer/transformer_layer_infer.py | 8 +- lightllm/models/qwen3_dflash2/model.py | 21 +- lightllm/models/qwen3_eagle/model.py | 3 - lightllm/models/qwen3_moe_mtp/model.py | 4 - lightllm/models/qwen3next/__init__.py | 10 +- .../models/qwen3next/att_resource_config.py | 157 ++ .../models/qwen3next/att_resource_manager.py | 94 + .../layer_infer/transformer_layer_infer.py | 3 - lightllm/models/qwen3next/model.py | 95 +- lightllm/models/starcoder/model.py | 27 +- lightllm/models/starcoder2/model.py | 21 +- .../cache_placement_controller.py | 2 +- .../multi_level_kv_cache/cpu_cache_client.py | 12 +- .../multi_level_kv_cache/disk_cache_worker.py | 3 +- lightllm/server/pd_io_struct.py | 5 +- .../dynamic_prompt/hybrid_att_radix_cache.py | 42 +- .../router/dynamic_prompt/radix_cache.py | 19 +- lightllm/server/router/manager.py | 10 +- .../server/router/model_infer/infer_batch.py | 105 +- .../model_infer/mode_backend/base_backend.py | 118 +- .../mode_backend/diverse_backend/impl.py | 2 +- .../dp_backend/dp_shared_kv_trans.py | 166 +- .../mode_backend/multi_level_kv_cache.py | 169 +- .../pd/decode_node_impl/decode_impl.py | 2 +- .../decode_node_impl/decode_trans_process.py | 60 +- .../mode_backend/pd/nccl_kv_transporter.py | 19 +- .../mode_backend/pd/nixl_kv_transporter.py | 9 +- .../prefill_trans_process.py | 51 +- .../mode_backend/pd/resource_io.py | 57 + .../mode_backend/pd/trans_process_obj.py | 2 +- .../model_infer/mtp_speculative/utils.py | 12 +- .../model_infer/shared_att_resources_utils.py | 75 + .../server/router/req_queue/base_queue.py | 8 +- lightllm/utils/att_resource_config.py | 215 +++ lightllm/utils/backend_validator.py | 4 +- lightllm/utils/config_utils.py | 49 - lightllm/utils/kv_cache_utils.py | 123 +- test/att_resources/validate_model_cache.py | 322 ++++ .../static_inference/static_benchmark.py | 46 +- test/cpu_cache_kernel/test_speed.py | 370 +--- unit_tests/common/att_resources/conftest.py | 44 + .../common/att_resources/manager_fixtures.py | 61 + .../common/att_resources/model_helpers.py | 152 ++ .../common/att_resources/test_allocators.py | 90 + .../att_resources/test_backend_startup.py | 305 +++ .../att_resources/test_bounded_state.py | 293 +++ .../common/att_resources/test_capacity.py | 692 +++++++ .../att_resources/test_compute_binding.py | 103 + .../att_resources/test_fp8_compute_binding.py | 150 ++ .../att_resources/test_full_att_cpu_kernel.py | 216 +++ .../att_resources/test_full_att_transfer.py | 281 +++ .../common/att_resources/test_full_state.py | 199 ++ .../test_full_storage_initialization.py | 96 + .../att_resources/test_gqa_pd_kernel.py | 223 +++ .../common/att_resources/test_gqa_transfer.py | 81 + .../att_resources/test_manager_helpers.py | 439 +++++ .../att_resources/test_model_calibration.py | 116 ++ .../test_model_initialization.py | 161 ++ .../att_resources/test_model_resources.py | 162 ++ .../test_model_storage_definition.py | 165 ++ .../att_resources/test_mtp_state_commit.py | 109 ++ .../att_resources/test_network_transport.py | 278 +++ .../att_resources/test_no_attention_draft.py | 76 + .../att_resources/test_page_pool_storage.py | 154 ++ .../att_resources/test_paged_storage.py | 22 + .../test_quantized_full_state.py | 234 +++ .../test_quantized_paged_storage.py | 30 + .../att_resources/test_role_transfers.py | 111 ++ .../att_resources/test_runtime_transport.py | 374 ++++ .../att_resources/test_state_creation.py | 132 ++ .../att_resources/test_state_identity.py | 287 +++ .../att_resources/test_state_queries.py | 181 ++ .../common/att_resources/test_tail_state.py | 81 + .../common/att_resources/test_tp_config.py | 165 ++ .../basemodel/attention/linear/test_gdn.py | 50 +- .../common/basemodel/test_model_output.py | 38 +- .../common/basemodel/test_overlap_utils.py | 125 +- unit_tests/common/test_paged_mem_manager.py | 59 - .../test_quantized_paged_mem_manager.py | 33 - unit_tests/common/test_req_manager_page.py | 149 +- unit_tests/conftest.py | 40 + unit_tests/models/glm5_next/test_cache.py | 266 ++- .../models/glm5_next/test_indexer_topk.py | 2 +- unit_tests/models/glm5_next/test_kernels.py | 53 +- unit_tests/models/glm5_next/test_layers.py | 48 +- unit_tests/models/glm5_next/test_mtp.py | 6 +- unit_tests/models/glm5_next/test_pd_cache.py | 282 ++- .../models/test_qwen3_dspark_model_output.py | 6 +- .../test_cache_placement_controller.py | 2 +- .../test_multimodal_cache_matching.py | 12 +- .../router/dynamic_prompt/test_radix_cache.py | 109 +- .../mode_backend/test_dp_big_pages.py | 186 ++ .../mode_backend/test_multi_level_kv_cache.py | 15 +- .../mode_backend/test_nccl_control_channel.py | 51 + .../test_paged_side_allocations.py | 109 +- .../mode_backend/test_pd_notify.py | 33 +- .../mode_backend/test_prompt_cache_preload.py | 52 - .../model_infer/test_cpu_resource_layout.py | 457 +++++ .../test_hybrid_cpu_state_cache.py | 202 +- .../model_infer/test_multimodal_chunking.py | 86 +- unit_tests/server/router/test_kv_capacity.py | 60 + unit_tests/server/test_dp_request_capacity.py | 1 - unit_tests/test_static_benchmark_paged_kv.py | 21 +- .../utils/test_config_utils_fixed_kv.py | 33 - 221 files changed, 18609 insertions(+), 5641 deletions(-) create mode 100644 docs/cache_manager_design.md create mode 100644 lightllm/common/att_resources/__init__.py create mode 100644 lightllm/common/att_resources/allocators/__init__.py create mode 100644 lightllm/common/att_resources/allocators/full_att_kv_allocator.py create mode 100644 lightllm/common/att_resources/allocators/state_page_allocator.py create mode 100644 lightllm/common/att_resources/att_resource_manager.py create mode 100644 lightllm/common/att_resources/att_tail/__init__.py create mode 100644 lightllm/common/att_resources/att_tail/base_state.py create mode 100644 lightllm/common/att_resources/att_tail/create_kpool_state.py create mode 100644 lightllm/common/att_resources/att_tail/kpool_state.py create mode 100644 lightllm/common/att_resources/att_tail/transfer_utils.py create mode 100644 lightllm/common/att_resources/bounded_att/__init__.py create mode 100644 lightllm/common/att_resources/bounded_att/base_state.py create mode 100644 lightllm/common/att_resources/bounded_att/conv_ssm_types.py create mode 100644 lightllm/common/att_resources/bounded_att/create_gdn_state.py create mode 100644 lightllm/common/att_resources/bounded_att/create_kda_state.py create mode 100644 lightllm/common/att_resources/bounded_att/gdn_pages_pool.py create mode 100644 lightllm/common/att_resources/bounded_att/gdn_state.py create mode 100644 lightllm/common/att_resources/bounded_att/kda_pages_pool.py create mode 100644 lightllm/common/att_resources/bounded_att/kda_state.py create mode 100644 lightllm/common/att_resources/bounded_att/transfer_utils.py create mode 100644 lightllm/common/att_resources/bounded_att/utils.py create mode 100644 lightllm/common/att_resources/full_att/__init__.py create mode 100644 lightllm/common/att_resources/full_att/base_state.py create mode 100644 lightllm/common/att_resources/full_att/create_gqa_state.py create mode 100644 lightllm/common/att_resources/full_att/create_mla_state.py create mode 100644 lightllm/common/att_resources/full_att/gqa_state.py create mode 100644 lightllm/common/att_resources/full_att/gqa_transfer.py create mode 100644 lightllm/common/att_resources/full_att/group_quantized_state.py create mode 100644 lightllm/common/att_resources/full_att/group_quantized_transfer.py create mode 100644 lightllm/common/att_resources/full_att/mla_state.py create mode 100644 lightllm/common/att_resources/full_att/mla_transfer.py create mode 100644 lightllm/common/att_resources/full_att/packed_mla_state.py create mode 100644 lightllm/common/att_resources/full_att/packed_mla_transfer.py create mode 100644 lightllm/common/att_resources/full_att/static_fp8_state.py create mode 100644 lightllm/common/att_resources/full_att/static_fp8_transfer.py create mode 100644 lightllm/common/att_resources/types.py create mode 100644 lightllm/common/att_resources/utils.py create mode 100644 lightllm/common/basemodel/basemodel_utils.py rename lightllm/{models/deepseek3_2/triton_kernel => common/basemodel/triton_kernel/kv_copy}/destindex_copy_kv_flashmla_fp8.py (82%) rename lightllm/{models/deepseek3_2/triton_kernel => common/basemodel/triton_kernel/kv_copy}/prefill_compact_kv_flashmla_fp8.py (87%) delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py delete mode 100644 lightllm/common/kv_cache_mem_manager/__init__.py delete mode 100644 lightllm/common/kv_cache_mem_manager/allocator.py delete mode 100644 lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py delete mode 100755 lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py delete mode 100755 lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py delete mode 100755 lightllm/common/kv_cache_mem_manager/mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/mem_utils.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/__init__.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/base.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/deepseek.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/linear_att.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/normal.py delete mode 100644 lightllm/common/kv_cache_mem_manager/operator/quant.py delete mode 100755 lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py delete mode 100755 lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py delete mode 100644 lightllm/common/req_manager/glm5_next.py delete mode 100644 lightllm/common/req_manager/hybrid_base.py delete mode 100644 lightllm/common/req_manager/linear_att.py delete mode 100644 lightllm/common/state_cache_manager/__init__.py delete mode 100644 lightllm/common/state_cache_manager/base.py delete mode 100644 lightllm/common/state_cache_manager/glm5_next.py delete mode 100644 lightllm/common/state_cache_manager/layer_cache.py delete mode 100644 lightllm/common/state_cache_manager/linear_att.py create mode 100644 lightllm/models/glm5_next/att_resource_config.py create mode 100644 lightllm/models/glm5_next/att_resource_manager.py create mode 100644 lightllm/models/qwen3next/att_resource_config.py create mode 100644 lightllm/models/qwen3next/att_resource_manager.py create mode 100644 lightllm/server/router/model_infer/mode_backend/pd/resource_io.py create mode 100644 lightllm/server/router/model_infer/shared_att_resources_utils.py create mode 100644 lightllm/utils/att_resource_config.py create mode 100644 test/att_resources/validate_model_cache.py create mode 100644 unit_tests/common/att_resources/conftest.py create mode 100644 unit_tests/common/att_resources/manager_fixtures.py create mode 100644 unit_tests/common/att_resources/model_helpers.py create mode 100644 unit_tests/common/att_resources/test_allocators.py create mode 100644 unit_tests/common/att_resources/test_backend_startup.py create mode 100644 unit_tests/common/att_resources/test_bounded_state.py create mode 100644 unit_tests/common/att_resources/test_capacity.py create mode 100644 unit_tests/common/att_resources/test_compute_binding.py create mode 100644 unit_tests/common/att_resources/test_fp8_compute_binding.py create mode 100644 unit_tests/common/att_resources/test_full_att_cpu_kernel.py create mode 100644 unit_tests/common/att_resources/test_full_att_transfer.py create mode 100644 unit_tests/common/att_resources/test_full_state.py create mode 100644 unit_tests/common/att_resources/test_full_storage_initialization.py create mode 100644 unit_tests/common/att_resources/test_gqa_pd_kernel.py create mode 100644 unit_tests/common/att_resources/test_gqa_transfer.py create mode 100644 unit_tests/common/att_resources/test_manager_helpers.py create mode 100644 unit_tests/common/att_resources/test_model_calibration.py create mode 100644 unit_tests/common/att_resources/test_model_initialization.py create mode 100644 unit_tests/common/att_resources/test_model_resources.py create mode 100644 unit_tests/common/att_resources/test_model_storage_definition.py create mode 100644 unit_tests/common/att_resources/test_mtp_state_commit.py create mode 100644 unit_tests/common/att_resources/test_network_transport.py create mode 100644 unit_tests/common/att_resources/test_no_attention_draft.py create mode 100644 unit_tests/common/att_resources/test_page_pool_storage.py create mode 100644 unit_tests/common/att_resources/test_paged_storage.py create mode 100644 unit_tests/common/att_resources/test_quantized_full_state.py create mode 100644 unit_tests/common/att_resources/test_quantized_paged_storage.py create mode 100644 unit_tests/common/att_resources/test_role_transfers.py create mode 100644 unit_tests/common/att_resources/test_runtime_transport.py create mode 100644 unit_tests/common/att_resources/test_state_creation.py create mode 100644 unit_tests/common/att_resources/test_state_identity.py create mode 100644 unit_tests/common/att_resources/test_state_queries.py create mode 100644 unit_tests/common/att_resources/test_tail_state.py create mode 100644 unit_tests/common/att_resources/test_tp_config.py delete mode 100644 unit_tests/common/test_paged_mem_manager.py delete mode 100644 unit_tests/common/test_quantized_paged_mem_manager.py create mode 100644 unit_tests/conftest.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py create mode 100644 unit_tests/server/router/model_infer/test_cpu_resource_layout.py create mode 100644 unit_tests/server/router/test_kv_capacity.py delete mode 100644 unit_tests/utils/test_config_utils_fixed_kv.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 0eb8a0583c..9792fe5778 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -109,7 +109,7 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_mem_manager() + self._init_req_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -124,10 +124,14 @@ class TpPartBaseModel: | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | 读取初始化模型的 config.json, 并进行一些 key 名的同名合法化操作 | | def _verify_params(self)： | 校验参数 | -| def _init_mem_manager(self): | 初始化 token attention 使用的 mem manager 对象 | +| def _init_att_resource_manager(self): | 由具体模型创建管理器并准备全部 attention State：KV meta State、GDN/KDA 固定运行态和 indexer tail；由模型侧 helper 根据 main 和前序 draft 的配置解析静态量化标定偏移 | | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_att_resources()` 准备自身 State，再统一分配共享资源。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 + +资源准备仍通过私有 hook `_init_att_resource_manager` 交给具体模型或模型族实现。标定偏移由模型侧根据 main 和前序 draft 的 full-attention 层定义计算，构造函数不调用该 hook。hook 内部可复用 `basemodel_utils.init_model_full_attention` 构建普通 KV 定义；helper 不构成独立启动阶段。 + ## 2. 添加 bloom 模型的示例说明 具体实现在 ***lightllm/models/bloom*** 目录下，下面的代码片段请对应源码进行阅读，其中 triton_kernel 目录下为推理类使用的一些 kernel，下文中不做详细介绍，同时 bloom 模型因为不需要传递特殊状态信息使用默认的状态类即可。如想更深入的理解整个框架，可以进一步参考 llama 和 llama2 等模型的接入实现源码。 @@ -319,9 +323,6 @@ class BloomPreLayerInfer(PreLayerInferTpl): self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1) return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) - def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight): total_token_num = infer_state.total_token_num input_ids = input_ids[0:total_token_num] @@ -346,6 +347,9 @@ class BloomPreLayerInfer(PreLayerInferTpl): dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False) input_embdings = self._norm(input_embdings, infer_state, layer_weight) return input_embdings + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) ~~~ ***transformer_layer_infer.py*** @@ -416,9 +420,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl): def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor: o_tensor = torch.empty_like(q) + k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_) token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_), - infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :], - infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :], + k_cache, + v_cache, o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_), layer_weight.tp_alibi, infer_state.b_loc, @@ -470,9 +475,6 @@ class BloomPostLayerInfer(PostLayerInferTpl): self.embed_dim_ = network_config["n_embed"] return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) - def soft_max(self, data): return torch.softmax(data.permute(1, 0).float(), dim=-1) @@ -507,6 +509,9 @@ class BloomPostLayerInfer(PostLayerInferTpl): ans_logics = gather_data.permute(1, 0).float() gather_data = None return ans_logics + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) ~~~ ### （3） 实现模型的框架类 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index d9436acf08..ee8e7d9f24 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -109,7 +109,7 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_mem_manager() + self._init_req_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -124,10 +124,14 @@ Commonly used interfaces that need to be inherited and implemented | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | Read the config.json for initializing the model and perform some key name legalization operations | | def _verify_params(self)： | Validate parameters | -| def _init_mem_manager(self): | Initialize the mem manager object used by token attention | +| def _init_att_resource_manager(self): | Create the concrete manager and prepare all attention States: KV meta States, fixed GDN/KDA runtime States and indexer tails; derive static quantization calibration offsets in model helpers from the main and preceding draft configurations | | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | +After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_att_resources()`, then allocates shared resources. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. + +Concrete models or model families still implement the private `_init_att_resource_manager` hook. Model helpers derive calibration offsets from the full-attention layer definitions of the main and preceding drafts; constructors do not call this hook. The hook can reuse `basemodel_utils.init_model_full_attention` to build one KV scheme; this helper is not a separate startup phase. + ## 2. Example of Adding Bloom Model The specific implementation is under the ***lightllm/models/bloom*** directory. Please read the source code for the code snippets below. The triton_kernel directory contains some kernels used by inference classes, which will not be introduced in detail in this article. At the same time, the bloom model uses the default state class because it doesn't need to pass special state information. For a deeper understanding of the entire framework, you can further refer to the source code implementation of llama and llama2 model integration. @@ -320,9 +324,6 @@ class BloomPreLayerInfer(PreLayerInferTpl): self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1) return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) - def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight): total_token_num = infer_state.total_token_num input_ids = input_ids[0:total_token_num] @@ -347,6 +348,9 @@ class BloomPreLayerInfer(PreLayerInferTpl): dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False) input_embdings = self._norm(input_embdings, infer_state, layer_weight) return input_embdings + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) ~~~ ***transformer_layer_infer.py*** @@ -417,9 +421,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl): def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor: o_tensor = torch.empty_like(q) + k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_) token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_), - infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :], - infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :], + k_cache, + v_cache, o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_), layer_weight.tp_alibi, infer_state.b_loc, @@ -471,9 +476,6 @@ class BloomPostLayerInfer(PostLayerInferTpl): self.embed_dim_ = network_config["n_embed"] return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) - def soft_max(self, data): return torch.softmax(data.permute(1, 0).float(), dim=-1) @@ -508,6 +510,9 @@ class BloomPostLayerInfer(PostLayerInferTpl): ans_logics = gather_data.permute(1, 0).float() gather_data = None return ans_logics + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) ~~~ ### （3） Add implementation model class diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md new file mode 100644 index 0000000000..9ec820c1b6 --- /dev/null +++ b/docs/cache_manager_design.md @@ -0,0 +1,1702 @@ +# Attention Resources 架构设计草案 + +本文定义 attention 资源的架构设计，包括 `AttResourceManager` 的职责、共享 allocator、各类 State 的接口、目录组织及 CPU/DP/PD 传输约定。文中的结构与接口为设计契约，待细化事项集中列于第 9 节。 + +## 1. 设计约定 + +统一在 `lightllm/common/att_resources/` 下组织可复用的资源实现，由 `AttResourceManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。模型可以直接使用该类，仅在资源创建或操作存在差异时提供子类。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 + +公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： + +```python +class AttResourceManager: + full_att_allocator: FullAttKVAllocator | None + bounded_att_big_page_allocator: StatePageAllocator | None + bounded_att_small_page_allocator: StatePageAllocator | None +``` + +这三个入口表示资源管理器对分配域的管理责任，聚合操作与传输接口见第 6.2 节。具体初始化必须保证所需 allocator 已创建或绑定，无需的 allocator 为 `None`。例如纯 full-attention 模型只初始化 `full_att_allocator`；启用 bounded 状态页的组合再按配置初始化大小页 allocator。GPU 运行态的存在本身不要求启用页池。 + +- 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att`、`main_att_tail`、`draft_att_tail` 成员。具体模型按需组织资源组合与实现类型，拥有的 State 成员统一遵守第 3.5 节的角色、类型和存储方案命名。 +- `full_att`、`bounded_att`、`att_tail` 作为可复用实现的目录分类。full attention 保存历史 KV；bounded attention 由 `BoundedAttentionState` 抽象基类约定操作，当前包含 `KDAState`、`GDNState`，以后可加入 sliding-window 等有界状态，各自定义运行态与可选页池；attention tail 保存索引、池化、压缩等模块在增量计算中需要跨步骤保留的尾部数据及控制信息；具体实现决定恢复边界和快照要求。 +- main/draft 可以采用不同方案；同一共享分配域的参与方使用同一个 allocator 实例，遵守统一的索引和回收约定。资源对象持有 buffer 和布局信息；管理器创建或绑定共享 allocator，调用方直接使用对应 allocator 分配与归还索引。 +- `AttResourceManager` 提供运行态清零、清零指定请求的 tail、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 +- `req_to_token_indexs` 和请求槽位仍由 `ReqManager` 持有。资源管理器接收显式传入的物理索引、请求槽位及恢复长度，提供资源操作接口。 +- 内部调用契约使用 `assert`，例如页类型、物理索引范围、Tensor 形状及共享分配域的一致性；用户配置、外部页面内容及格式兼容性继续显式校验，资源不足等运行失败保留相应异常。GPU 上的 MTP 状态选择保留异步断言，避免引入主机同步。 +- 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 + +本文区分已确定的职责边界和接口草案。main/draft 的资源由该模型组合的具体管理器统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 + +### 1.1 接口命名、排列与注释规范 + +新增 Python 文件、类和接口按 **公共函数在前、私有辅助函数在后** 排列。已有基础流程文件以保留功能、减少 review 差异为优先，保留原有函数相对顺序，不为满足分组规范整体重排。 + +- 模块中先定义公共函数，再定义以下划线开头的私有辅助函数（双下划线协议函数除外）。文件名以 `_` 开头不改变文件内函数的排列规则。脚本的 `if __name__ == "__main__"` 执行入口仍放在全部定义之后。 +- 类的方法区依次排列：`__new__`、`__init__`、其他公共方法与属性、私有辅助方法。`__getstate__`、`__post_init__` 等 Python 双下划线协议方法属于公共协议，不作为私有辅助方法处理。 +- `async def`、`classmethod`、`staticmethod`、`property`、抽象方法均按方法名分类。属性 getter/setter 等存在定义依赖的方法保持必要顺序，装饰器和关联注释随函数一起移动。 +- 同一组内保留原有相对顺序，便于按接口职责阅读；不按字母排序，不为了调整位置修改函数名称或增加转发层。 +- 模块说明、导入、常量、类型定义和类级数据声明保留必要的依赖顺序。重排须检查装饰器、默认参数、类型注解及模块初始化语句的定义时依赖，并保持函数体与业务语义不变。 + +本规范适用于资源管理器、State、allocator、helper、相关模型与 server 调用代码，以及对应的测试和验证脚本。新增文件和新增类必须遵守上述完整分组规则；在存量文件中增加接口时，新增公共接口放在入口区域，新增私有辅助函数放在相关实现区域，新增代码之间仍按公共在前、私有在后排列，不牵动已有函数的位置。 + +模型初始化、推理执行、调度、缓存和传输等已有基础流程，只保留实现功能所需的改动，不混入纯函数排序、接口改名或无关格式整理。函数名称及 public/private 属性是否调整由功能需要决定，不因排列规范修改；功能重构与存量代码的风格整理分别进行。 + +State 的基类与具体实现同样使用中文大块分隔注释，依次组织基础定义、该类资源的数据操作、CPU Cache、DP、PD、私有辅助实现。KDA/GDN 直接继承 `BoundedAttentionState`；所有具体 full-attention State 直接继承 `FullAttentionState`。具体 State 之间不继承实现，分别在自己的文件中按上述分块组织完整操作。 + +- 分类数据操作统一带 `full_att`、`bounded_att` 或 `att_tail` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_pages`、`reset_attention_tail_state`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 +- `get_layer_view` 用于计算访问，CPU/DP/PD 保持按传输方向命名。State 不增加仅用于统一名字的转发方法。 +- 所有 attention 配置字段必须提供中文注释，说明业务含义、单位或对应 Tensor 轴、全局/TP 本地范围，以及适用方案和默认值的特殊语义；派生属性还需说明计算关系。资源身份只保存在 State，不放入 config。 +- 已知的调用参数使用显式 `key=value`，如 `size=...`、`dtype=...`、`device=...`、`pin_memory=...`；不先组装 `options` 字典再通过 `**options` 传入。 +- 资源身份、布局范围和分配域选择等必要检查直接在使用处展开，不再包装成只有断言、布尔判断或转发循环的小函数。同一操作已校验的槽位继续复用，避免逐页、逐 State 重复校验分配器。 +- 每个公共接口使用中文 docstring 说明操作对象、索引含义和数据范围；在有差异时明确量化参数、前缀长度、MTP 游标及完整 tail 的处理。私有实现使用独立分块，关键布局或状态转换另加说明。 +- bounded 运行态重置、大小页保存与恢复、MTP 提交和运行态复制集中排列。pool 只保存快照 buffer；前缀与槽位的关联、页面发布和引用生命周期由调用流程管理。 + +## 2. 整体结构与职责边界 + +```mermaid +flowchart TD + Model[模型初始化] --> Concrete[模型专属 AttResourceManager] + Concrete -.->|实现公共接口| API[AttResourceManager] + Context[InferenceContext] --> Req[ReqManager] + Req --> Slots[请求槽位与 KV 映射] + Context -->|前缀复用| Radix[Radix Cache 功能] + Context -->|多级缓存| CPU[CPU Cache 功能] + Radix -->|调用通用状态页操作| Concrete + CPU -->|缓存数据读写| Concrete + Concrete --> Alloc[按需初始化或绑定的共享 allocator] + Concrete --> Internal[模型自行组织的内部资源] + Internal --> Full[full attention 实现] + Internal --> Bounded[bounded attention：KDAState / GDNState 等] + Internal --> Indexer[attention tail 实现] + Concrete -->|绑定具体视图| Forward[模型与 attention backend] + Transport[DP/PD 调度] -->|调用具体资源操作| Concrete +``` + +图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享该组合的具体管理器，分别绑定各自的计算 State。 + +| 层次 | 组件 | 职责 | +| --- | --- | --- | +| 请求流程 | `InferenceContext`、`ReqManager` | 创建、暂停、恢复和完成请求；由 ReqManager 管理请求槽位、KV 索引映射及采样参数 | +| 资源数据操作 | `AttResourceManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | +| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用具体管理器操作资源数据 | +| 资源实现 | 模型所需的 State、存储布局、checkpoint pool | 持有 buffer 和布局，实现读写、复制、重置和快照；IPC 导出范围由管理器统一处理 | +| 传输调度 | DP/PD 与多级缓存模块 | 管理任务、队列、通信、事件同步及完成通知 | + +`AttResourceManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，具体管理器执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 + +代码依赖从调用流程向资源层单向传递。radix cache 依赖管理器的通用状态页操作；管理器与 State 不导入或访问 radix 节点、树和策略。资源实现也不反向导入请求调度模块；传输调度使用数据读写与复制接口。 + +## 3. 对象结构与资源归属 + +```text +ReqManager +├── request_slots +└── req_to_token_indexs # 共享请求到 KV 的索引映射 + +AttResourceManager # 公共接口与 allocator 入口 +├── full_att_allocator: FullAttKVAllocator | None +├── bounded_att_big_page_allocator: StatePageAllocator | None +└── bounded_att_small_page_allocator: StatePageAllocator | None + +AttResourceManager # 模型实现，内部结构由其自行定义 +├── 按需创建或绑定所需 allocator +├── 初始化模型所需的资源组合 +├── 实现内部状态页聚合操作与 CPU / DP / PD 数据传输 +└── 绑定推理视图、组织传输所需状态 # 搬运直接调用 State 方法 +``` + +例如纯 full-attention 管理器按角色持有具体 KV 方案；混合模型管理器组合不同算法、存储方案的运行态及 indexer tail。main/draft 分别持有各自的 State，共享 allocator 按分配域创建。公共调用方通过接口使用这些组合，通用聚合操作复用基类实现。 + +以下列出可复用资源组件的结构示例。linear attention 直接使用 `KDAState`、`GDNState` 等具体状态对象，由 `BoundedAttentionState` 约定必需操作，不增加资源包装实例。管理器内部的成员遵守统一命名及唯一性约束，checkpoint pool 仅在相应功能启用时存在： + +```text +FullAttentionState +├── capacity / page_size / device # 普通 token 容量、HOLD 页大小及分配设备 +├── storage_dtype # 实际 KV 编码类型；普通方案使用 config.dtype，INT4/8 为 int8，静态 FP8 为 uint8 +├── k_storage / v_storage # GQA 及其量化方案分别分配 K/V；MLA 使用单份 latent storage +└── layer_mapping # 本模型的层映射 + +KDAState # 一种具体 attention 状态实现 +├── conv_state: torch.Tensor # 卷积历史窗口及 MTP 尾部空间 +├── ssm_state: torch.Tensor # KDA 递推状态及 MTP 状态槽位 +├── mtp_state_index: Tensor | None +├── big_pages_pool # 按需创建，使用共享 allocator 分配的槽位 +│ ├── conv_state +│ └── ssm_state +└── small_pages_pool # 按需创建 + ├── conv_state + └── ssm_state + +GDNState # 独立定义 GDN 状态、布局和快照操作 +└── 按其实现组织 Conv / SSM / MTP 状态及可选大小页 pool + +AttentionTailState # 抽象操作契约，不固定内部存储 +└── KPoolAttentionTailState # GLM K-pool 的具体状态 + └── tail_buffer # 原始 key/gate ring，按层、请求组织 +``` + +同一共享分配域中的 allocator 只初始化一次，由资源管理器负责其完整生命周期。若 main/draft 分别使用管理器实例，装配时应注入共享 allocator 引用，并明确唯一的创建和销毁责任，以及调用方解除全部引用并向 allocator 单次归还的约定；若采用一个组合管理器，则由该实现协调内部资源。两种装配方式均不能因模型数增加而重复分配或释放同一共享索引。实例装配形式保留为待定项。 + +| 资源 | 持有者 | 分配单位 / 生命周期 | +| --- | --- | --- | +| full-attention 空闲页与容量状态 | `AttResourceManager.full_att_allocator` | 每次为请求统一分配/回收物理页索引，main/draft 共用 | +| 大页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_big_page_allocator` | 为 main/draft 参与的大页 pool 统一分配槽位 ID | +| 小页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | +| 请求 KV 索引表 | `ReqManager` | 记录逻辑 token 到共享物理索引的映射；负责绑定、更新和清理 | +| HOLD 保留物理页及索引描述 | `AttResourceManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | +| KV、量化 scale、packed indexer buffer 与层映射 | 具体管理器组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | +| KDA/GDN 的 Conv/SSM、MTP 游标 | 具体管理器持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | +| 大小页 checkpoint buffer | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID 保存数据；不持有独立空闲列表或长度记录 | +| indexer 原始尾部窗口 | 具体管理器组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | +| 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | +| 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复关联资源数据，解除引用后直接归还 allocator | + +### 3.1 main/draft 共享 full-attention allocator + +调用方直接通过 `AttResourceManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，具体管理器恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： + +```text +逻辑 token 位置 t + ↓ ReqManager.req_to_token_indexs +统一物理索引 i + ├── main 对应的 KV 存储位置 + └── draft 对应的 KV 存储位置 +``` + +二者的层数、head 维度、dtype 和量化布局可以不同；同一索引域内需要使用兼容的 page_size，并且各存储都覆盖 allocator 可返回的索引范围。同一分配域统一容量和物理页索引规则，不引入模型间索引重映射。 + +共享索引表示存储位置的对应关系，不表示 main/draft 在每个位置都已经写入有效 KV。MTP 预留区和回退后的有效长度仍由请求执行流程维护。 + +释放时，请求流程与具体管理器确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 + +### 3.2 main/draft 共享 linear checkpoint allocator + +main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAState`，draft 使用 `GDNState` 或无需 linear 状态。各状态对象持有自身运行态及启用的大页或小页 pool。对同一前缀的一个大页或小页 checkpoint，调用方直接从管理器持有的对应 allocator 分配一次槽位 ID，再将该 ID 传给管理器的聚合数据操作，供参与的两侧 pool 使用。两侧的状态布局、dtype 和单槽位字节数可以不同，各自的存储必须覆盖共享 allocator 可返回的槽位范围。 + +例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。具体管理器负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 + +共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 + +linear 运行态与两侧 indexer tail 均由 `ReqManager` 提供的 `req_idx` 定位。indexer tail 不单独分配 checkpoint 槽位；其状态布局和恢复规则由各自实现描述。 + +### 3.3 分配状态与数据状态分离 + +allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；具体管理器组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 + +同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 + +对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由具体管理器维护。 + +### 3.4 按具体 attention 实现定义状态 + +定义只约束操作的 `BoundedAttentionState` 抽象基类，由 `KDAState`、`GDNState` 等子类负责自己的成员、布局、重置、快照与恢复语义。`AttResourceManager` 初始化并组合所需对象，公共缓存流程通过管理器访问它们。不同模型也可以复用同一种状态实现。 + +`KDAState` 的成员定义如下： + +```python +class KDAState(BoundedAttentionState): + config: ConvSSMStateConfig # 仅包含算法布局和 TP 信息 + resource_id: str | None # 构造时为 None，由管理器绑定模型角色时设置 + request_capacity: int # 包含 HOLD 请求槽位 + conv_state: torch.Tensor + ssm_state: torch.Tensor + mtp_state_index: torch.Tensor | None + mtp_step: int # 该实例的 MTP 布局参数 + big_pages_pool: KDAPagesPool | None + small_pages_pool: KDAPagesPool | None +``` + +State 直接持有按层组织的 Tensor，不额外设置通用状态管理对象。该示例中的 Conv、SSM 和 MTP 索引驻留 GPU；checkpoint pool 独立持有快照 buffer。设请求槽位容量为 `N`（包含保留槽位），该实例的 `mtp_step` 为 `S`： + +| 成员 | 内容与索引方式 | +| --- | --- | +| `conv_state` | 对应算法的卷积历史窗口，窗口宽度为 `conv_kernel_size - 1 + S`；用 `buffer[layer, req_idx, ...]` 访问 | +| `ssm_state` | KDA 的递推矩阵状态，每请求预留 `S + 1` 个状态槽位；第 `j` 个槽位位于 `buffer[layer, req_idx * (S + 1) + j, ...]` | +| `mtp_state_index` | 形状为 `[N]` 的 `int32` Tensor；记录各请求接受的 MTP 状态索引。`S == 0` 时为 `None` | +| 大小页 pool | 按需创建，保存该算法指定前缀位置的状态；槽位来自管理器的对应 `bounded_att_*_allocator` | + +`GDNState` 按 GDN 算法独立定义 Conv/SSM 和 MTP 状态，其大小页使用 `GDNPagesPool`，支持不同的 K/V head 数及维度。`KDAPagesPool` 按 KDA 的相同 Q/K/V head 数和方阵递推状态定义布局。两类 pool 各自分配快照存储，不额外记录长度或页面有效性；前缀与槽位的关联由缓存调用流程维护。两种 State 可以复用兼容的 Conv/SSM 数据操作，共同实现 `BoundedAttentionState` 接口；基类不规定 Conv/SSM 等成员，也不意味着两种算法的状态内容可以互换。 + +运行态按 `ReqManager` 的槽位容量预分配，状态对象接收 `req_idx` 执行操作。main/draft 可以持有同一种状态类的两个独立实例，也可以使用不同状态类型；各实例的层数、维度、dtype 和 MTP 配置由对应模型决定。allocator 仍由 `AttResourceManager` 创建或绑定并统一管理。 + +后续 sliding-window attention 可以定义实现 `BoundedAttentionState` 的 `SlidingWindowState`，按窗口 KV、有效范围和游标等实际需求组织成员及自己的页池。它可以使用适用的 bounded-attention 分配域，其页布局、恢复边界和已接受状态选择须按窗口语义实现；不要求包含 Conv/SSM 字段，也不复用 KDA/GDN 的页池。统一操作接口见 6.7 节。 + +### 3.5 State 成员命名、资源身份与唯一性 + +`AttResourceManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `set_att_state(model_name, state)` 绑定。`model_name` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.resource_id`。`FullAttentionConfig`（含量化与 MLA 子类）、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `resource_id` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_creation()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 + +管理器拥有的 State 成员名与 `state.resource_id` 完全一致，统一使用 `
__state`。draft 按加载顺序从 0 编号，最多 36 个；例如 `draft0_kda_state` 与 `draft1_kda_state` 是两个独立对象。三类 State 基类统一声明抽象接口 `get_state_type() -> str`，各具体方案必须实现，返回算法及存储方案标记，不再保存 `state_type` 属性。GQA、MLA、量化与 K-pool 实现根据 config 返回标记，KDA/GDN 分别返回 `"kda"` / `"gdn"`；定制方案通过重写接口区分。现有标记字符串保持不变，资源绑定后返回值必须稳定，不随容量或设备变化；CPU/DP/PD 布局直接使用相同的 `resource_id` 关联资源，无需转换命名格式。同一管理器内每个 `(role, state_type)` 最多一个实例，启用时必须有且只有一个。同一模型同方案的全部层放入同一个 State，不按层拆成同类型列表，不通过别名重复持有同一 State。 + +| 方案 | State 成员名及 `resource_id` | +| --- | --- | +| 普通 BF16 KV | `main_full_att_bf16_state` | +| INT8，group size 8，BF16 scale | `main_full_att_int8_group8_bf16_scale_state` | +| 静态 FP8，per-head scale | `main_full_att_fp8_static_per_head_state` | +| MLA BF16 | `draft0_mla_bf16_state` | +| 带 indexer 字段的 packed MLA FP8 | `main_mla_packed_fp8_state` | +| KDA / GDN | `main_kda_state` / `draft0_gdn_state` | +| K-pool C4 indexer tail | `main_indexer_kpool_c4_tail_state` | + +普通 KV 不加 `dense`；MLA 直接使用 `mla` 标记，不叠加 `full_att_mla`。普通浮点存储区分 `bf16`、`fp16`、`fp32`；量化方案区分 INT4/INT8、group size 与 scale dtype，静态 FP8 区分 per-head/per-tensor。`packed` 表示 KV 与 indexer 字段联合存储，与纯 MLA 区分。Indexer 使用算法及压缩比例标记，例如 `indexer_kpool_c4_tail`，不以通用 `att_tail` 掩盖具体实现。完整 shape、标定身份等仍保存在 State 配置及存储描述中。 + +以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： + +```text +Glm5NextAttResourceManager +├── full_att_allocator +├── bounded_att_big_page_allocator +├── bounded_att_small_page_allocator +├── main_mla_packed_bf16_state +├── main_kda_state +├── main_indexer_kpool_c4_tail_state +├── draft0_mla_packed_bf16_state +└── draft0_indexer_kpool_c4_tail_state +``` + +当前 GLM draft 不使用 KDA，因此不创建 `draft0_kda_state`。Qwen 的线性状态为 `main_gdn_state`；Full Attention 按实际配置创建 `main_full_att_bf16_state` 等成员。MLA 与普通 KV 都可使用同一个 `full_att_allocator` 分配域。 + +`vanilla_no_att`、`eagle_no_att` 等仅执行 FFN 的 draft 不创建 KV State，也不初始化 attention backend；其计算引用可以为 `None`，仍共享管理器的 allocator 与请求映射。缺少实际需要的 draft attention State 则应在初始化时断言失败。 + +因此同一角色可以同时拥有 `main_full_att_bf16_state` 和 `main_full_att_int8_group8_bf16_scale_state`；这是两种不同方案，各自覆盖对应的层，聚合复制及传输必须覆盖二者。现有内置模型按配置选择自己的方案；未来混合量化模型在管理器中显式定义多个规范成员，基类通过聚合查询将它们传给相同 helper，不需要增加通用 State 注册表。 + +main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 State 同时命名为 main 和 draft。当前启动流程为各模型分别创建 State 和存储；GLM 的每个模型也分别创建 indexer ring。三种 allocator 的名称不带模型角色，同一分配域只初始化或绑定一个实例。 + +具体管理器只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_resource_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `set_att_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 + +`InferStateInfo` 保存 `att_resource_manager`、`is_draft_model` 和资源角色 `att_resource_role`，不额外绑定 `full_att_state`。各层直接访问管理器上对应方案的 State 成员，再通过局部引用取得层视图或写入 KV；FP8 backend 从具体量化方案的成员读取固定参数。每个 draft 的层由自身编号对应的 State 管理；`att_resource_role` 为 `main` 或 `draft0` 到 `draft35`，不能仅根据 `is_draft_model` 选择资源。是否初始化 full-attention backend 根据管理器中该角色实际启用的资源判断；FFN-only 和 linear-only 模型不创建 full-attention backend，linear backend 由相应模型独立初始化。 + +资源身份用于 CPU/PD 资源布局与 DP 映射，所有传输按完整角色和方案匹配。部署配置确定各资源的页布局；PD 参数兼容性由 PD master 在连接阶段校验，CPU 服务绑定时核对页大小，磁盘缓存按 main/draft 模型版本隔离。 + +## 4. 目录设计 + +目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器、公共 State、allocator 和 pool 文件按对应类名命名；具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 + +```text +lightllm/common/ +├── basemodel/ +│ ├── basemodel.py # 模型基础流程与模型专属配置 hook +│ └── basemodel_utils.py # 配置归一化、KV 选项与 full/linear State 定义构建 +├── att_resources/ +│ ├── __init__.py +│ ├── utils.py # State 角色与精度命名 +│ ├── att_resource_manager.py # AttResourceManager：State 装配、容量与存储初始化、缓存和传输接口 +│ ├── types.py # TensorField / StateTransferLayout 等 +│ ├── allocators/ # 管理器直接持有，同一分配域共享实例 +│ │ ├── __init__.py +│ │ ├── full_att_kv_allocator.py # FullAttKVAllocator / KVCapacityReader：full_att 分配域及共享容量读取 +│ │ └── state_page_allocator.py # StatePageAllocator：bounded_att 大小页各一个实例 +│ ├── full_att/ +│ │ ├── __init__.py +│ │ ├── base_state.py # FullAttentionConfig / FullAttentionState：通用配置与抽象接口 +│ │ ├── create_gqa_state.py # GQA 统一创建入口：选择普通精度、INT4/INT8 或静态 FP8 +│ │ ├── gqa_state.py # GQAFullAttentionState：兼容 MHA / GQA / MQA +│ │ ├── gqa_transfer.py # GQA CPU/PD 共用的 K/V 页传输入口与 kernel +│ │ ├── group_quantized_state.py # GroupQuantizedFullAttentionState：INT4 / INT8 等 +│ │ ├── group_quantized_transfer.py # INT4/INT8 KV 与动态 scale 的定制传输 +│ │ ├── static_fp8_state.py # StaticFP8FullAttentionState +│ │ ├── static_fp8_transfer.py # FP8 K/V 原始编码的定制传输 +│ │ ├── create_mla_state.py # MLA 统一创建入口：选择普通、packed 或 packed FP8 +│ │ ├── mla_state.py # MLAFullAttentionState：压缩 latent 存储 +│ │ ├── mla_transfer.py # latent KV 的定制传输 +│ │ ├── packed_mla_state.py # PackedMLAFullAttentionState：含 indexer key/scale 的 packed 布局 +│ │ └── packed_mla_transfer.py # packed KV、RoPE 与 indexer 字段的定制传输 +│ ├── bounded_att/ +│ │ ├── __init__.py +│ │ ├── base_state.py # BoundedAttentionState：有界状态的抽象操作契约 +│ │ ├── utils.py # 页池分配策略：DP fetch 使用共享 pinned 存储 +│ │ ├── transfer_utils.py # bounded_att_buffer_copies：复制完整状态快照 +│ │ ├── conv_ssm_types.py # KDA/GDN 的 Conv/SSM 布局配置 +│ │ ├── create_kda_state.py # KDA 统一创建入口 +│ │ ├── kda_state.py # KDAState：Conv / SSM / MTP 与页操作 +│ │ ├── kda_pages_pool.py # KDAPagesPool：KDA Conv 历史和方阵递推快照 +│ │ ├── create_gdn_state.py # GDN 统一创建入口 +│ │ ├── gdn_state.py # GDNState：独立算法状态与页操作 +│ │ └── gdn_pages_pool.py # GDNPagesPool：支持不同 K/V head 数和维度的快照 +│ └── att_tail/ +│ ├── __init__.py +│ ├── base_state.py # AttentionTailState：尾部状态抽象操作契约 +│ ├── transfer_utils.py # att_tail_buffer_copies：本地 ring 与固定逻辑窗口之间搬运 +│ ├── create_kpool_state.py # K-pool 统一创建入口 +│ └── kpool_state.py # KPoolAttentionTailState：原生 ring 与恢复规则 +└── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 + +lightllm/models// +├── model.py # 选择、装配本模型的资源管理器 +├── att_resource_config.py # 模型维度与资源组合配置，供模型和缓存服务共用 +└── att_resource_manager.py # 模型资源组合、功能数据操作及具体资源绑定 +``` + +`allocators/` 与三类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 + +`full_att/base_state.py` 同时定义 `FullAttentionConfig` 和 `FullAttentionState`；`bounded_att/base_state.py`、`att_tail/base_state.py` 分别定义另外两类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 + +`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_pages`、`restore_bounded_att_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 + +`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；CPU meta 定义入口为 `create_bounded_attention_definition(...)`。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 + +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；其余复制视图函数返回源、目标视图列表，由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 + +CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttResourceManager.get_cpu_transfer_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 + +无需模型特有初始化的场景直接使用 `AttResourceManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `set_att_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 + +模型特有的 State 创建与存储绑定保留在 `lightllm/models//att_resource_manager.py`。Qwen3Next 创建 `main_gdn_state`；Glm5Next 创建 `main_kda_state` 和 `main_indexer_kpool_c4_tail_state`，压缩比例标记随配置变化。main/draft 的 KV State 按各自方案命名，draft 自行添加需要的 bounded/indexer State。模型管理器继承基类的数据操作和传输布局实现，仅在确有差异时覆盖，不重复实现相同的聚合转发。 + +token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttResourceManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttResourceManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 resource_id 匹配时可构建临时映射，它不决定模型的资源组织方式。 + +#### 按逻辑算法创建 State + +每种逻辑算法提供一个公开的 `create` 入口。量化属于该算法的存储选择，不按量化类型拆分创建入口；调用方提供逻辑配置及编码参数，由入口返回具体 State。 + +| 逻辑类型 | 创建入口 | 参数决定的实现 | +| --- | --- | --- | +| GQA（含 MHA/MQA） | `full_att.create_gqa_state(config, *, kv_type, page_size, device, ...)` | 普通浮点、INT4/INT8 分组量化、per-tensor/per-head 静态 FP8 | +| MLA | `full_att.create_mla_state(config, *, kv_type, packed, page_size, device, ...)` | 普通 MLA、带 indexer 的 packed MLA、packed FP8 MLA | +| KDA | `bounded_att.create_kda_state(config, request_capacity, mtp_step=0, *, device)` | KDA Conv/SSM 运行态 | +| GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | +| K-pool | `att_tail.create_kpool_state(config, request_capacity, mtp_step=0, *, device)` | 带池化尾部窗口及 MTP 历史槽位的 K-pool State | + +每个算法的创建入口独立放在 `create_gqa_state.py`、`create_mla_state.py`、`create_kda_state.py`、`create_gdn_state.py` 和 `create_kpool_state.py`，通过各分类包导出。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `set_att_state`。 + +GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 由调用方提供真实的 `scales` 和 `calibration_id`，per-head 还需 `q_scales`。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 + +`lightllm/utils/att_resource_config.py` 负责把启动参数和标定文件转换为这些显式参数，并调用对应算法入口。full-attention 创建时仍为零普通 token 容量，后续由管理器统一分配；bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 + +```python +state = create_gqa_state( + config=full_config, kv_type="int4kv", group_size=8, + page_size=16, device="meta", +) +manager.set_att_state("main", state) +# 全部 State 准备完成后，由 manager 统一规划容量并分配存储。 +``` + +根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttResourceManager.initialize_resources()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 + +管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过具体管理器操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 + +`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisSlice`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 + +## 5. allocator 与状态页 pool + +### 5.1 allocator 的公共入口与按需初始化 + +| 公共入口 | 类型 | 初始化条件 | +| --- | --- | --- | +| `full_att_allocator` | `FullAttKVAllocator \| None` | 组合中存在需要页式 KV 分配的 full-attention 资源 | +| `bounded_att_big_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 大页 | +| `bounded_att_small_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 小页 | + +具体管理器初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 + +allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由模型管理器明确描述。 + +大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 + +`FullAttKVAllocator` 的申请量和返回索引均以 token 表示，分配和回收须满足整页约束；`StatePageAllocator` 返回状态页槽位 ID。GPU 运行态和 indexer tail 使用请求槽位，不使用上述状态页分配器。 + +HOLD 页位于普通 KV 分配域之后，由 full-attention 存储预留。初始化流程将 allocator 的 `hold_token_indices` 显式交给 `ReqManager.bind_hold_page` 建立 padding 请求映射；ReqManager 不持有 allocator，也不负责释放 KV。`FullAttKVAllocator` 构造时直接创建 `shared_can_use_token_num: SharedInt` 并发布初始容量，在 `alloc`、`free` 和 `free_all` 中直接更新。对应的 `KVCapacityReader` 定义在同一文件中，router 按本地 DP 分组读取共享计数；Reader 不重置计数。读写双方使用同一服务前缀与 `rank_in_node` 命名，因此构造 allocator 前须完成服务名和 rank 初始化，无需外部绑定或容量回调。 + +`FullAttKVAllocator` 沿用 token 索引栈 `mem_state`、`mark_start/mark_end` 游标和循环返回缓冲区；每次分配复制索引到预分配的返回缓冲区，避免逐次申请 pinned memory。分配和归还均以整页为单位，归还时保持每页内部 token 索引连续且有序，页之间可以任意排列。调用方负责资源所有权、单次归还及索引的类型、范围与页内顺序；allocator 只保留容量与页对齐、int32 索引上限及分配/归还数量边界的断言，不逐元素校验、不另建逐页分配位图，也不对归还索引排序。返回缓冲区按原有三倍容量循环复用，调用方须在缓冲区复用前完成异步索引读取；需要长期保留索引时自行复制。 + +### 5.2 管理器、allocator 与 pool 分工 + +CPU 大小页池通过 `bounded_att/utils.py` 的 `create_page_pool_tensor()` 统一选择分配方式。未启用 `enable_dp_prompt_cache_fetch` 时沿用 `torch.empty(..., pin_memory=True)`；启用后直接创建共享 Storage 并原地调用 `cudaHostRegister(Portable | Mapped)`，避免先申请普通 pinned 页池再转共享存储造成整块复制。显式关闭 pin_memory 时只创建共享 CPU 存储，GPU 存储保持原分配方式。注册随本地 Storage 生命周期管理，最后一个视图释放时等待 GPU 使用结束再注销；接收进程需要自行注册其映射地址。KDA/GDN 页池均提供 `register_pin_memory()`，复用 `register_page_pool_tensor()` 注册 Conv/SSM 的整个 Storage；已 pinned 或空存储重复调用为空操作。`create_dp_ipc_copy()` 保留共享 CPU 页池，`create_pd_ipc_copy()` 始终排除全部大小页池。DP 接收端设置本地设备并读出资源后，调用 manager 的 `register_page_pools_pin_memory()`，统一注册所有 bounded State 的大小页池。`Portable` 使本进程其他 CUDA context 也识别其 pinned 属性，`Mapped` 建立设备可访问映射；它们不会替另一个进程注册地址。 + +`StatePageAllocator` 只维护可用槽位。`KDAPagesPool`、`GDNPagesPool` 等具体页池只负责分配和持有各自的状态 buffer；State 使用外部传入的槽位完成保存和恢复。`AttResourceManager` 持有 allocator 并提供跨 State 的聚合数据操作，调用方直接使用 allocator。操作顺序如下： + +1. 调用方直接从 `bounded_att_big_page_allocator` 或 `bounded_att_small_page_allocator` 分配槽位。新分配的槽位不代表已有可复用内容，必须先完成保存或载入才能发布。 +2. `save_bounded_att_pages` 聚合调用参与 State 保存；异步操作通过 stream/event 保证后续读取顺序。保存失败直接向上传递异常，由调用方停止发布并处理槽位生命周期。 +3. `restore_bounded_att_pages` 按调用方选定的槽位恢复全部参与 State，调用方保证前缀关联正确并已满足同步条件。 +4. 调用方确认引用已解除且异步使用完成，直接向同一个 allocator 归还共享槽位一次；pool 中的数据可保留至下一次写入。 + +CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排除在大页 allocator 的普通分配范围外；参与操作的 main/draft 大页 pool 分别为这些槽位提供存储。成员名为 `bounded_att_cpu_load_big_page_slot`（从 CPU 加载末尾不足一页的状态时使用）和 `bounded_att_cpu_offload_big_page_slot`（为卸载到 CPU 预留，当前尚无调用点），均属于 bounded-attention 大页池。 + +具体管理器声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 + +### 5.3 前缀长度与状态页的关联 + +前缀长度与大小页槽位的关联由请求及缓存流程维护。`big_pages_pool` 和 `small_pages_pool` 仅保存算法快照，不重复持有长度数组或页有效性标记。 + +大小页保存/恢复和 bounded 运行态复制只接收索引与布局，不接收前缀长度。`reset_attention_tail_state(req_indices)` 只清零指定请求的 tail，manager 的 CPU/DP 操作不传递专用于重置的长度。PD tail 读写单独接收 `state_token_len`，用于把 prefill 末尾的逻辑位置映射到两端各自的 ring；该值来自 PD 任务的 `end_kv_index`，不保存到 State 或传输页内。 + +调用方只能发布已完成全部参与 State 保存的页面，并通过 stream/event 保证异步操作顺序。尚未捕获的中间边界不能发布为可复用快照；失败时不发布对应页面，等待已提交操作结束后再回收槽位。前缀身份、引用与就绪状态均由调用流程管理。 + +## 6. 对外接口与 attention 资源暴露方式 + +模型专属管理器负责资源初始化、内部组合及缓存和传输功能所需的数据操作。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，具体管理器组织内部 State 完成操作。 + +### 6.1 FullAttentionState：公共接口与不同 storage 的完整处理 + +本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttResourceManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 + +#### 继承与内部 storage 的关系 + +建议以 `FullAttentionState` 作为公共接口，由具体布局子类实现差异。子类内部的 storage 持有 buffer 和布局信息，operator 可以复用兼容的写入、量化及复制 kernel。布局选择集中在具体子类及其配套实现，公共调用方使用稳定接口。 + +```text +FullAttentionState +├── GQAFullAttentionState # MHA / GQA / MQA 共用的 KV 存储 +├── GroupQuantizedFullAttentionState # INT4 / INT8 等 +├── StaticFP8FullAttentionState +├── MLAFullAttentionState +└── PackedMLAFullAttentionState # 附带 indexer 数据的 packed 布局 +``` + +所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。`FullAttentionConfig` 与抽象接口共同定义在 `full_att/base_state.py`，量化和 packed 专属配置可扩展该配置；配置结构的继承不代表 State 实现的继承。 + +`gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 + +该 State 只保存 KV，`config.head_num` 表示当前 TP rank 切分或复制后持有的 KV 头数。K 和 V 各自独立分配一个 `[layer, token_slot, local_kv_heads, head_dim]` buffer，不再保存 `heads_per_kv`。写入接口的输入仍按 K heads、V heads 前后排列，State 将两部分分别写入各自存储；计算 backend 直接取得 K/V 视图。Query 头到 KV 头的分组关系由 attention backend 处理，因此三者共用同一套 State 接口与存储实现。 + +#### 必须提供的公共接口 + +```python +# ViewT / WriteT 为具体布局定义的类型，下面是公共签名草案。 +class FullAttentionState: + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + @abstractmethod + def get_state_type(self) -> str: ... + + def __init__( + self, layout_config: FullAttentionConfig, + page_size: int = 1, *, + device: str | torch.device = "cuda", + ) -> None: + ... + + def initialize_storage(self, token_capacity: int) -> None: + ... + + def get_token_bytes(self) -> int: + ... + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + def get_layer_view(self, layer_id: int) -> ViewT: + ... + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: WriteT) -> None: + ... + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + ... + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + raise NotImplementedError + + # layout 描述该 State 在外部页中的位置及本地轴到页内轴的映射。 + def write_to_cpu( + self, token_indices, cpu_page_buffers, cpu_page_indices, + *, layout: StateTransferLayout, cpu_page_readies: List[bool] | None = None, + ) -> None: + raise NotImplementedError + + def read_from_cpu( + self, cpu_page_buffers, cpu_page_indices, token_indices, + *, layout: StateTransferLayout, + ) -> None: + raise NotImplementedError + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + # 其他 DP 的同一 TP 分片到当前 State；调用方准备各自设备上的 CUDA 索引，目标为 int64。 + def copy_tokens_from_other_dp( + self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, + remote_dp_full_att_state: FullAttentionState, + ) -> None: + raise NotImplementedError + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + raise NotImplementedError + + # pd_page_buffers 为页集合，pd_page_index 选择物理页；layout 区域偏移相对该页起点。 + def write_to_pd_buffer( + self, token_indices, pd_page_buffers, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + raise NotImplementedError + + def read_from_pd_buffer( + self, pd_page_buffers, token_indices, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + raise NotImplementedError +``` + +“必须提供”表示每个具体布局直接实现相同的抽象操作契约，不能从其他具体 State 继承实现。buffer 分配属于初始化；请求使用哪些物理 token 槽位由管理器的 allocator 决定。 + +| 接口 | 完整语义 | +| --- | --- | +| 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | +| `initialize_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | +| `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | +| `get_layer_view(layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | +| `write_full_att_kv(layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | +| `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | + +容量、device、层映射及布局标识应从同一份稳定配置和描述中取得，避免多处维护互相矛盾的副本。配置相同且只改变容量时，容量计算和真实 buffer 创建应覆盖相同的数据组成。共享底层内存按实际分配去重记账。 + +#### 存储字段与 State 元数据 + +三类 State 不提供完整存储枚举接口，各实现显式确定参与传输的固定字段。Full attention 的布局直接根据配置计算形状、dtype 和字节区间；PD/DP 调用处显式列出实际字段：GQA 为 K/V，MLA 为 latent，分组量化增加动态 scale，packed MLA 包含对应的 KV、scale、RoPE 与 indexer 字段。静态标定参数独立保留，不加入 token 传输字段。KDA/GDN 直接从配置描述 Conv/SSM，K-pool 直接描述完整 tail ring。 + +`TensorField` 仅作为内部搬运数据描述，仅保留 `tensor`、`axes`，字段名称由所在字典的 key 提供；Tensor 自身提供 dtype、shape、stride、device 和底层存储视图。运行态、MTP 游标及大小页池由具体 State 直接持有，不额外导出一份完整字段列表。 + +CPU/PD 的字段选择由 State 显式定义，公共区域接口如下： + +| State 类型 | CPU 区域接口 | PD 区域接口 | +| --- | --- | --- | +| `FullAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | +| `BoundedAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | +| `AttentionTailState` | 无；CPU 缓存恢复通过边界重置 | `get_pd_transfer_regions()` | + +所有接口返回 `tuple[TransferRegion, ...]`，可以基于 meta State 调用，不分配传输 buffer。State 决定哪些字段参与传输、字段区域的 dtype/shape、TP 切片和共享页的副本写入者；每份描述以 `byte_start` / `byte_end` 记录左闭右开的字节区间，两个位置均从该 State 的区域起点计算。三类 State 的 CPU/PD 区域接口统一无参数。full-attention 从全局启动参数读取 `cpu_cache_token_page_size` 或 `pd_kv_page_size`，按固定页容量展开形状与字段偏移；全局 head 数与本地 TP 切片均由自身 config 推导。bounded/indexer 每次描述一份运行态。 + +本次有效 token 数属于读写流程，不参与布局生成。CPU cache 使用固定的全局 CPU 页布局；State 的传输接口在不足整页时仍使用相同的字段形状和偏移，只搬运有效 token。缓存页的保存边界和发布条件由外层流程控制。 + +GQA、MLA 及各量化方案只导出 KV、动态 scale 和随 token 保存的 indexer 字段，不传静态标定参数或 packed padding。KDA/GDN 只导出有效 Conv 窗口和一份已接受 SSM，不传 MTP 候选区、游标或整个大小页池。K-pool 导出完整原生 key/gate ring。CPU 与 PD 可以在具体子类中分别覆盖；当前字段布局相同的实现共用自己的私有构建方法。 + +manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字段的实际最大结束位置安排后续资源，并通过布局终点计算整页大小；manager 不解析 Conv/SSM/K-pool 的内部结构。CPU 服务复用相同 State 接口；当前 CPU 联合页仍只包含 full/bounded，tail 按恢复边界重置。DP 映射只包含 full/bounded attention，调用它们无参数的 PD 区域接口，只取字段名和本地轴切片，不使用区域的页形状与字节偏移；本次复制的 token/request 数量由调用参数决定。 + +资源身份使用 `state.resource_id`；层映射、容量和 TP 分片事实继续由 State 的 config、capacity、page_size 及计算接口维护。CPU/DP/PD 由抽象接口约定并由具体 State 实现，管理器直接调用对应方法，不再维护额外的传输能力集合。信息保留在其所属对象，传输布局仅描述当前路径所需字段。 + +字段描述只读，引用的 Tensor 可以由合法计算操作更新。字段列表不表示数据已就绪，stream/event 同步和前缀关联仍由调用流程负责。每次增加辅助 buffer 或 packed 字段，都需要检查固定传输字段、容量计算、写入和复制/传输是否一起覆盖它。公共层不根据固定属性名推测资源组成。 + +#### 计算视图与完整存储范围 + +| 布局 | 计算视图的主要内容 | 完整复制需要覆盖的内容 | +| --- | --- | --- | +| 普通 K/V | K、V | 对应位置的全部 K/V | +| 分组 INT4/INT8 | K/V 与动态 scale | packed/量化 K/V 加上对应动态 scale | +| 静态 FP8 | 以正确数值类型解释的 KV 及静态 scale | 本地 token 复制只搬 KV；跨资源使用须验证或传递固定量化参数 | +| MLA | 压缩 KV 或其所需分段 | 本 State 负责的完整压缩 KV 表示 | +| 带 indexer 的 packed 布局 | attention KV 子视图、indexer key/scale 子视图 | KV 及全部相关 packed 辅助数据；按实际存储区域处理别名 | + +`get_layer_view()` 可以返回 `DenseKVView`、`QuantizedKVView`、`MLAView` 等具体小对象。视图类型由模型/backend 在初始化绑定时确认。各布局不必返回相同 Tensor 个数，也不必填充包含所有可选字段的大对象。 + +带 indexer 的 packed 布局中，attention 与 indexer 可以分别读取 buffer 的不同区域。因此不能用 attention 的读取视图代替完整复制范围。类似地，静态 FP8 的物理存储 dtype 和计算时解释的 dtype 可以不同，视图必须明确相应参数。 + +`write_full_att_kv` 更新其输入所描述的 KV 及关联字段；indexer 池化等专用 kernel 可以通过专用视图写入另一区域。管理器与执行流程在相关写入都满足同步条件后才允许复用或发布。融合写入路径可以直接使用预先绑定的布局视图，同时遵守相同的数据表示与完整性约定。 + +#### 基类与子类的分工 + +| 部分 | 基类可复用逻辑 | 子类必须负责的内容 | +| --- | --- | --- | +| 初始化 | 生命周期、配置与映射校验、稳定视图管理 | `_create_storage(...)`：实际 Tensor/区域和完整描述 | +| 层访问 | 逻辑层到存储层的映射、视图缓存 | `_make_layer_view(storage_layer_id)`：该布局的计算视图 | +| 写入 | 参数契约、层映射和公共调用入口 | `_write_full_att_kv(...)`：编码/量化/写入 kernel 及输入格式校验 | +| 完整复制 | 索引操作契约、公共调用入口 | `_copy_full_att_tokens(...)`：布局正确的复制，或显式复用与其描述兼容的公共工具 | +| 传输 | 定义 CPU/DP/PD 直接读写接口，未支持的操作显式报错 | 覆盖所支持的读写方法，直接使用自身 storage 调用相应 kernel | + +受保护方法名为候选。一个子类可以复用满足其布局要求的 kernel；通用复制工具只有在理解 token 维、stride、数据区域和别名关系时才适用。复杂 packed 布局应由专用实现处理，公共层不尝试仅凭 Tensor shape 推导复制语义。 + +#### 索引、异步执行与生命周期 + +构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `initialize_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `initialize_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `initialize_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式和静态标定参数的专属检查在对应 State 构造时完成。 + +State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 显式接收必需的 `scales` 和可选的 `q_scales`；meta 阶段保留真实标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 + +`layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 + +混合模型将实际 full-attention 逻辑层号及追加的 draft 层号写入 State 配置；例如逻辑层 3、7、8 可以分别对应物理行 0、1、2。计算侧直接以逻辑层号调用 `write_full_att_kv` 和 `get_layer_view`，由 State 完成映射，CPU/PD 描述复用同一份层标识。KV/Q 标定按物理存储行组织，backend 必须使用对应行的参数。 + +写入与复制接受与资源同设备的一维整数索引 Tensor。源和目标索引长度相同，目标唯一；基础复制先约定源目标不重叠，重叠移动另行定义。索引还须满足具体布局的对齐和粒度条件。 + +数据操作在调用方当前 CUDA stream 上提交，方法返回不代表设备操作完成；跨 stream 消费由调用方建立依赖。销毁存储前，管理器等待相关计算和传输结束并解除推理视图及跨进程存储引用。热路径不支持隐式 resize；需要重建存储时必须重新绑定视图及相关执行结构。 + +#### CPU / DP / PD 直接操作 + +三种路径直接作为 `FullAttentionState` 的方法暴露，由子类覆盖实现。CPU/PD 接口通过 `StateTransferLayout` 描述联合页与 TP 切片；DP State 接口直接复制同一 TP 分片的完整 token，不需要字段映射。布局是只读数据参数，不包含算子、buffer 所有权或执行方法。 + +| 接口 | 数据方向与位置参数 | +| --- | --- | +| `write_to_cpu(..., layout)` | 当前 State → CPU 缓存；layout 描述外部页区域、层/轴映射和写入副本选择 | +| `read_from_cpu(..., layout)` | CPU 缓存 → 当前 State；layout 描述该目标 State 从页中读取的对应范围 | +| `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` | 其他 DP 同一 TP 分片 → 当前 State，复制完整 token 的 KV、动态量化参数及 indexer 字段，不接收 mapping | +| `write_to_pd_buffer(..., layout)` | 当前 State → `pd_page_index` 指定的 PD 页；按 layout 写入指定区域 | +| `read_from_pd_buffer(..., layout)` | `pd_page_index` 指定的 PD 页 → 当前 State；按 layout 恢复对应区域 | + +`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，具体管理器按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 + +基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttResourceManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 + +每次调用完整处理其映射涉及的数据及关联的动态 scale、packed 字段；一次组合传输可以由若干次调用完成，上层检查所有必需数据区域均已完成。固定参数按约定传递或验证兼容性。位置映射的组成、单位及生成责任见 6.5 节。 + +方法统一在调用方当前 CUDA stream 上提交操作，返回 `None`。若内部使用其他 stream，子类须将完成依赖汇合到调用方 stream，使调用方随后记录的 event 覆盖本次操作。网络收发、任务排队和完成通知由调度层负责。 + +形状相同不能证明不同模型的 KV 语义可互换。模型/资源身份和 token 前缀匹配由管理器与缓存流程校验；State 检查布局及固定参数兼容性。跨布局重编码需要明确实现，本地 `copy_full_att_tokens` 不承担这一职责。 + +### 6.2 AttResourceManager：内部资源的聚合操作与数据传输 + +`AttResourceManager` 对内部资源组合提供统一的数据操作。一个模型可以包含不同角色或不同具体方案的 FullAttentionState、KDAState、GDNState 等状态实现，同一角色和具体方案最多一个；token 复制覆盖 full-attention 参与集合，清零操作覆盖需要初始运行态的参与集合，状态页操作覆盖指定大小页域的参与集合。外部调用方无需知道具体数量、类型或成员名称。 + +数据操作的依赖方向是 `radix cache / CPU cache / 其他调用方 → AttResourceManager → 各 State`；槽位分配与归还由调用方直接访问管理器持有的共享 allocator。状态页接口不接收 radix 节点、树或前缀匹配对象，也不读取 radix 内部状态；radix 只保存并使用管理器提供的页类别、槽位 ID 和状态长度等数据。状态页在未启用 radix 时同样可以保存、恢复或用于 CPU 缓存。 + +#### 资源初始化、State 查询与传输布局查询 + +| 接口 | 语义 | +| --- | --- | +| `AttResourceManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | +| `set_att_state(model_name, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | +| `finish_state_creation()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_states_created = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | +| `initialize_resources(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：标记 State 创建完成，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | +| `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | +| `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | +| `get_attention_tail_states(role=None)` | 返回指定角色或全部 `AttentionTailState` 的 tuple | +| `get_full_att_token_bytes()` | 返回当前 TP rank 全部 full-attention State 的单 token 存储字节数，包含动态 scale 和行 padding | +| `get_cpu_page_bytes()` | 依赖 `get_cpu_transfer_layouts()`，取全部 State 的最大 `byte_end` 并按 16 字节对齐；包含字段间及 State 间的必要 padding，页内不保存格式指纹、长度或版本 | +| `get_cpu_transfer_layouts()` | 按全局 `cpu_cache_token_page_size` 生成 CPU 联合页中各 State 的字段偏移和 TP 映射，使用 `lru_cache` 缓存 | +| `get_pd_full_att_transfer_layouts()` | 从全局启动参数读取 `pd_kv_page_size`，按 State 字段的实际偏移和结束位置生成 `kv` 布局及最终页跨度；不依赖 `get_pd_page_bytes()`，按 `resource_id` 索引并缓存 | +| `get_pd_bounded_att_or_tail_transfer_layouts()` | 按实际字段结束位置独立生成 `att_state` 布局及其自身所需大小，不查询 KV 布局或执行组合容量校验；每个 State 保存一份运行态，`tokens_per_page=None` | +| `get_pd_page_bytes()` | 查询两类布局大小并集中决定实际 PD 块容量：返回 KV 页与完整运行态页占用的最大值；有 KV 且 KV 页较小时记录 warning 并扩大物理页；无资源时返回 0。包含 field 与 State 间的 padding，可基于 meta State 查询，不分配传输 buffer | +| `initialize_storage(token_capacity, page_size, *, big_page_tokens, small_page_capacity, device)` | 为已创建 State 分配实际 full-attention 存储、共享 allocator 和 bounded 大小页 | +| `create_pd_ipc_copy()` | PD 导出：保留运行态 Tensor 和容量信息，排除 allocator 及所有大小页池 | +| `create_dp_ipc_copy()` | DP shared 导出：保留运行态和共享 CPU 大小页池，仅排除 allocator | +| `register_page_pools_pin_memory()` | 在当前 GPU 上统一注册全部 bounded State 的大小页池，已注册存储可重复调用 | + +聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `set_att_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `resource_id` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数，也不提供单对象选择接口。模型推理直接访问具体 State 成员。 + +State 装配期间由管理器内部遍历 `__dict__`，公共查询会断言失败，避免缓存不完整的 main/draft 集合。装配完成后先调用 `finish_state_creation()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 + +容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `initialize_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded/indexer 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 + +实际存储分配完成后执行 warmup。CPU 布局由第一次容量查询或读写时按需生成；PD 容量规划会通过 `get_pd_page_bytes()` 在 meta State 上提前生成 PD 布局，不分配实际传输 buffer。PD 页的 token 容量由全局 `get_env_start_args().pd_kv_page_size` 唯一确定，两个 PD 布局 getter 均不接收页 token 数，按 manager 实例缓存并直接返回 `Mapping[str, StateTransferLayout]`，没有对应资源时返回空字典。 + +PD 的依赖方向为：`get_pd_page_bytes()` 分别查询两个独立的布局 getter。full-attention 和 bounded/indexer 各自生成偏移及本类布局所需大小，互不查询，也不调用 `get_pd_page_bytes()`。实际块大小取两类布局占用的最大值，统一在 `get_pd_page_bytes()` 完成。State 提供包含内部字段 padding 的相对偏移，manager 按字段的实际最大结束位置确定下一个 State 的起点，不再用裸存储字节数预留 PD 区域。PD 收发流程按 `kv` / `att_state` 任务选择布局。 + +CPU getter 从全局启动参数读取 `cpu_cache_token_page_size`，同样按 manager 实例缓存。不保存独立的 `cpu_page_bytes`、`pd_page_bytes`、`cpu_transfer_layouts`、`pd_transfer_layouts` 或 `cpu_page_tokens` 成员。存储分配可能改变 State 的物理页规格，因此分配完成后同时清除布局及派生 CPU/PD 页大小缓存。初始化时的异常直接向上传播，初始化失败即终止该流程；不增加属性拦截或回滚逻辑。 + +查询检查与遍历集中在 `_get_states`，通过 `@lru_cache(maxsize=None)` 缓存结果，键包含 manager 实例、资源类别和角色。返回不可变 tuple 中的对象引用，不复制存储，也不修改资源归属，不再维护实例缓存字典。manager 按进程生命周期使用；若在进程内卸载或重建管理器，需对 State 查询、页大小和布局查询方法调用 `cache_clear()`，释放缓存持有的 manager 和 State 引用，该操作会清除对应方法所有实例的查询缓存。 + +IPC 发布按用途调用 `create_pd_ipc_copy()` 或 `create_dp_ipc_copy()`，两者均浅拷贝管理器并将三个 allocator 字段置为 `None`。PD 导出会分别浅拷贝 bounded State，再清空大小页池；即使源页池是共享内存，PD 副本也不携带。DP shared 导出保留运行态及大小页池引用，要求页池已使用共享 CPU 存储，避免序列化私有页池时隐式分配、复制。两种导出都不修改源管理器，也不转移槽位所有权。 + +`shared_att_resources_utils.py` 提供 `publish_att_resources` / `load_att_resources`，用于发布和读取 attention 资源及请求索引映射的 IPC 副本。PD/DP 的发布与读取显式传入 `kind`，分别使用 `rank_storage_pd_` 和 `rank_storage_dp_`，锁及剩余读取计数也各自独立，允许同时启用且不依赖读者到达顺序。DP 仅导入本节点其他 DP 中对应 TP rank 的副本，每个发布者的读者数为 `dp_size_in_node - 1`；资源管理器和请求映射使用按节点内 DP rank 索引的字典，每项均对应本进程的 TP 分片；节点内进程 rank 仅用于加载 IPC 资源。DP 接收端设置当前设备后读取 DP 副本，再调用 `register_page_pools_pin_memory()`；manager 遍历各 State 的大小页池，复用 pool 的 `register_pin_memory()`。注册只作用于接收进程的映射，不随对象序列化；PD 接收端无需注册页池。 + +`lru_cache` 属于方法，不进入实例的序列化数据；副本及接收方首次查询时以各自实例为键建立缓存。是否允许查询由创建完成标记判断，不依赖 allocator 是否存在。 + +#### 数据操作入口 + +| 操作 | 接口 | 聚合语义 | +| --- | --- | --- | +| full-attention token 复制 | `copy_full_att_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | +| bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页 | +| bounded 状态页保存 | `save_bounded_att_pages(...)` | 将指定运行态保存到全部参与 bounded State 的同类 pool、同一组槽位 | +| bounded 状态页恢复 | `restore_bounded_att_pages(...)` | 按调用方选定的槽位，将各 bounded State 保存的页恢复到对应运行态位置 | +| bounded 验证后提交 | `commit_bounded_att_state(req_indices, accepted_steps, *, role)` | 将指定角色的接受步数分发给参与 bounded State；由 State 选择已接受运行态，不暴露游标 Tensor | +| attention tail 初始化与无快照恢复 | `reset_attention_tail_state(req_indices)` | 清零指定请求在全部参与 tail 中的状态；调用方保证空前缀或可重置的缓存边界，保留其他状态 | +| CPU 写出 / 读回 | `write_req_to_cpu(...)` / `read_req_from_cpu(...)` | 组合 KV 和状态页，按调用参数提供的长度读写各 State 在外部页中的区域 | +| DP 复制 | `copy_from_other_dp(remote_att_manager, src_token_indexes, dst_token_indexes, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，清零目标 tail；不分配资源 | +| PD 写出 / 读回 | `write_to_pd_buffer(...)` / `read_from_pd_buffer(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | + +#### 必须实现的聚合接口 + +基类源码使用醒目的分隔线划分八个大块：基础管理、Full attention、Bounded attention、Indexer tail、CPU Cache 数据交换、DP 迁移、PD 分离传输、私有辅助实现。Bounded attention 分块依次放置运行态重置、大小页保存与恢复、MTP 提交接口。各块内保留接口的具体说明，接口顺序保持 public 在前、private 在后。 + +管理器的分类操作使用 `full_att`、`bounded_att`、`att_tail` 标明作用范围；CPU/DP/PD 接口继续表示跨资源传输。具体 State 的分类操作与管理器使用相同命名；State 只操作自身数据，管理器负责聚合参与对象。 + +以下仅展示接口签名，实际方法在基类中提供聚合实现。 + +```python +from collections.abc import Mapping +from typing import Any, List, Literal + +PageKind = Literal["big", "small"] + + +class AttResourceManager: + full_att_allocator: FullAttKVAllocator | None + bounded_att_big_page_allocator: StatePageAllocator | None + bounded_att_small_page_allocator: StatePageAllocator | None + + def __init__(self, request_capacity: int): ... + + def set_att_state( + self, model_name: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState, + ) -> None: ... + + # ======================================================================== + # Full attention：token 数据复制 + # ======================================================================== + def copy_full_att_tokens( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, + ) -> None: ... + + # ======================================================================== + # Bounded attention：运行态重置、大小页保存与恢复、MTP 提交 + # ======================================================================== + + # 空前缀初始化，只处理 bounded 运行态；req_indices 定位已分配的槽位。 + def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... + + # Bounded Attention 状态页：供 radix/CPU cache 复用，调用方直接分配和归还槽位。 + def save_bounded_att_pages( + self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + ) -> None: ... + + def restore_bounded_att_pages( + self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + ) -> None: ... + + # MTP 推测解码：只提交指定角色；接受步数包含本轮已经执行的基准输入步。 + def commit_bounded_att_state( + self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, *, role: str, + ) -> None: ... + + # ======================================================================== + # Indexer tail：空前缀初始化与无快照恢复 + # ======================================================================== + def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... + + # 基类使用 lru_cache 缓存布局，存储初始化后清除缓存。 + def get_cpu_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + + # CPU 联合页：KV + 全部参与 linear State 的有效页。 + # 纯 full-attention 模型没有 linear 页；混合模型不能省略它们。 + def write_req_to_cpu( + self, token_indices, cpu_page_buffers, cpu_page_indices, + *, layouts: Mapping[str, StateTransferLayout], + cpu_page_readies: List[bool] | None = None, + big_page_slot_ids: List[int] | None = None, + small_page_slot_ids: List[int] | None = None, + ) -> None: ... + + # 一次读取同一前缀的有序页面；全部页面读回后清零 req_idx 的 tail。 + def read_req_from_cpu( + self, cpu_page_buffers, cpu_page_indices, token_indices, + *, layouts: Mapping[str, StateTransferLayout], + big_page_slot_ids: List[int] | None = None, + req_idx: int | None = None, + ) -> None: ... + + # DP 缓存迁移：索引已由调用方分配，空 Tensor/列表表示不搬运该类数据。 + def copy_from_other_dp( + self, remote_att_manager: AttResourceManager, + src_token_indexes: torch.Tensor, dst_token_indexes: torch.Tensor, + src_big_page_ids: List[int], dst_big_page_ids: List[int], + src_req_idx: int, dst_req_idx: int, + ) -> None: ... + + # 基类直接构建两类 PD 布局，各自使用 lru_cache 缓存；无对应资源时返回空字典。 + def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + + def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + + # PD 分离：pd_page_buffers 为页集合，pd_page_index 必填；kind 沿用 kv / att_state 两类任务。 + def write_to_pd_buffer( + self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, + ) -> None: ... + + def read_from_pd_buffer( + self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, + ) -> None: ... +``` + +以上方法均有基类默认实现，直接按已创建的 State 聚合执行。子类仅实现模型特有部分。未支持的 CPU/DP/PD 路径须明确报错，启用相应功能时检查参与 State 的完整支持能力。未启用某种页分配域时，保存和恢复该域明确报错；纯 full-attention 模型无需创建大小页 pool，其 `reset_bounded_att_state` 可以为空操作；没有 indexer tail 的组合，其 `reset_attention_tail_state` 也可以为空操作。 + +`page_kind="big"` 对应 `bounded_att_big_page_allocator` 和参与 State 的 `big_pages_pool`；`"small"` 对应小页 allocator 与 `small_pages_pool`。两者是独立分配域，即使槽位数字相同，也不是同一份资源。 + +调用方根据 `page_kind` 直接选择相应 allocator；该成员为 None 时不能申请该域。分配与归还使用 allocator 自身的接口，管理器不再增加包装方法。allocator 只管理空闲索引，不修改 State 的页数据或有效长度。 + +`req_indices` 只定位运行态存储，不传递请求对象。保存/恢复时，请求状态槽位与页槽位逐项对应；目标槽位唯一且已分配。调用方维护快照与实际前缀的关联，资源方法不再接收重复的长度参数。 + +#### bounded 空前缀运行态初始化 + +`reset_bounded_att_state(req_indices: List[int])` 用于从前缀长度 0 开始计算时初始化 bounded 运行态。`req_indices` 仅接受 Python 整数列表，不支持 `torch.Tensor`；实现直接使用列表，不做类型转换。调用方已通过 ReqManager 取得请求槽位，并确认这些槽位此前的计算和传输使用结束；管理器只按传入索引初始化 bounded 数据，不查询缓存命中情况或请求对象。 + +- 对全部参与的 KDA/GDN 等状态调用 `reset_bounded_att_state(req_indices)`，将 Conv/SSM 恢复到零初值，并重置 MTP 游标及影响后续读取的候选信息。未来算法按自身契约恢复空前缀初态。 +- 不修改 AttentionTailState。没有缓存命中时，调用方另行调用 `reset_attention_tail_state(req_indices)`，完成 tail 的空前缀初始化。 +- full-attention 没有需要清零的空前缀递推状态；该操作不遍历或清空历史 KV，不更改请求映射，也不失效大小页快照或归还索引。纯 full-attention 管理器合法地不做数据操作。 +- 目标请求槽位须唯一且在有效范围内，main/draft 所需运行态全部初始化后才可使用。重复执行在没有并发访问时具有相同结果；方法返回仍遵守 stream 完成约定。 + +这是长度 0 的数据初始化接口，不接收非零前缀长度。命中非零前缀时使用 `restore_bounded_att_pages` 或 CPU/DP/PD 读回；没有 tail 快照时在保证可重置边界后通过 `reset_attention_tail_state` 清零 tail；不能在 bounded 运行态恢复后调用此方法，否则会清除已经恢复的 bounded 数据。 + +#### 验证后运行态提交 + +`commit_bounded_att_state(req_indices, accepted_steps, *, role)` 只执行 bounded 运行态选择，不承担 token 验证、采样、KV 回收或请求长度更新。调用方在验证完成后传入逐请求的槽位与接受步数，两个参数必须是当前 State 设备上的一维 CUDA Tensor；请求索引为 int32/int64，接受步数为 int32，长度一致，接口内部不转换设备或类型。语义与 State 的 `commit_bounded_att_state` 接口一致；接受步数包含已经执行的基准输入步。 + +main/draft 共享管理器，因此必须显式指定 `role="main"` 或具体 draft 编号（例如 `role="draft1"`）。管理器只调用该角色参与 bounded State 的 `commit_bounded_att_state`，不将主模型验证结果直接提交到 draft 的候选状态；该角色没有需要提交的 State 时为空操作。 + +`mtp_state_index` 及其他提交控制信息只由具体 State 持有。BaseModel 不维护 `mtp_state_cursors` 或其他游标引用集合，验证流程也不直接读写游标。State 自行决定通过游标选择、数据搬移或其他方式提交；计算 backend 可通过 State 或层视图读取所需信息。 + +#### 无 tail 快照时清零尾部状态 + +`reset_attention_tail_state(req_indices)` 清零指定请求的 tail，用于空前缀初始化或可重置的缓存前缀恢复。请求槽位使用 `list[int]`，目标唯一且已分配；接口不接收前缀长度，不读取请求映射或缓存节点。 + +管理器依次调用各 AttentionTailState 的 `reset_attention_tail_state(req_indices)`。K-pool 直接清零指定请求的完整 ring，不在 reset 中重复检查长度。模型初始化与缓存流程保证恢复边界对齐：缓存页大小必须是各参与 tail 池化步长的整数倍，main/draft 使用不同方案时分别满足各自约束。异常直接透传，不保证跨 State 或整个 CPU/DP 恢复流程的原子性。 + +该操作只重置 tail 及其必要控制信息，保留已恢复的 Conv/SSM、MTP 游标、KV 和大小页快照。没有 tail 的管理器合法地不做数据操作。调用方保证此前使用已结束，并在相关数据操作完成后继续计算；这不是运行中 MTP 回退的替代操作。 + +前缀页命中时，调用方恢复 KV 和 bounded 页后，调用本接口清零未保存的 tail。PD 已传入完整 tail 时直接使用该快照，不再重置。DP fetch 只迁移可恢复的缓存前缀，由管理器在目标端清零 tail；边界选择属于缓存流程。 + +这里的“最后一页”按本次命中前缀的逻辑顺序确定，与 CPU 物理页编号大小无关。全部页面先恢复 KV 和 linear 页数据，再由末页恢复 linear 运行态；中间长度无需逐个重置 tail，只有最终恢复长度需要满足 tail 的可重置边界。批量恢复多个请求时，调用方按请求分别调用，各自选择对应前缀的末页。 + +#### full-attention 的聚合 token 复制 + +`copy_full_att_tokens(src_indices, dst_indices)` 的两个索引均为存储设备上的一维 CUDA 整数 Tensor；直接使用输入 dtype 和 stride，不转换设备、类型或连续布局。源和目标位置不重叠。该接口在本管理器持有的 full-attention 存储内部执行位置复制。参与集合在初始化时按共享 token 索引域确定，可以包含 main/draft 的不同布局，也可以包含同一模型的多个 FullAttentionState。管理器将同一组物理索引交给各参与 State 的 `copy_full_att_tokens`；每个 State 在自己的存储中从源位置复制到目标位置，不在不同模型之间互拷数据。 + +- 复制覆盖全部参与存储层的 KV、动态 scale 和 packed indexer key/scale 等随 token 变化的数据，保持各自存储表示；共享物理区域的别名按所有权去重。 +- 源和目标为逐项对应的一维整数索引 Tensor，长度相等、目标唯一，且满足各布局的复制粒度与对齐要求；基础契约要求源目标不重叠。源数据须已就绪并保持有效，目标物理位置须已分配且允许写入。 +- 只搬运 full-attention token 数据，不分配或回收索引、不修改 ReqManager 映射，也不复制 linear 运行态、大小页快照或 AttentionTailState 的原始 tail。 +- 调用前校验全部参与 State 的参数和能力；无 full-attention 参与者时明确报错。方法在当前 stream 提交并返回 None，全部参与复制完成前不能消费目标数据；中途失败时目标视为未就绪。 + +例如 main 和 draft 各有一个 FullAttentionState，复制物理位置 `[10, 11]` 到 `[20, 21]` 时，一次管理器调用会分别搬运 main 和 draft 各自的对应数据。共享的是位置变换，不要求两者具有相同存储布局。 + +#### 多个 State 的参与规则 + +具体管理器在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 + +例如，一个管理器同时管理两个 KDAState 和一个 GDNState，大页槽位 17 对应： + +```text +bounded_att_big_page_allocator 分配一次 slot_id = 17 + ├── KDAState A.big_pages_pool[17] + ├── KDAState B.big_pages_pool[17] + └── GDNState C.big_pages_pool[17] +``` + +调用 `save_bounded_att_pages(..., page_kind="big", slot_ids=[17])` 保存三个 State；`restore_bounded_att_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 + +调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 + +full-attention 历史 KV 和没有独立 pool 的 AttentionTailState 不属于这组状态页操作。需要同时恢复这些数据的调用方，使用相应资源操作，并保证历史 KV、bounded 快照和 tail 对应同一前缀且均已就绪。 + +#### 保存与直接归还槽位的完整性 + +- 调用方只从共享 allocator 分配一次槽位，全部参与 pool 覆盖同一索引范围。分配成功仅表示槽位可用，不表示其中已有有效状态。后续数据操作失败时，本次数据不得发布；调用方在已提交操作结束后决定重写或归还槽位。 +- `save_bounded_att_pages` 将同一组槽位传给全部参与 State 保存。只有所有参与者都完成对应位置的数据保存，调用方才能发布页面；一项失败时直接透传异常，不发布本次结果。 +- `restore_bounded_att_pages` 使用调用方选定的快照恢复各参与 State。复制阶段失败时直接透传异常，源页保持不变，调用方不得消费半恢复的数据。 +- 归还槽位前，调用方必须解除相关使用引用并满足完成条件，再直接向对应 allocator 归还一次。只归还指定大小页域的槽位，不操作相同数字的 KV 索引或请求槽位；槽位中的旧内容不能作为重新分配后的有效快照使用。 + +保存时全部参与 State 必须实际具备所要求位置的状态；不能从较新的状态补造较早快照。main/draft 可以分别计算，但聚合操作的完整性条件不变；若具体实现分阶段捕获参与者，未全部完成前必须保持聚合页不可用。 + +GPU 数据操作遵守当前 stream 的依赖约定。调用方通过 stream/event 在读取、发布或释放前确认异步复制完成。数据操作失败后，等待已提交操作结束再复用或释放相关存储。 + +共享域跨多个管理器实例时，调用方使用同一个 allocator 实例分配和归还。只有所有使用方均已结束且全部相关引用已解除后，才归还该槽位一次；不能仅因某个管理器已经结束使用就提前回收。allocator 的创建和销毁责任仍由装配时约定的唯一所有者承担。 + +#### CPU / DP / PD 聚合接口的参数与数据路径 + +`layouts` 是按完整 `resource_id`（包含 main/draft 角色）索引的 `StateTransferLayout` 映射，用于 CPU/PD 外部页。DP 不构造字段映射：manager 直接遍历本地 State，通过同名成员取得远端 State，将对应 token、页或请求索引交给具体实现。DP 不再维护映射缓存，也不通过 helper 转发。 + +CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。CPU 读写由管理器按 full-attention、bounded 状态页、末页运行态恢复及 tail 重置分阶段组织，高频读写路径不记录流程日志；buffer 与字段搬运约束交给 State，管理器检查恢复所需的目标请求、非空前缀及完整页槽位数量。异步搬运的完成仍由调用方的 stream/event 保证。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 + +**CPU 联合页** + +包含 linear attention 的模型，每个已发布 CPU 缓存页都必须包含全部参与 linear State 对应页终点的有效快照，且长度一致。调用方在写出前通过前缀与槽位的关联选取完整快照，读回时按布局写入目标 pool；不完整状态不能发布。纯 full-attention 模型仍使用自身的 KV 页格式。接口不增加缺失 linear 快照的 CPU 页语义，也无需为选择有效快照终点新增元数据查询方法。 + +- `token_indices: torch.Tensor` 是由调用方准备在 CUDA 上的 full-attention 物理 token 索引，`cpu_page_indices: List[int]` 是 CPU 页池的页 ID，仅接受 Python 列表，不接受 Tensor 或其他序列，也不在接口内部转换。索引按列表顺序排列，从各页槽位 0 开始；长度与页数遵守 6.5 节的整页/尾页约定。页列表按逻辑前缀顺序排列，不要求物理页 ID 连续或递增。只传状态的 CPU 格式使用空整数 Tensor 作为 token 索引。 +- 格式包含 linear 页时，管理器写出接口分别接收 `big_page_slot_ids` 与 `small_page_slot_ids`；`cpu_page_indices` 前段对应大页、后段对应小页，每段与槽位列表一一对应。未参与的页池传 None 或空列表。full-attention KV 按完整页面列表写出一次，各 bounded State 再分别从大小页池写出快照；读回接口固定写入大页池，不接收 `page_kind`。管理器的 `big_page_slot_ids: List[int]` 只包含完整 CPU 页对应的普通大页槽位，不包含末尾碎页的临时槽位；没有 linear 页时为 None。槽位参数仅接受 Python 列表，不接受 Tensor。快照搬运不需要逐页长度数组。 +- `cpu_page_readies: List[bool] | None` 是本批 CPU 页的跳过写入标记，顺序与 `cpu_page_indices` 一致；True 保留整张联合页，None 表示全部写出。缓存层保留完整 token/page 序列并一次调用资源管理器，不逐页发起传输，也不通过过滤页 ID 改变 token 映射。资源层不查询或发布 ready 状态，全部 ready 时直接返回。 +- FullAttentionState 将 ready 标记传给各自 CPU/PD 共用的传输入口；CPU 写出按 State 构造一次 GPU bool Tensor，全部字段复用。在 kernel 内按逻辑页屏蔽源 KV/量化字段读取和目标写入；若当前 block 没有待写元素，直接跳过后续索引读取、地址计算和数据搬运。跨页 block 混有 ready 与未就绪元素时保留逐元素 mask。PD 与读回不传标记、不分配此 Tensor，也不执行 ready 归约判断。bounded State 接收对应大小页段的 ready 列表，在取快照和构造复制视图前跳过 ready 页。 +- CPU 页不保存 tail 数据，写出接口不接收请求索引；布局只包含 full-attention 与 bounded-attention。 +- `read_req_from_cpu` 一次处理同一前缀的有序页面，将全部 KV 写到给定物理位置，将各页 linear 快照写入对应 pool 槽位。管理器根据 `token_indices.numel()` 和布局的页 token 数判断是否存在末尾碎页，将其完整 linear 快照写入预留 LOAD 大页；完整页写入 `big_page_slot_ids` 指定的大页槽位。读回后直接使用最后一页恢复 `req_idx` 的 linear 运行态，不修改调用方的槽位列表。仅含碎页时传 `big_page_slot_ids=[]`；只传状态的格式使用空 token 索引，按传入的普通槽位恢复。 +- 存在 bounded 或 tail 时，调用方提供单个目标 `req_idx`；管理器在全部页面读回后调用 `reset_attention_tail_state([req_idx])`。调用方保证最终前缀处于可重置边界。仅有 full-attention 时，`req_idx=None`。 +- 恢复终点必须对齐 tail 的可重置边界。最后一页可以仅使用大页的一部分，但仍须落在允许恢复的小页边界；不能恢复需要真实 tail 快照的非对齐位置。 +- 任一资源读入失败时异常直接透传，已写入的数据不得作为就绪数据使用。写出失败时不得发布不完整联合页。 + +例如恢复三个 CPU 页，终点长度依次为 `[P, 2P, 3P]`：全部 KV 和 linear 快照读回后，只对最终长度 `3P` 重置 tail，不逐页重置。管理器在重置 tail 前已将最后一页的 linear 快照恢复为运行态，调用方不再额外调用恢复接口。 + +**DP 直接复制** + +当前 DP fetch 的来源是 radix cache 命中的前缀。缓存页大小须为每个 tail 池化步长的整数倍；GLM 启动时检查 `linear_att_hash_page_size % index_kpool == 0`。管理器复制 KV 与 bounded 运行态后清零目标 tail，不额外接收前缀长度。tail 不参与 DP 映射，也不读取源 tail；PD 的结束位置可能未对齐，仍传输完整快照。 + +`src_manager` 明确来源 DP/TP 存储实例，也可以是传输进程中已绑定共享存储的管理器。源和目标 token 索引等长、逐项对应；`src_req_indices`、`dst_req_indices` 仅接受 `List[int]`，成对提供且数量相等，用来迁移 prefill 基准 bounded 运行态；KDA/GDN 源游标必须为 0。有 tail 时直接清零目标请求的 tail，不传递额外长度。token 索引与请求槽位是不同的寻址维度，无需等长；前缀对应关系及可重置边界由调用方保证。 + +只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，具体管理器将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 + +**PD 页读写** + +| `kind` | 必需参数 | 处理对象 | +| --- | --- | --- | +| `"kv"` | `token_indices`；`req_idx=None` | 按单页容量读写全部参与 full-attention 的 KV、scale 和 packed 字段 | +| `"att_state"` | `req_idx`；有 tail 时还需 `state_token_len`；`token_indices=None` | 读写 prefill 基准 bounded 状态与固定 tail 逻辑窗口，不经过大小页 pool | + +PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再仅转交 tail State 完成 ring 重排。full/bounded State 的接口不增加长度参数；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 + +CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 + +#### 调用方与资源实现的边界 + +radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存和恢复状态页时调用上述通用接口，前缀与槽位的对应关系由 radix 调用流程保证；槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 + +CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，具体管理器负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 + +### 6.3 模型推理使用具体资源视图 + +Model 和 InferStateInfo 通过 `att_resource_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_full_att_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_resource_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 + +视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与 tail 的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 + +普通 KV 层可支持多种存储精度。模型准备 State 时，将该层使用的规范成员名保存为 `layer.full_att_state_name`，关联的 indexer 同样记录该名称；推理通过 `getattr(att_resource_manager, self.full_att_state_name)` 直接访问成员。这里仅保存字符串，不持有 State 别名，不建立资源注册表，也不在 forward 中遍历候选 State。不同层可以记录不同方案的成员名，由具体模型负责确定对应关系。 + +Gemma4 共享 KV 层使用被共享源层所属 State 的成员名，并通过源层号取得视图。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按自身压缩比例访问对应角色的 tail 成员。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 + +成员访问只读取 Python 属性，不复制 buffer，也不执行 GPU→CPU 同步。KV 写入直接调用该 State 的方法，packed KV/indexer 通过其层视图访问。DFlash 的特殊 KV 写入上下文同样设置管理器和 draft 标记。需要主机标量的 backend 在 State 初始化时准备标定镜像，避免 forward 中逐层执行 GPU→CPU 同步。 + +### 6.4 调用边界与独立传输进程 + +单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttResourceManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 + +具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttResourceManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该模型管理器的具体实现提供。 + +跨进程共享通过管理器的 `create_pd_ipc_copy()` / `create_dp_ipc_copy()` 保留 State Tensor 引用并由 IPC 传递句柄，不调用 `initialize_storage()` 重新分配 KV；两者均排除 allocator，PD 副本排除全部大小页池，DP 副本保留共享 CPU 大小页池，不导出整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 + +### 6.5 复杂 attention 组合与 TP 的位置映射 + +#### 位置映射的范围 + +传输需要分别描述逻辑层到物理层的映射、组合页中的资源区域、各字段的 TP 分片与副本关系，以及本次操作的 token/page 索引。KV、scale、Conv、SSM 和 tail 可以使用不同的 dtype、stride、对齐和轴划分，不能统一简化为一个 `offset` 或一个全局 `tp_rank`。 + +#### 信息由谁维护 + +| 信息 | 维护位置 | 示例 | +| --- | --- | --- | +| 本地存储布局与分片事实 | 具体 State 的配置、存储参数及内部 Tensor | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | +| 传输字段及本地 TP 切片 | 具体 State 的 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` | 字段选择、全局形状、State 内字段的起止区间和各 rank 的读写切片 | +| 组合页布局及资源间分区 | 模型专属 `AttResourceManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_full_att_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | +| 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `StateTransferLayout`；按资源与 CPU/PD 路径分别保存 | +| 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | + +资源的全局身份须包括模型角色及资源标识，不能只用层号识别 main/draft。同一模型角色可以有多个存储或量化方案不同的 State，但每个 `(role, state_type)` 最多一个；同方案的所有层归入该 State。管理器按需组织成员，不增加固定的六个成员变量。共享 allocator 只保证索引域，不能推导层、head、字段或传输字节位置。 + +同一 State 可能服务 CPU cache 和多个 PD 对端，不应把组合页的绝对位置固化成唯一的 `self.pd_offset`。本地存储事实固定在 State，外部位置由该次传输选择的 layout/mapping 提供。算子不得用进程全局的 TP rank 替代被操作 State 的分片描述，尤其是在传输进程导入多个 rank 的 State 时。 + +#### StateTransferLayout：单个 State 的传输布局 + +`StateTransferLayout` 通过 `resource_id` 绑定一个具体 State，描述它在 CPU/PD 外部页中的传输布局。它通过 `regions` 组织字段级映射，每个 `TransferRegion` 描述 State 内一个字段与外部页区域之间的对应关系，`AxisSlice` 进一步描述轴上的拷贝范围： + +```text +StateTransferLayout # 一个 State 的传输布局 +├── resource_id # 对应的 State +├── tokens_per_page # 外部页的 token 容量 +├── byte_start / byte_end # 本 State 的页内预留区间 [start, end) +└── regions: tuple[TransferRegion, ...] + └── TransferRegion # 一个字段与外部区域的映射 + ├── field / byte_start / byte_end / dtype / axes / shape + ├── axis_slices: tuple[AxisSlice, ...] + └── write_enabled +``` + +同一类布局页可包含多个 State 的 layout，共用该类布局的 token 容量。每个 layout 的 `[byte_start, byte_end)` 只描述本 State 的预留区间，包含该 State 的字段、对齐及预留空间；其中每个 region 的区间只包含对应字段的实际数据。后一个 State 从前一个 State 的终点开始排列，所有 layout 的最大终点表示该类布局的总占用。 + +CPU buffer 使用 `(page_num, page_bytes)` 的二维页池，通过 `cpu_page_buffers[page_id]` 选择物理页；页间距由实际 buffer 决定，不保存在 State 布局中。CPU 整页大小按全部 State 的占用向上对齐到 16 字节。PD 的 `kv` 与 `att_state` 分别生成自身区间，`get_pd_page_bytes()` 统一选择实际传输块大小，可能大于运行态布局的占用。同一字段也可通过多个 region 描述不同范围，例如 Conv 的 Q/K/V channel 切片。读回时反向使用字段与外部区域之间的映射。 + +这些类型均为只读数据参数，不持有 buffer 或执行搬运。具体 State 的方法根据 layout 完成数据操作。 + +| 字段 | 语义与单位 | +| --- | --- | +| `resource_id` | 与管理器成员名和 State 的资源名称一致，用于选择当前资源的页区域 | +| `tokens_per_page` | 含 KV 的页格式为正整数 token 容量；纯递推状态页为 `None`，按请求状态或 checkpoint 槽位处理 | +| `byte_start` | 本 State 相对整页的起始字节，包含该位置；region 已使用页内坐标，搬运时不再重复加上此起点 | +| `byte_end` | 本 State 的结束字节，不包含该位置，涵盖自身字段、对齐及预留空间；与起点之差是本 State 占用，不是物理页间距 | +| `regions` | 本 State 参与传输的物理字段或片段列表，同一字段可分为多个映射项 | +| 每项的 `field`、`dtype`、`axes`、`shape` | 标识 State 中的字段及外部页区域；axes 命名各维度，shape 用元素数量；外部区域采用连续布局，stride 由 shape 和 dtype 推导 | +| 每项的 `byte_start`、`byte_end` | 字段的左闭右开字节区间，组合后相对整页；长度等于 `prod(shape) * dtype.itemsize`，不包含字段间的对齐空隙，必须位于所属 State 区间内 | +| 每项的 `axis_slices` | `tuple[AxisSlice, ...]`，通过 `axis`、`src_start`、`dst_start`、`length` 描述轴名、源起点、目标起点和元素数量；一个 region 内同一轴只出现一次，同轴不连续范围拆成多个 region | +| 每项的 `write_enabled` | 指定 TP 副本是否负责写入该区域，避免多个副本重复写入；恢复时各目标副本均可读取 | + +以下冻结数据类只描述传输位置，不包含搬运方法： + +```python +from dataclasses import dataclass, replace +import torch + + +@dataclass(frozen=True) +class AxisSlice: + axis: str + src_start: int + dst_start: int + length: int + + +@dataclass(frozen=True) +class TransferRegion: + field: str # 对应 State 内部搬运字段，如 k/v/scale/packed + byte_start: int # 字段起点；State 返回局部位置，manager 组合后相对整页 + byte_end: int # 字段终点，不包含；区间长度等于 shape 对应的数据字节数 + dtype: torch.dtype # 外部区域的数据类型 + axes: tuple[str, ...] # 如 (token, layer, head, dim) + shape: tuple[int, ...] # 外部区域在各轴上的完整尺寸，不含页池的 page 轴 + axis_slices: tuple[AxisSlice, ...] + write_enabled: bool = True + + +@dataclass(frozen=True) +class StateTransferLayout: + resource_id: str # 含模型角色，如 main_full_att_fp16_state + tokens_per_page: int | None # KV 页为 token 容量，纯状态页为 None + byte_start: int # 本 State 的页内起点，包含 + byte_end: int # 本 State 的页内终点，不包含，含自身对齐与预留空间 + regions: tuple[TransferRegion, ...] +``` + +`field` 在初始化时与 State 字段匹配；本地 Tensor 继续保留自身的 dtype、stride 和字段起点，可以是 packed buffer 的非连续视图。外部页中每个 region 的完整区域采用连续布局，`axes` 与 `shape` 长度相同且轴名唯一，逐维 stride 由后续维度大小乘 dtype 字节数推导，不再作为输入字段。多个 region 可以通过切片引用同一外部区域，例如 Conv 的 Q/K/V 分片。部署配置确定布局，CPU cache 读写直接使用字段区域与 TP 切片。 + +`AxisSlice` 的起点和长度均以元素为单位。在 `TransferRegion` 中，方向固定为本地 State 到外部页，读回时交换源、目标起点。构造时使用字段名，搬运时通过属性访问，不再依赖四元组的位置。 + +层维度统一完整、同序复制，搬运前检查源目标层数一致。逻辑层到物理存储位置的映射由 State 自身维护，例如逻辑层 `[3, 7]` 紧凑存为 `[0, 1]`，外部页沿用相同层顺序。CPU/PD 区域和 DP State 复制均不提供层子集选择或层重排参数。 + +token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_slices`。其余搬运轴通过 `AxisSlice` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。使用多个连续片段表达同一轴的非连续范围。 + +`shape` 表达该外部字段区域的完整范围，而 `axis_slices` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： + +```python +replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) +``` + +最终 `StateTransferLayout.regions` 中的区间相对整张外部页，搬运时直接使用。共用搬运函数要求 buffer 至少为二维，并沿第一维选择对应页；PD State 入口直接接收页集合与显式页索引。取页后直接按字段字节区间、dtype 和 shape 构造视图，不重复断言布局边界。 + +下面将前述结构用于 P 端 TP=2、rank 1：本地 K/V 各有 2 个 heads，外部页使用 4 个全局 heads。假设组合页共 64 KiB，其中前 32 KiB 存 main K/V，后 32 KiB 用于其他资源： + +```python +k_region = TransferRegion( + field="k", + byte_start=0, + byte_end=16384, + dtype=torch.float16, + axes=("token", "layer", "head", "dim"), + shape=(8, 2, 4, 128), + axis_slices=( + AxisSlice(axis="head", src_start=0, dst_start=2, length=2), + AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + ), +) +v_region = replace(k_region, field="v", byte_start=16384, byte_end=32768) + +layout = StateTransferLayout( + resource_id="main_full_att_fp16_state", + tokens_per_page=8, + byte_start=0, + byte_end=32768, + regions=(k_region, v_region), +) +``` + +示例中的配置标识字符串只为说明关联关系，不规定标识编码方式。K/V 每个完整区域为 `8 * 2 * 4 * 128 * 2 = 16384` 字节；当前 rank 通过 head 映射只写各区域的一半，另一半由另一个 TP rank 写入。逻辑层 `[3, 7]` 到本地层 `[0, 1]` 的关系保留在 State，外部区域按相同顺序完整存放这两层。 + +例如当前 rank 的本地 head 0、token 在页内位置 0、物理层 0、dim 0，对应 K 区域地址偏移为 `0 + 2 * 256 = 512` 字节，V 为 `16384 + 512 = 16896` 字节。CPU buffer 每行 `65536` 字节，选择下一行即进入下一物理页；本 State 的 `byte_end - byte_start = 32768` 只表示自身占用。 + +D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸，把本地 head 映射改成 `AxisSlice(axis="head", src_start=0, dst_start=3, length=1)`。P/D 根据各自的 TP 配置生成切片，完整 layout 对象可以不同。 + +字段列表按实际物理表示组织。例如 K、V、scale 可分别描述；packed 数据也可整体描述为一个字节区域。别名区域只搬运一次。Conv 的 Q/K 和 V、SSM 可能采用不同 head 映射，应分别描述，不能把一个资源的 `head_start` 套用于所有字段。数据类型、量化组粒度及索引单位由对应格式约束。 + +具体管理器从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 + +FullAttentionState 的 token 页操作要求 `tokens_per_page` 为正整数。CPU 参数 `cpu_page_buffers` 为完整二维页池；PD 参数 `pd_page_buffers` 同样为页集合，`pd_page_index` 指定物理页。传入完整页集合，不提前切到单页或 State 区域。对于常规 K/V，一个外部地址可以写成： + +```text +page_base = cpu_page_buffers[page_id].data_ptr() # CPU 选择对应行 +page_base = pd_page_buffers[pd_page_index].data_ptr() # PD 选择对应行 + +page_base ++ region.byte_start ++ token_in_page * token_stride_bytes ++ buffer_layer * layer_stride_bytes ++ buffer_head * head_stride_bytes ++ dim * dim_stride_bytes +``` + +`data_ptr()` 已指向传入 Tensor 的起点，不重复加 Tensor 的 storage offset。State 内部源/目标地址使用该 State 自身 Tensor 的基址及 stride；PyTorch 元素 stride 进入字节描述前乘以对应 dtype 大小。非连续区域必须保留真实 stride，不能通过会产生拷贝的 reshape/contiguous 假装仍指向原写入目标。 + +CPU/PD 的 token 页操作统一从外部页内槽位 0 开始，公共接口不设置页内 token 起点参数。令 `i` 为本次索引向量位置： + +```text +# CPU：按给定页 ID 顺序，从每页槽位 0 开始 +page_id = cpu_page_indices[i // tokens_per_page] +token_in_page = i % tokens_per_page + +# PD：通过显式索引选择本次传输页 +page_id = pd_page_index +token_in_page = i +``` + +有效 token 数由索引长度给出。普通 CPU 路径要求整页操作；支持尾页的路径允许最后一页只有前部有效，CPU 页 ID 数量须与覆盖页数一致，PD 单页要求 `len(token_indices) <= tokens_per_page`。checkpoint 对应的前缀长度由调用方关联到槽位，不作为页拷贝参数，也不保存在 pool 中。 + +上述固定的是 token 维度起点。组合资源的 `byte_start` / `byte_end`、TP head/channel 范围以及 State 的实际物理 token 索引仍可非零，应继续保留在各自的位置映射中。 + +#### DP 同名 State 的直接复制 + +DP 请求共享信息分为两个数组，均按请求在当前批次中的位置和节点内 DP 编号索引： + +- `shared_req_infos` 保持 `[max_req_num, dp_size_in_node, 2]`，两项分别为 KV 长度和请求槽位。 +- `shared_req_big_page_ids` 独立使用 `[max_req_num, dp_size_in_node, max_big_pages]`，其中 `max_big_pages = max_req_total_len // big_page_tokens`；未启用大页时不创建。 + +大页列表第 `i` 项对应前缀长度 `(i + 1) * big_page_tokens`。主 TP rank 将 radix 已命中的大页与请求持有的大页按对应边界写入，每次先把当前行清为 `-1`，避免较短请求复用共享行时读到旧页 ID。发布和读取沿用节点内 barrier，源请求引用保持到所有复制完成。 + +`TransTask` 使用 `src_token_indexes` / `dst_token_indexes`、`src_big_page_ids` / `dst_big_page_ids` 明确配对关系；`dst_req` 持有目标请求，`src_req_idx` 指定来源运行态，`src_dp_rank` 以节点内 DP 编号定位对应 TP 分片的远端 manager 和请求索引表；节点内进程 rank 只在初始化加载 IPC 资源时使用。目标 token 索引来自 CPU allocator，在调用资源接口前一次性转成 CUDA int64。 + +目标 DP 只复制从本地已有 KV 长度到来源 KV 长度之间新增的完整大页，不重复复制已有前缀页。先申请目标大页再申请 KV；目标大页分配必须成功，通过断言检查。复制完成后，将目标页按前缀长度登记到 `req.hybrid_len_to_big_page_id`，后续 radix 插入和释放沿用请求的现有归属管理。未跨越新大页边界时，大页列表为空。 + +DP 副本之间使用相同的 State 成员名。`copy_from_other_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.resource_id)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_run_time_state_from_other_dp`，再清零目标 tail。调用方保证源 MTP 游标为 0、前缀处于 tail 可重置边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 + +full-attention 使用 `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_page_state_from_other_dp` 与 `copy_run_time_state_from_other_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源索引可为 int32/int64，目标索引须为 int64（当前 `index_copy_` 的要求）。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 + +CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=36)` 控制每个字段发射的 program 数；kernel 从 `program_id * BLOCK` 开始，按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 保持各自的数据操作路径。 + + + +DP 直接复制两端 State 的字段，buffer 基址、容量和 stride 从各自 State 获取。层维度完整、同序复制；字段类型、形状、编码及固定参数的检查仍由具体实现完成。不保存额外的复制映射。 + +DP 只复制对应 TP rank 的完整分片，不支持通过 State 接口拼接多个 TP 分片。不同 token 来自不同 DP 时，先按来源分组索引，再逐项调用 `copy_tokens_from_other_dp`；同一批内配对源与目标 token 索引。 + +源与目标布局完全一致、分片一一对应时，可以采用同构布局的复制 kernel，但须先确认源目标 stride 等条件成立。重分片或编码变化由具体子类明确声明支持并提供实现。 + +#### main/draft 与不同 TP 的例子 + +假设 main 的 full-attention 逻辑层为 `[3, 7]`，本地压缩存储层为 `[0, 1]`；一页 16 个 token，全局 4 个 KV heads，每 head 64 个 FP16 元素。以 K、V 分区的示例格式组织 `kv` 页： + +```text +byte 0 : main_full_att_fp16_state K [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 16384 : main_full_att_fp16_state V [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 32768 : draft0_mla_fp8_state 的数据区域，按 draft 自身格式定义 +``` + +在该示例格式下，main 的 P 端 TP=2 时，rank 1 的本地 head `[0, 2)` 映射到页内全局 head `[2, 4)`；D 端 TP=4 时，rank 3 的本地 head `[0, 1)` 从页内全局 head `[3, 4)` 恢复。两端共享页格式，使用不同的局部分片映射。 + +如果 draft 使用 TP 副本形式的 MLA，draft 区域由选定的 P 端副本写一次，D 端各副本读取该区域。它不沿用 main 的 head 偏移公式。main 与 draft 即使共享 token allocator，也分别取得自己的区域映射。 + +CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索引数量确定、格式由外层描述；PD 使用 `kv` 与 `att_state` 两类任务，各自拥有布局，不强制把 KV、Conv、SSM 和 tail 放进同一页。采用按 TP rank 分块的格式时，必须明确其 TP 配置约束；跨 TP 配置恢复需要相应格式和转换能力。 + +#### 初始化与提交时的校验 + +布局由部署配置和 State 生成；搬运时校验资源名称、字段覆盖、dtype、字节对齐、页内地址边界及层/轴映射范围，TP 切片指定分片范围与副本写入者。各资源的实际写入区域不得相互覆盖；packed 别名按实际物理范围去重。涉及多个布局的描述以实际访问地址为准，不能仅比较 shape。 + +数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。具体管理器校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 + +### 6.6 子类实现与直接调用示例 + +以下为伪代码，layout/mapping 在初始化时由管理器组织并经具体 State 校验，调用过程保持直接。 + +```python +class GQAFullAttentionState(FullAttentionState): + def write_to_pd_buffer( + self, token_indices, pd_page_buffers, *, pd_page_index: int, layout, + ) -> None: + gqa_kv_to_pd_buffer( + kv_buffer=self.kv_buffer, + token_indices=token_indices, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + layout=layout, + ) + + def read_from_pd_buffer( + self, pd_page_buffers, token_indices, *, pd_page_index: int, layout, + ) -> None: + gqa_pd_buffer_to_kv( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kv_buffer=self.kv_buffer, + token_indices=token_indices, + layout=layout, + ) +``` + +CPU 操作按路径选取该 State 的位置映射： + +```python +state.write_to_cpu( + src_token_indices, cpu_page_buffers, cpu_page_indices, + layout=cpu_write_layout, +) +# 当前 stream 上记录 event；所有资源写入完成后发布 CPU 页 ready。 + +state.read_from_cpu( + cpu_page_buffers, cpu_page_indices, dst_token_indices, + layout=cpu_read_layout, +) +# 后续计算与写入建立 stream/event 依赖。 +``` + +DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由具体管理器组织： + +```python +dst_state.copy_tokens_from_other_dp( + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + remote_dp_full_att_state=src_state, +) +``` + +PD 两端分别传入完整发送/接收页集合与本地页索引；两端页索引可以不同，区域与 TP 位置来自各自 layout： + +```python +p_state.write_to_pd_buffer( + src_token_indices, send_pages, + pd_page_index=send_page_index, + layout=p_layout, +) +# 全部参与资源/分片的打包完成后，调度层发送。 + +# D 端接收完成并建立 stream 依赖后执行。 +d_state.read_from_pd_buffer( + recv_pages, dst_token_indices, + pd_page_index=recv_page_index, + layout=d_layout, +) +# 所有目标资源/分片恢复完成后允许 decode。 +``` + +调度层负责网络、同步和生命周期。方法返回不等于数据已经就绪，发送页在网络发送完成前不能复用，源槽位在读取完成前不能释放。KDA/GDN 状态的统一操作见 6.7 节，indexer tail 的操作见 6.8 节；它们共用适用的位置描述，不继承 `FullAttentionState`。 + +### 6.7 BoundedAttentionState：有界状态的统一操作接口 + +`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。模型管理器直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 + +```text +BoundedAttentionState # 抽象操作契约 +├── KDAState # Conv + KDA 递推状态 + 原生 MTP 布局 +├── GDNState # Conv + GDN 递推状态 + 原生 MTP 布局 +└── 后续 SlidingWindowState 等有界状态实现（各自定义布局与页池） +``` + +`bounded_att_*_allocator` 的范围仍可覆盖 sliding-window 等有界状态，但使用该 allocator 不要求继承 `BoundedAttentionState`。 + +#### 状态操作约定 + +- forward 通过层视图更新 Conv/SSM；MTP 提交选择已接受状态。提交可以只更新游标。当前大小页保存和 DP/PD 导出仅处理 prefill 基准状态，KDA/GDN 的源游标必须为 0，不支持 MTP 提交后保存或导出。 +- 页保存只接受实际捕获的状态位置。跨过中间边界而未保留其状态时，只能预留槽位，不能发布为有效页。 +- 大小页 pool 保存独立快照，默认使用 CPU pinned memory；DP/PD 操作请求运行态。两者使用不同索引和生命周期。 + +#### KDAState 内部布局示例 + +设线性层数为 `L`，请求槽位容量为 `N`（已经包含 HOLD），本地 head 数为 `H`，head 维度为 `D`，Conv 历史宽度 `W = kernel_size - 1`，MTP 额外候选数为 `S`。该 KDA 布局中，Q/K/V 均为 `H * D` 通道： + +| 存储 | 示例形状 | 内容 | +| --- | --- | --- | +| GPU `conv_state` | `[L, N, 3*H*D, W+S]` | 历史窗口及验证时暂存的候选尾部 | +| GPU `ssm_state` | `[L, N*(S+1), H, D, D]` | 每请求 `S+1` 个递推状态槽位 | +| GPU `mtp_state_index` | `[N]`，int32；S=0 时可为 None | 当前已接受状态的原生选择索引 | +| 某个 pool 的 Conv | `[C, L, 3*H*D, W]` | 一个确定前缀的规范化历史窗口 | +| 某个 pool 的 SSM | `[C, L, H, D, D]` | 同一前缀对应的递推矩阵 | + +大、小 pool 的 C 分别取管理器给出的容量；两者保存同种状态快照，区别主要在分配域、保存边界和容量，不是按 token 页大小重复展开运行态。pool 不持有空闲列表。默认使用 CPU pinned 快照存储，设备和 pin memory 选项由 `initialize_page_pools()` 的调用方指定，不在 State 中重复保存。 + +GDN 的 Q/K 与 V head 数和维度可以不同，其 Conv/SSM 形状由自身实现确定，不把上述 `3*H*D` 或方阵要求放进基类。KDA/GDN 即使 shape 相同，也不能互换状态内容。 + +#### 必须实现的公共接口 + +所有下列方法均为抽象方法。子类可继承已验证兼容的具体实现；对于未支持的 checkpoint、MTP 或传输功能，也必须明确实现为报错并在对应接口中说明限制，不能依靠基类空方法静默成功。初始化只要求实际启用的功能可用。 + +```python +from abc import ABC, abstractmethod +from typing import Any, List, Literal + +PageKind = Literal["big", "small"] + + +class BoundedAttentionState(ABC): + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + @abstractmethod + def get_state_type(self) -> str: ... + + @abstractmethod + def initialize_page_pools( + self, big_page_capacity: int, small_page_capacity: int = 0, *, device="cpu", pin_memory: bool = True + ) -> None: ... + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + @abstractmethod + def get_layer_view(self, layer_id: int) -> Any: ... + + @abstractmethod + def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... + + @abstractmethod + def save_bounded_att_pages( + self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + ) -> None: ... + + @abstractmethod + def restore_bounded_att_pages( + self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + ) -> None: ... + + @abstractmethod + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: ... + + # ======================================================================== + # 三、CPU Cache 数据交换：读写大小页快照 + # ======================================================================== + @abstractmethod + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + + @abstractmethod + def write_page_state_to_cpu( + self, page_kind: PageKind, slot_ids, cpu_page_buffers, cpu_page_indices, + *, layout: StateTransferLayout, cpu_page_readies: List[bool] | None = None, + ) -> None: ... + + @abstractmethod + def read_page_state_from_cpu( + self, cpu_page_buffers, cpu_page_indices, + *, big_page_slot_ids: List[int], layout: StateTransferLayout, + ) -> None: ... + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + @abstractmethod + def copy_big_page_state_from_other_dp( + self, src_page_ids: List[int], dst_page_ids: List[int], + remote_dp_bounded_state: BoundedAttentionState, + ) -> None: ... + + @abstractmethod + def copy_run_time_state_from_other_dp( + self, src_req_indices: List[int], dst_req_indices: List[int], + remote_dp_bounded_state: BoundedAttentionState, + ) -> None: ... + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + @abstractmethod + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: ... +``` + +bounded State 使用 `write_page_state_to_cpu` / `read_page_state_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 + +CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。PD 状态页按单请求处理。`req_indices` 是 ReqManager 已分配的请求槽位，`slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_other_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 + +#### 运行态与推理访问 + +| 接口 | 必需语义 | +| --- | --- | +| 初始化 | 构造函数按 `request_capacity` 和 `mtp_step` 创建运行态及层映射，大小页池先为 `None`；不创建 allocator，不分配请求号或页槽位 | +| `initialize_page_pools(...)` | 容量规划完成后创建大小页池，覆盖共享 allocator 的槽位域；容量为零的页池保持 `None`，不重建运行态或重置 MTP 游标 | +| `get_layer_view(layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | +| `reset_bounded_att_state(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | +| `commit_bounded_att_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | + +`accepted_steps` 为逐请求接受的、已经执行过 forward 的输入步数，包含本轮基准输入步，不是“接受的 draft token 数”或新采样输出数。取值范围为 `[1, 本轮实际验证宽度]`；KDA/GDN 的原生索引为 `accepted_steps - 1`。实现不应把不存在的零步回滚当作有效操作。未启用 MTP 时，合法的一步提交可为无数据移动的操作。 + +KDA 的提交可以只更新 `mtp_state_index`，下一轮 kernel 根据该游标选择输入状态。当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0，直接读取 Conv 的 `[:W]` 窗口和 SSM 的 `req_idx*(S+1)` 基准槽位。导入时写到目标的基准位置，并将游标设为 0。实现中以 TODO 标明：若以后支持 MTP 提交后传输，需要按已接受索引 j 选择 Conv 的 `[j:j+W]` 窗口和 SSM 的 `req_idx*(S+1)+j` 槽位。 + +prefill 的最终状态、普通 decode 的结果和已提交验证状态都可以成为“当前已接受状态”；执行流程必须保证所有相关层已完成且选择信息与数据一致。验证结果尚未提交时，不允许将临时候选作为当前状态发布。大小页保存仅发生在 prefill 阶段，此时 `mtp_state_index` 为 0，直接复制 Conv 的 `[:W]` 窗口及 SSM 的 `req_idx*(S+1)` 基准槽位，不需要先捕获候选状态。实现中以 TODO 标明：若以后支持 decode/MTP 后保存，需要改为按已接受游标选取数据。 + +#### 状态页：只保存真实存在的前缀状态 + +| 接口 | 必需语义 | +| --- | --- | +| `save_bounded_att_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req 与 slot 逐项对齐，前缀关联由调用方维护 | +| `restore_bounded_att_pages(...)` | 从调用方选定的槽位恢复运行态和算法控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | + +调用方负责将快照关联到实际执行后的位置；KDA/GDN 的页保存契约只覆盖传入请求的 prefill 完成状态，不接收用于声称历史位置的长度参数。chunk 跨过中间边界但没有保留对应运行态时，预留槽位不能作为已保存的页面发布。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 + +保存或 CPU 导入时直接写入指定槽位；异步完成前由上层保持 pending，完成后才发布或读取。发生失败直接透传异常，调用方不发布该批页面，并等已提交操作结束后才可回收。恢复前由调用方保证前缀关联和来源就绪，State 按兼容布局搬运数据。 + +同一槽位数字在 big/small pool 中属于不同分配域。main/draft 的相关 State 共享分配的槽位 ID，各自保存快照数据；管理器组织所有参与 State 的保存和恢复，缓存调用流程在数据就绪后发布；最终释放时，调用方解除引用并确认使用结束，再直接向 allocator 归还一次。 + +#### CPU、DP、PD 的状态单位 + +| 路径 | 传输单位与接口职责 | +| --- | --- | +| CPU 写出 | `write_page_state_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，长度由缓存任务携带，不写入页内；每个待发布页的参与状态必须完整；CPU pool 到 CPU 页同样是独立的格式搬运 | +| CPU 读回 | State 的 `read_page_state_from_cpu` 按 `big_page_slot_ids` 将快照写入大页池；管理器补上碎页 LOAD 槽位，并调用 `restore_bounded_att_pages` 恢复最后一页运行态 | +| DP 大页复制 | `copy_big_page_state_from_other_dp(src_page_ids, dst_page_ids, remote_dp_bounded_state)` 在同一 TP 分片的 DP 副本之间复制大页 Conv/SSM 快照；不分配页、不修改请求运行态或小页，前缀关联与源页生命周期由调用方维护 | +| DP 运行态复制 | `copy_run_time_state_from_other_dp(src_req_indices, dst_req_indices, remote_dp_bounded_state)` 复制同一 TP 分片的 prefill 基准 Conv/SSM；KDA/GDN 源游标必须为 0，目标游标归零，不修改大小页或其他候选区 | +| PD 写出 | `write_to_pd_buffer` 将一个请求的 prefill 基准状态写入 `pd_page_index` 指定的状态页区域，按协议关联其前缀长度 | +| PD 读回 | `read_from_pd_buffer` 将完整快照恢复到指定请求槽位，不接收前缀长度或读取页内长度、版本字段；不自动生成历史 checkpoint | + +CPU 尾页若需要保存 prefill 完成后的运行态，管理器先使用保留槽位调用 `save_bounded_att_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 + +CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字段映射，直接复制同一算法、同一配置及 TP 分片的对应存储。Conv 区域可用 `(layer, channel, window)` 轴，SSM 用 `(layer, head, key_dim, value_dim)` 轴。它们没有历史 KV 的 token 轴。含 KV 的 CPU 页仍有 token 容量，checkpoint 组件不按该容量重复展开；纯 PD 状态页的 `tokens_per_page` 为 None。各资源偏移、TP 范围和副本写入者由管理器组合，State 直接执行方法。 + +当前 DP/PD 导出 prefill 基准状态，调用方保证 KDA/GDN 源游标为 0，不传 MTP 候选区，也不支持 MTP 提交后导出。目标可按自己的 MTP 步数初始化工作区；DP 两端配置及 TP 分片一致，按各自 MTP 步数计算基准槽位。PD 则通过全局布局与 TP 切片处理两端分片差异。不能因为 KDA/GDN 的 Conv/SSM 形状相同就跨算法恢复。 + +多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；具体管理器完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 + +#### 依赖、同步与使用示例 + +资源方法接收请求槽位、checkpoint 槽位和明确的资源引用，只有 tail 无快照恢复额外接收前缀长度，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；模型管理器单独协调 tail 在恢复边界上的重置或迁移。 + +GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 + +```python +# ReqManager 已分配请求槽位；调用方已从管理器持有的 allocator 分配状态页槽位。 +state.reset_bounded_att_state(req_indices) +# backend 经层视图完成 prefill；调用方维护实际 chunk 终点。 +state.save_bounded_att_pages( + req_indices, + page_kind="big", slot_ids=slot_ids, +) +# 确认所有参与状态已按指定长度保存并完成传输后，缓存调用方发布联合页。 + +state.restore_bounded_att_pages( + dst_req_indices, + page_kind="big", slot_ids=slot_ids, +) +# target 请求在相关 KV、linear、tail 均就绪后继续执行。 + +# 验证 forward 完成后，将接受步数按请求提交。 +state.commit_bounded_att_state(req_indices, accepted_steps) +# 后续计算 kernel 使用已提交游标；当前不支持在 MTP 提交后保存大小页或进行 DP/PD 导出。 +``` + +各具体实现必须满足数据完整性、接受位置选择和跨布局复制约束。 + +### 6.8 AttentionTailState：原始尾部窗口的操作接口 + +`AttentionTailState(ABC)` 约束注意力相关模块的尾部运行态操作，具体实现负责自身窗口、位置规则和搬运 kernel。例如 `KPoolAttentionTailState` 定义 GLM K-pool 的尾部布局。通用分类使用 `att_tail`，具体类及 `indexer_kpool_c4_tail` 等资源标记保留算法信息。该分类不包含算子临时 workspace、完整递推状态或随历史增长的索引 key cache；对齐后可重置是 K-pool 的具体规则，不作为所有 tail 实现的共同假设。它不继承 `FullAttentionState` 或 `BoundedAttentionState`，也不持有 allocator、请求 KV 索引表或独立大小页 pool。 + +#### KPoolAttentionTailState 布局与状态边界 + +设本实例层数为 `L`、请求容量为 `N`（包含 HOLD）、池化步长为 `P`、MTP 额外容量为 `S`、index head 维度为 `D`，原生 ring 布局为 `[L, N, P+S, 2*D]`，末维分别存 raw key 与 gate score。GLM K-pool 示例取 `P=4`、`D=128`，具体实现须声明其 kernel 支持的配置范围。 + +计算先用此前 tail 和本轮 raw 输入完成跨 chunk 池化，再更新 tail。绝对位置 `t` 写入 `t % (P+S)`。完成池化的压缩 key/scale 写入 full-attention packed buffer，由 `FullAttentionState` 管理；tail 保留增量计算所需的原始数据。 + +PD 只传输 prefill 结束时最近 `P` 个逻辑位置，不传输完整 `P+S` 物理 ring。CPU cache 与 DP fetch 恢复对齐的缓存前缀，直接重置目标 tail。D 端仍按本地 `P+S` 分配 ring，保留后续 MTP 验证、拒绝回退和 draft 回访所需的候选历史。PD 逻辑窗口仅用于 prefill 状态交接，不用于迁移已经执行 MTP 的运行态；前缀恢复的可重置边界由模型配置和缓存流程保证。 + +#### 必须实现的接口 + +构造函数直接接收配置、请求容量 `request_capacity`、MTP 步数 `mtp_step` 和分配设备 `device`，不再使用 `TailStorageSpec`。请求容量包含 HOLD 槽位；容量与 MTP 步数各在 State 上保存一份，设备从 `tail_buffer.device` 获取。K-pool 在构造阶段分配完整的 `pool_size + mtp_step` ring，不依赖后续 KV token 容量规划。 + +基类接口仅包含类型标记、层访问、边界重置及 PD 区域/读写；构造参数由各具体 State 和算法 create 入口定义，不要求统一使用 KPoolStateConfig。CPU cache 不保存 tail，CPU/DP 前缀恢复统一调用重置接口；tail 不提供 CPU 读写或 DP 复制接口。 + +```python +from abc import ABC, abstractmethod +from typing import Any + + +class AttentionTailState(ABC): + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + @abstractmethod + def get_state_type(self) -> str: ... + + # ======================================================================== + # 二、Attention tail：计算视图与无快照重置 + # ======================================================================== + @abstractmethod + def get_layer_view(self, layer_id: int) -> Any: ... + + @abstractmethod + def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... + + # ======================================================================== + # 三、PD 分离传输 + # ======================================================================== + @abstractmethod + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx, pd_page_buffers, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_page_buffers, req_idx, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: ... +``` + +重置接口只接收请求索引列表并清零对应 tail，不接收前缀长度。PD 按单请求和显式 `pd_page_index` 搬运固定逻辑窗口，`state_token_len` 表示该 State 已完成 prefill 的前缀长度。请求槽位由 ReqManager 分配，State 不保存请求长度表。 + +#### 容量、访问与恢复 + +| 接口 | 语义 | +| --- | --- | +| 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | +| `get_layer_view(layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | +| `reset_attention_tail_state(req_indices)` | 清零指定请求的 tail，适用于新请求或可重置的缓存前缀；K-pool 清零完整 ring，其他实现清零对应数据与控制信息。模型配置和缓存流程负责保证边界，接口不接收长度 | + +`get_layer_view` 供计算 backend 直接使用；池化与写 tail 仍在模型 kernel 中完成，不增加通用的 `append_token` 或 attention forward 包装。执行顺序必须满足“先池化、再更新 tail”。 + +`reset_attention_tail_state` 是新请求/前缀恢复入口，不是 MTP 回滚入口。即使已接受长度恰好对齐 P，运行中的 draft 重访仍可能需要 ring 内历史，不能用该条件代替 MTP 状态处理。模型初始化与缓存流程分别保证 main/draft 的恢复条件；调用方还需确认对应 full-attention、linear 等资源全部恢复到约定位置。 + +#### PD：传输 prefill 逻辑窗口并适配本地 ring + +| 接口 | 数据路径及约束 | +| --- | --- | +| `write_to_pd_buffer(...)` | 按 `state_token_len` 提取指定请求的逻辑窗口 → `pd_page_index` 指定的 PD 状态页区域，使用 layout 的资源偏移与 TP 写入规则 | +| `read_from_pd_buffer(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；清零该请求的其余候选历史，调用方建立接收完成依赖 | + +这些方法传输 tail 自身的 prefill 逻辑窗口，不承担历史 KV、池化 key/scale、bounded 状态或 MTP proposal 的迁移。发送端取最近 `P` 个位置并按时间顺序排列；接收端先清零该请求的本地 ring，再把有效位置放回对应物理槽位，其他请求保持不变。后续 MTP 候选历史由 D 端执行过程中产生。 + +`StateTransferLayout` 描述外部页位置与切片。K-pool region 的轴为 `("layer", "window", "channel")`，形状为 `[layers, P, 2*head_dim]`，与本地 `mtp_step` 无关。`window` 按时间顺序对应 `[state_token_len-P, state_token_len)`；短前缀中小于 0 的位置在发送时补零，接收时忽略。独立状态页的 `tokens_per_page=None`，该窗口不是 full-attention 的页内 token 轴。 + +P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_step`。位置 `pos` 在 P 端从 `pos % (P+S_p)` 读取，在 D 端写入 `pos % (P+S_d)`；不能直接截取 ring 的前 P 个物理槽位。模型算法、池化步长、层顺序、字段精度和逻辑维度仍须兼容。现有 main/draft prefill 保持相同的序列位置，因此可共用任务的 `end_kv_index`；位置变化不影响缓存的布局描述。 + +例如 `P=4`、`S_p=0`、`S_d=2`、前缀长度 11，外部窗口保存位置 `[7,8,9,10]`；P 端读取槽位 `[3,0,1,2]`，D 端恢复到 `[1,2,3,4]`。D 端其余 ring 槽位清零。PD 块仍使用 KV 页确定的固定容量，减少 tail 有效区域和本地复制量不代表网络整页传输字节数同步减少。 + +长度、前缀身份和格式来自保存/传输控制元数据，不能从原始 tail 字节反推。调用层负责核对元数据并维持 main/draft 的一致恢复点；State 负责布局和数据操作。所有 GPU 操作遵循当前 stream，返回不表示完成；源请求与目标页在读写完成前不能复用。 + +#### 与 linear 页接口和 MTP 提交的区别 + +- 不提供 `save_bounded_att_pages`、`restore_bounded_att_pages`：tail 没有独立大小页 pool。CPU/DP 缓存前缀在对齐边界通过 reset 恢复；非对齐的 PD 交接使用 prefill 逻辑窗口。 +- 不要求 `commit_bounded_att_state`：K-pool 没有类似 KDA 的 SSM 选择游标，按绝对位置访问 ring，MTP 拒绝后由执行流程选择后续位置。增加空的提交方法不能替代窗口容量与重访范围的正确性；未来出现确实需要提交控制状态的实现时，再扩展对应契约。 + +## 7. 资源初始化与调用约定 + +### 7.1 初始化与绑定 + +BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己的接口，各 draft 调用各自的接口，模型方法不接收其他模型或组织跨模型循环。`ModeBackend.init_model()` 直接组织 main/draft 构造、资源准备、共享分配和各模型运行时初始化，并在全部初始化成功后发布就绪标志。benchmark 和资源验证脚本在各自调用点展开相同流程，不另设跨模型运行时启动封装。模型构造保留配置、权重、请求管理器和计算层的准备。 + +| 单模型公共接口 | 封装的内部操作 | +| --- | --- | +| `prepare_att_resources()` | 在本模型 KV 内存分区内调用 `_init_att_resource_manager`，根据本模型配置准备自身全部 State | +| `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | + +1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。ReqManager 仍由 main/draft 共用，其构造不依赖 attention 资源已经初始化。 +2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。main 的资源 hook 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，其资源 hook 通过 `set_att_state` 只添加本模型的 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。静态 FP8 标定偏移由模型侧配置 helper 计算：main 从 0 开始，draft 累加 main 及前序 draft 的实际 full-attention 层数；混合模型排除 bounded-attention 层，原生 MTP 使用自己的层定义，无 attention 的 draft 不占标定行。公共准备接口及资源 hook 均不接收标定偏移，也不返回层数。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 +3. backend 检查 main 和所有 draft 使用相同请求容量且引用同一个管理器，随后调用 `manager.initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)`。管理器先标记 State 创建完成；每个模型的成员已经直接保存在共享实例中，无需合并管理器或 State。不同 draft 即使方案相同，也可拥有不同的配置、存储和游标。 +4. `finish_state_creation()` 完成后，通过管理器接口计算容量并初始化实际存储。FullAttentionState 原地分配 buffer；不同 State 可以具有不同物理布局，但共享 allocator 的索引域及容量；同一分配域只创建一个 allocator。实际存储分配完成后再配置 CPU/PD 布局。 +5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 +6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 + +BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由单模型公共阶段接口调用，构造函数不调用运行时初始化。main 在资源 hook 中选择具体管理器，各模型分别在自己的 hook 中准备所需 attention State，子类可以继承适用的模型族实现。原生 MTP 如需复用 main 的资源配置，也在自身资源准备阶段读取；ReqManager 初始化只处理请求映射共享。 + +需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。分配完成后的 backend 绑定、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方可按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 + +单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 + +`AttResourceManager.initialize_resources()` 由基类提供通用实现，保留三个阶段：标记 State 创建完成、估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 + +自动容量规划使用所有模型实际加载后的剩余显存，不再按 main 权重比例预估 draft 占用。静态量化标定须在规划前加载和校验；KV、动态 scale、packed padding 及 HOLD 页计入预算，静态 KV/Q 标定参数占比较小，不在容量规划中单独扣减。按完整物理页选择可容纳的最大 token 容量，各 rank 再统一取最小容量。共享 allocator 的不同 State 存储分别计算，不能因共享索引而漏算 draft buffer。 + +固定配置下，`get_full_att_token_bytes()` 返回所有 full-attention State 的单 token 开销，页开销直接计算为 `page_bytes = token_bytes * page_size`。可分配页数为 `budget // page_bytes - 1`，其中扣除的一页用于 HOLD，再乘以 page size 得到 token 容量；容量查询接口不接收 token 容量或页数，也不通过两次总占用查询求差值。索引范围由 `FullAttKVAllocator` 的已有断言校验。 + +PD 容量规划使用 `manager.get_pd_page_bytes() * page_count` 计算预留空间，依赖方向是先布局、后页大小。FullAttentionState 的 `get_pd_transfer_regions()` 根据全局 `pd_kv_page_size` 返回实际字段区域，包含各 field 的 dtype 对齐及实现要求的内部 padding。manager 同时平移字段区间的起点与终点，以 `max(region.byte_end)` 确定当前 State 的数据结束位置，向上对齐到 8 字节后记录为 layout 的 `byte_end`，并作为下一个 State 的 `byte_start`。所有 State 排列完成后，`max(layout.byte_end)` 就是 KV 页大小。GPU packed 行中未被任何传输字段覆盖的 padding 不进入 PD 页。 + +bounded/indexer 的运行态布局使用同一规则：按字段的实际最大结束位置确定 State 占用，不另外维护纯数据字节数估算，也不累加多个映射共用的同一区域。它始终返回各运行态 State 自身的起止区间，不依赖 full-attention 布局，运行态总占用取 `max(layout.byte_end)`。`get_pd_page_bytes()` 返回 KV 页与完整运行态页占用的最大值；没有参与资源时返回 0。有 KV 且 KV 页小于运行态页时记录英文 warning，自动扩大实际传输页，保持 `pd_kv_page_size` 的 token 数和两类布局偏移不变，不再因两类页大小关系而中止初始化。PD 调用方按这个返回值分配和选取实际传输页，再以该页起点执行对应任务的布局；例如运行态布局为 80 字节、KV 页为 128 字节时，实际传输块为 128 字节，运行态仅操作自己的 80 字节范围。 + +PD master 负责校验连接参数，长度通过传输任务的 `end_kv_index` 传递；资源层不附加 PD 页布局指纹及长度、版本字段。两个布局 getter 与页大小查询均使用 `lru_cache`，存储初始化后一起清除缓存。 + +CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_transfer_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 + +State 的全局存储字节数查询继续描述存储开销，不用于推导 CPU/PD 传输容量。CPU 联合页不预留页头、长度或版本字段。布局及其派生页大小均按需缓存，实际存储分配后一起清理，避免继续使用 meta 准备阶段的布局尺寸。 + +State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_size` 及具体分布规则，本地数量使用只读属性计算，不保存第二份可独立修改的数值。CPU/PD 全局页和 DP 运行态映射直接读取对应 State 的 config，管理器不维护 `_full_head_counts`、`_bounded_head_counts`。每个 draft 的 State 独立保存完整 config，本地形状相同也不意味着全局布局相同，容量规划和传输均读取各自配置。 + +| 配置 | 保存的信息 | 计算属性 | +| --- | --- | --- | +| `FullAttentionConfig` | `global_head_num`、`replicated`、TP 坐标及原有层/维度/dtype 信息 | `head_num` 返回当前 rank 的 KV head 数 | +| `ConvSSMStateConfig` | 层、`global_k_heads`、`global_v_heads`、TP 坐标及 Conv/SSM 维度/dtype 信息；不保存资源身份 | `num_k_heads`、`num_v_heads` 分别返回本地 head 数，`conv_channels` 返回本地卷积通道数 | + +full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例如复制的 MLA latent；否则按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 + +`basemodel_utils.py` 提供 `get_resource_config(model_config, *, model_class=None)` 和 `get_full_attention_options(config, args, *, model_class=None)`，统一处理字段别名、text config 和默认 KV 维度。它们是普通工具函数，BaseModel 不再声明对应的 `_resource_config`、`_full_attention_options` 方法；具体模型仅在存在专有规则时提供可选 hook，并调用工具函数取得默认配置。传入 `model_class` 时，工具函数选择该模型的专有实现，否则使用通用逻辑。 + +同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域；磁盘缓存命名空间包含 main/draft 模型版本。静态量化标定按本模型实际 full-attention 层区间选取，CPU 服务和推理端复用相同的布局构建流程。 + +资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 + +### 7.2 缓存保存与恢复 + +没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)` 和 `reset_attention_tail_state(req_indices)`，分别初始化 bounded 与 tail 运行态。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 + +调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_pages`。具体管理器将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 + +radix 沿用原有判断：大页节点具有状态，小页节点通过 `small_page_buffer_idx is not None` 判断是否可恢复；图片边界仍通过 `clamp_mm_cache_len` 独立检查。匹配得到节点引用后，调用方直接调用 `restore_bounded_att_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_attention_tail_state(req_indices)` 清零 tail，保留已恢复的 linear 状态。调用方还需确认历史 KV、tail 和 bounded 快照均对应恢复前缀，并保证数据就绪。 + +使用 K-pool 的 tail 在完整池化边界上通过 `reset_attention_tail_state` 重置；非对齐恢复需要原始 tail 数据。CPU 读回后，管理器在最终前缀数据恢复后调用 `reset_attention_tail_state`，不读取 tail 快照。DP fetch 在对齐的缓存前缀上重置 tail；PD 按第 6.8 节传输固定逻辑窗口并重排到本地 ring。前缀恢复的清零规则不用于 MTP 回退。 + +MTP 接受位置由模型执行流程选择，供后续计算使用；当前大小页保存和 DP/PD 导出发生在 prefill 基准状态阶段。KDA/GDN 只保存实际捕获的状态；跨越但未保存的中间位置不可作为有效页，也不能发布为 CPU 缓存页；执行流程需在待缓存页边界安排真实状态捕获。管理器的数据操作不承担一次 forward 的提交或请求长度更新。 + +### 7.3 同步、引用与回收 + +请求流程负责请求暂停、结束及请求槽位回收；缓存功能流程负责节点引用和淘汰时机。归还资源索引前，调用方应停止相关使用、满足异步完成条件并解除相应引用。ReqManager 清理请求映射；对于状态页，调用方直接向对应 allocator 归还一次。 + +有缓存、请求或传输引用的数据必须保持有效。释放一份缓存数据不等于结束某个请求；请求结束也不表示其已共享缓存应被释放。调用方负责将具备回收条件的索引归还 allocator，管理器内部保留的页内容不参与引用计数。 + +保存、导入或复制失败时，不发布本次不完整数据；调用方在操作完成且引用解除后，直接回收本次不再使用的分配，已完成的共享缓存保持有效。CPU/DP/PD 操作遵守 stream/event 依赖，调用方在继续计算、复用 buffer 或回收索引前确认完成。 + +## 8. CPU、DP 与 PD 传输设计 + +### 8.1 具体 State 实现传输，模型管理器组织组合 + +CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合缓存辅助文件： + +- `read_req_from_cpu(req, cpu_page_indices, page_lengths, *, start_token_len)` 为新增 KV 分配索引，补齐首个 CPU 页已有命中部分，为完整 bounded 页分配大页槽位，再调用资源管理器 `read_req_from_cpu`。成功提交后登记请求的大页引用并更新 KV 长度；CPU cache 加载按必须成功的流程执行，异常直接向上传播，不在此处捕获、同步回收或继续执行。 +- `write_req_to_cpu(req, cpu_kv_cache_stream)` 先按 KV 长度和碎页规则申请 CPU 页；实际页数确定后，合并 radix 共享大页和请求自有页，仅查找本批页面所需槽位。完整大页必须存在，直接索引获取，缺失时暴露错误，不截断缓存前缀。将整批页面、完整 token 索引、大小页槽位和 `cpu_page_readies` 一次传给资源管理器 `write_req_to_cpu`，由各 State 跳过 ready 页面。返回含完成事件的 `TransTask`，由已有任务完成流程发布 ready 并释放引用。 + +写出入口仅通过 `cpu_kv_cache_stream.wait_stream(producer)` 等待当前生产流已提交的写入，不额外等待 overlap 流，也不对 hybrid 模型执行主机同步。TODO：bounded CPU→CPU 快照搬运仍需改为当前流上的 kernel；当前主机 copy 不受流等待约束，kernel 落地前源快照必须预先就绪。 + +命中边界、容量准入和 CPU 页引用释放仍由外层缓存调度负责。以上入口不把 `InferReq` 传入 `AttResourceManager`；资源层只接收显式索引与槽位，并组合各 State 的数据操作。 + +外部通过 6.2 节的五个管理器抽象方法执行 CPU/DP/PD 操作。具体管理器将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 + +- full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 +- `KDAState`、`GDNState` 实现 `BoundedAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移 prefill 基准请求状态；管理器组织参与对象及位置映射。 +- `AttentionTailState` 提供层访问、按请求清零及 PD 直接读写方法，详见 6.8 节；CPU cache 和 DP fetch 通过重置恢复对齐前缀的 tail。GLM K-pool indexer 在 TP ranks 间为副本，其他实现自行描述其布局。 +- 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttResourceManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 + +跨进程导入的管理器持有借用的运行态存储和索引容量描述，不持有 allocator 或 CPU 状态池。请求映射由 server 层作为 ReqManager 所有的独立存储导出，不能附加到资源管理器。CUDA IPC 句柄按实际读取进程数发布，每个读取者消费独立引用；原进程在全部读取者结束使用前保持存储存活。 + +### 8.2 PD + +PD 任务分为 `kv` 和 `att_state` 两类；涉及 main/draft 时按模型角色组合所需资源： + +| 任务 | 组成 | 处理方式 | +| --- | --- | --- | +| `kv` | full-attention 历史 KV，包含布局所需的 packed indexer 数据 | 按 token 页复制 | +| `att_state` | linear Conv/SSM + 可选 indexer tail | 模型专属 `AttResourceManager` 按请求恢复位置调用对应 State 方法 | + +P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，indexer tail 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将 tail 写入目标 TP 各 rank 的本地请求槽位。 + +模型专属 `AttResourceManager` 实现状态页的布局组织和容量校验，调用 linear State 与 indexer tail State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 + +网络层只接收固定容量的字节页，不按 KV 的 head 数或 dtype 推断组合布局。两端在创建传输句柄前核对页容量，PD master 在连接时校验任务类别、资源身份、算法与编码、层映射和全局字段形状；页内不附加格式指纹、长度或版本。布局使用全局字段描述以支持不同 TP 大小，tail 布局只描述固定逻辑窗口，不包含本地 MTP 候选容量。运行态的前缀长度由传输任务提供。 + +P/D 使用相同模型和 main/draft 资源集合，允许不同 TP 与 `mtp_step`。full-attention 传输有效 KV 和量化字段；bounded 状态按 6.7 节传输 prefill 基准 Conv/SSM，恢复到本地基准位置并将游标归零；tail 按 6.8 节传输固定逻辑窗口并按本地 ring 容量重排。MTP 候选容量不进入 PD 字段形状，单个 task 仍使用既有的前缀位置描述。 + +### 8.3 CPU/Disk 与 DP fetch + +CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布局保存数据，不预留格式页头或状态长度、版本字段。各资源区域由模型配置与 State 字段确定，TP 副本由明确的单个写入者更新共享区域。CPU 服务绑定时核对页大小，读回长度由任务参数提供，前缀与槽位的关联由调用方维护。磁盘文件按 main/draft 模型路径及文件版本隔离。 + +GPU→CPU 保存状态页完成后，才能在 CPU 上重排或复制这些快照；主机侧读取必须等待保存完成，不能只向另一个 CUDA stream 排入依赖。CPU 页的 ready 发布仍由调用流程在全部参与资源与 TP 分片传输完成后执行。 + +模型专属 `AttResourceManager` 实现 CPU 联合页的布局与读写组织，调用所需 main/draft State 的 CPU 方法，组合 full-attention KV 和 linear 页；按调用方提供的槽位写出或读回快照，pool 和 CPU buffer 均不保存长度副本。CPU 缓存调度提供外部页和任务生命周期；每个已发布 CPU 页都具备该模型组合所需的有效 linear 状态。一次读回处理同一前缀的页面，并通过 `reset_attention_tail_state` 清零目标 tail；恢复终点必须满足具体 tail 的可重置边界。 + +CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大页 pool 分别提供存储。各参与方的暂存数据消费完成后，后续请求才能复用该槽位。磁盘存储使用 CPU 页格式；main/draft 异构组合需定义各自资源段的位置、字段形状及 TP 切片。 + +DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并在对齐边界重置目标 tail。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 + +## 9. 待细化的设计问题 + +| 问题 | 设计约定 | 待确定内容 | +| --- | --- | --- | +| main/draft 管理器装配 | 模型组合实现自己的管理器；各角色和具体方案拥有唯一 State，同一分配域共享唯一 allocator 实例 | 新模型的 main/draft 装配时序及独立存储分配；调用方覆盖全部参与资源的引用解除及单次归还责任 | +| AttResourceManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、清零指定请求的 tail、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | +| 内部 attention 资源的暴露 | 具体模型在初始化阶段绑定 State 或层视图，不要求统一模型视图获取接口 | 各模型/backend 的具体绑定方式与按层访问结构 | +| FullAttentionState 接口与继承 | 必需语义：容量计算、初始化、层视图、KV 写入、完整 token 复制及固定字段传输布局；CPU/DP/PD 由子类按支持路径直接实现读写方法 | 新布局的 TensorField 覆盖范围、StateTransferLayout 映射约束及具体搬运实现 | +| BoundedAttentionState 接口与继承 | KDAState、GDNState 实现统一抽象操作；已接受运行态、checkpoint 和传输分别定义，allocator 仍归管理器 | 按 6.7 节细化各子类的描述字段和 kernel，核对 MTP 状态选择、跨 TP/CPU 格式兼容性 | +| AttentionTailState 接口与继承 | 按 6.8 节实现层访问、按请求清零及 PD 传输；prefill 逻辑窗口迁移，无独立页 pool | 各子类的恢复规则及跨 ring 布局转换能力 | +| allocator 能力描述 | 公共入口按需初始化，不需要的 allocator 为 None | 恢复边界和传输能力使用哪些最小描述，避免重复配置导致不一致 | +| checkpoint 引用形式 | 调用方维护大小页类别、共享槽位 ID、前缀长度及参与模型资源；各侧 pool 只保存快照数据 | 引用对象或参数组合的具体类型，以及参与资源集合的表达 | +| 配置类型拆分 | main/draft 分别产生实现和布局参数，共享 allocator 接收统一容量与索引配置 | 各类 State 的配置字段、跨资源共享字段与校验责任 | +| PD tail 的协议布局 | 固定 P 个逻辑位置，按任务前缀长度映射两端 ring；兼容不同 MTP 步数 | 保持相同模型资源集合与 prefill 状态交接约束 | +| CPU 缓存页的组合方式 | 每页包含全部有效 linear 状态；管理器组织读写，并在全部页面恢复后清零目标 tail | 不同 main/draft 实现的资源段布局及恢复边界兼容 | + +以上待细化事项须遵守本文的资源归属、共享分配域和恢复语义约定。 diff --git a/lightllm/common/att_resources/__init__.py b/lightllm/common/att_resources/__init__.py new file mode 100644 index 0000000000..428c8aa2a1 --- /dev/null +++ b/lightllm/common/att_resources/__init__.py @@ -0,0 +1,5 @@ +"""Attention data resources, independent of request and cache policy.""" + +from .att_resource_manager import AttResourceManager + +__all__ = ["AttResourceManager"] diff --git a/lightllm/common/att_resources/allocators/__init__.py b/lightllm/common/att_resources/allocators/__init__.py new file mode 100644 index 0000000000..036f71c997 --- /dev/null +++ b/lightllm/common/att_resources/allocators/__init__.py @@ -0,0 +1,4 @@ +from .state_page_allocator import StatePageAllocator +from .full_att_kv_allocator import KVCapacityReader, FullAttKVAllocator + +__all__ = ["StatePageAllocator", "KVCapacityReader", "FullAttKVAllocator"] diff --git a/lightllm/common/att_resources/allocators/full_att_kv_allocator.py b/lightllm/common/att_resources/allocators/full_att_kv_allocator.py new file mode 100644 index 0000000000..8bb62d5fa6 --- /dev/null +++ b/lightllm/common/att_resources/allocators/full_att_kv_allocator.py @@ -0,0 +1,100 @@ +"""Full attention KV 分配器：接口使用 token 数量和索引，分配与回收须满足整页约束。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import numpy as np +import torch + +from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt +from lightllm.utils.dist_utils import get_current_rank_in_node +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +if TYPE_CHECKING: + from collections.abc import Sequence + + +logger = init_logger(__name__) + + +class FullAttKVAllocator: + def __init__(self, size: int, page_size: int = 1, *, pin_memory: bool = True): + assert size >= 0 and page_size > 0 and size % page_size == 0, "KV capacity must contain complete pages" + assert size + page_size - 1 <= np.iinfo(np.int32).max, "KV token indices including HOLD must fit in int32" + self.size = size + self.page_size = page_size + self.mem_state = torch.arange(self.size, dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self.mark_start = 0 + self.mark_end = self.size + self._mem_state_return = torch.empty(self.size * 3, dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self._return_start = 0 + self.can_use_mem_size = self.size + rank_in_node = get_current_rank_in_node() + self.shared_can_use_token_num = SharedInt(f"kv_available_tokens_{rank_in_node}") + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + @property + def hold_token_indices(self) -> tuple[int, ...]: + """One reserved page after the ordinary allocation domain.""" + return tuple(range(self.size, self.size + self.page_size)) + + def alloc(self, need_size: int) -> torch.Tensor: + assert need_size >= 0 and need_size % self.page_size == 0, "allocation must contain complete KV pages" + assert need_size <= self.mark_end - self.mark_start, ( + f"need {need_size} KV tokens, only {self.can_use_mem_size} available" + ) + start = self.mark_start + end = start + need_size + self.mark_start = end + self.can_use_mem_size -= need_size + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + # Reuse a rolling buffer so frees do not overwrite indices + # still being consumed by asynchronous H2D copies. + if self._return_start + need_size > self._mem_state_return.shape[0]: + self._return_start = 0 + ans = self._mem_state_return[self._return_start : self._return_start + need_size] + ans.copy_(self.mem_state[start:end]) + self._return_start += need_size + return ans + + def free(self, token_indices: Sequence[int] | torch.Tensor) -> None: + num_tokens = len(token_indices) + if num_tokens == 0: + return + assert num_tokens % self.page_size == 0, "free must contain complete KV pages" + end = self.mark_start + start = end - num_tokens + assert start >= 0, f"free {num_tokens} KV tokens exceeds allocated {end}" + if isinstance(token_indices, torch.Tensor): + self.mem_state[start:end].copy_(token_indices) + else: + self.mem_state.numpy()[start:end] = token_indices + self.mark_start = start + self.can_use_mem_size += num_tokens + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.can_use_mem_size == self.size: + logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") + + def free_all(self) -> None: + """Reset after all users of this allocation domain have finished.""" + self.mem_state.numpy()[:] = np.arange(self.size, dtype=np.int32) + self.mark_start = 0 + self.mark_end = self.size + self.can_use_mem_size = self.size + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + +class KVCapacityReader: + """Read each local DP group's allocator capacity without modifying it.""" + + def __init__(self): + args = get_env_start_args() + node_world_size = args.tp // args.nnodes + tp_group_size = args.tp // args.dp + self.is_multinode_tp = args.dp == 1 and args.nnodes > 1 + self.counters = [SharedInt(f"kv_available_tokens_{rank}") for rank in range(0, node_world_size, tp_group_size)] + + def get_available_tokens(self, dp_rank_in_node: int) -> int: + return self.counters[0 if self.is_multinode_tp else dp_rank_in_node].get_value() diff --git a/lightllm/common/att_resources/allocators/state_page_allocator.py b/lightllm/common/att_resources/allocators/state_page_allocator.py new file mode 100644 index 0000000000..cfc58ff9c2 --- /dev/null +++ b/lightllm/common/att_resources/allocators/state_page_allocator.py @@ -0,0 +1,39 @@ +"""状态页槽位分配器，仅管理槽位，不持有状态 buffer 或缓存策略。""" + +from __future__ import annotations + +from typing import List +from collections import deque + + +class StatePageAllocator: + def __init__(self, size: int, *, reserved_slots: int = 0): + self.size = size + self.reserved_slots = reserved_slots + if not 0 <= self.reserved_slots <= self.size: + raise ValueError("reserved_slots must be between zero and capacity") + self.capacity = self.size - self.reserved_slots + self._free = deque(range(self.capacity)) + self._allocated = bytearray(self.capacity) + + @property + def free_count(self) -> int: + return len(self._free) + + def alloc(self, num_pages: int) -> list[int] | None: + assert num_pages >= 0, "num_pages must be nonnegative" + if num_pages > self.free_count: + return None + slots = [self._free.popleft() for _ in range(num_pages)] + for slot in slots: + self._allocated[slot] = 1 + return slots + + def free(self, slot_ids: List[int]) -> None: + assert isinstance(slot_ids, list), "slot indices must be a list" + assert len(set(slot_ids)) == len(slot_ids), "cannot free duplicate slots" + assert all(0 <= slot < self.capacity for slot in slot_ids), "slot outside the allocatable domain" + assert all(self._allocated[slot] for slot in slot_ids), "cannot free an unallocated slot" + for slot in slot_ids: + self._allocated[slot] = 0 + self._free.extend(slot_ids) diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/att_resources/att_resource_manager.py new file mode 100644 index 0000000000..d242434ae3 --- /dev/null +++ b/lightllm/common/att_resources/att_resource_manager.py @@ -0,0 +1,711 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from collections.abc import Mapping +from copy import copy +from dataclasses import replace +from functools import lru_cache +from typing import List, Literal + +import torch + +from lightllm.common.att_resources.types import ( + StateTransferLayout, + PageKind, +) + +from lightllm.utils.dist_utils import get_current_device_id +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from .allocators import StatePageAllocator, FullAttKVAllocator +from .full_att import FullAttentionState +from .bounded_att import BoundedAttentionState +from .att_tail import AttentionTailState +from .utils import MODEL_ROLES + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +logger = init_logger(__name__) + + +class AttResourceManager: + """管理模型组合的注意力资源，提供缓存复用及 CPU/DP/PD 传输所需的聚合操作。 + + 默认通过聚合查询操作各类 State；模型子类仅在资源创建或操作有差异时扩展。 + 未支持的传输路径必须明确报错。 + """ + + # ======================================================================== + # 一、基础管理：共享分配器、资源初始化与聚合查询 + # ======================================================================== + + # 共享分配器：按模型需要启用；调用方直接分配和归还索引。 + full_att_allocator: FullAttKVAllocator | None + bounded_att_big_page_allocator: StatePageAllocator | None + bounded_att_small_page_allocator: StatePageAllocator | None + _states_created: bool = False + + def __init__(self, request_capacity: int): + """记录共享请求容量；State 和 allocator 由后续初始化流程设置。""" + self.full_att_allocator = None + self.bounded_att_big_page_allocator = None + self.bounded_att_small_page_allocator = None + self._request_capacity = request_capacity + self.cpu_tail_restore = False + + def set_att_state( + self, model_name: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState + ) -> None: + """按模型角色和具体方案绑定已创建的 State，保留原对象及其存储。""" + assert not self._states_created, "State creation is already finished" + assert model_name in MODEL_ROLES, f"unknown model role {model_name}" + state_type = state.get_state_type() + assert state_type.isidentifier(), f"invalid State type {state_type}" + name = f"{model_name}_{state_type}_state" + assert name not in self.__dict__, f"State {name} is already bound" + assert all(state is not value for value in self.__dict__.values()), "a State cannot belong to multiple model roles" + # config 只描述布局，资源身份仅保存在 State；绑定时保留原配置对象。 + state.resource_id = name + setattr(self, name, state) + + def initialize_resources( + self, + args: StartArgs | EasyDict, + *, + token_capacity: int | None, + mem_fraction: float, + device: str | torch.device | int = "cuda", + ) -> None: + """全部模型完成资源准备后，规划共享容量并初始化缓存存储。 + + 调用前，所有模型的权重和绑定请求容量的运行态已创建，full-attention + 仅持有 meta 定义，大小页池尚未分配。调用方在本方法返回后， + 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 + """ + # 阶段一：标记全部 State 对象创建完成，后续通过聚合接口访问。 + # main、draft0、draft1 等模型已分别将自己的 State 添加到当前管理器。 + self.finish_state_creation() + + # 阶段二：规划共享 token 容量，并向下对齐到完整 KV 页。 + # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 + # PD 模式还需为传输 buffer 预留空间，布局按全部 full-attention State 计算。 + assert self.get_full_attention_states(), "the shared token index domain requires a full-attention State" + pd_reserved_bytes = 0 + if args.run_mode in ("prefill", "decode"): + pd_reserved_bytes = self.get_pd_page_bytes() * args.pd_kv_page_num + if token_capacity is None: + # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。 + torch.cuda.empty_cache() + gpu_device = torch.device("cuda", get_current_device_id()) + free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_device) + budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes + + # 按每页开销直接计算容量，预留一个 HOLD 页；静态量化参数占比较小，不单独扣减。 + assert args.page_size > 0, "capacity planning requires a positive page size" + page_bytes = self.get_full_att_token_bytes() * args.page_size + if page_bytes > budget: + raise MemoryError("attention storage budget cannot fit HOLD pages") + token_capacity = (budget // page_bytes - 1) * args.page_size + + # 各 rank 的显存余量可能不同，统一采用最小容量。 + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + capacity = torch.tensor(token_capacity, dtype=torch.int64, device=gpu_device) + torch.distributed.all_reduce(capacity, op=torch.distributed.ReduceOp.MIN) + token_capacity = capacity.item() + logger.info( + f"Attention storage budget: {budget / 1024**3:.3f} GiB, " + f"transfer reserve: {pd_reserved_bytes / 1024**2:.3f} MiB, " + f"profiled token capacity: {token_capacity} (mem_fraction={mem_fraction})" + ) + token_capacity = token_capacity // args.page_size * args.page_size + assert token_capacity > 0, "attention capacity must contain at least one complete page" + + # 阶段三：管理器为已创建的 State 分配实际存储与共享 allocator。 + # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 + big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + self.initialize_storage( + token_capacity, + args.page_size, + big_page_tokens=big_page_tokens, + small_page_capacity=args.linear_att_cache_size or 0, + device=device, + ) + + # 创建完成只表示对象集合就绪；实际容量和存储在后续阶段初始化。 + def finish_state_creation(self) -> None: + """确认全部模型的 State 对象已创建并装配完成，允许查询 meta 定义并开始容量规划。""" + assert not self._states_created, "State creation is already finished" + self._states_created = True + + def initialize_storage( + self, + token_capacity: int, + page_size: int, + *, + big_page_tokens: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cuda", + ) -> None: + """在 State 对象创建完成后分配 full-attention 存储、共享 allocator 和 bounded 大小页。 + + 请求容量对应的 bounded/indexer 运行态已由各模型创建，这里仅补齐缓存存储。 + """ + assert self._states_created, "State creation must finish before storage initialization" + assert self.full_att_allocator is None, "resources already initialized" + full_states = self.get_full_attention_states() + assert full_states, "the shared token index domain requires a full-attention State" + for state in full_states: + # 管理器统一设置实际页规格和设备；State 的分配接口只接收 token 容量。 + state.page_size = page_size + state.device = torch.device(device) + state.initialize_storage(token_capacity) + self.full_att_allocator = FullAttKVAllocator(token_capacity, page_size) + + # bounded-attention 大页池末尾保留 CPU 加载/卸载暂存槽位，普通分配不使用它们。 + bounded_states = self.get_bounded_attention_states() + if bounded_states: + assert big_page_tokens > 0, "bounded attention pages require a positive token boundary" + big_page_capacity = (token_capacity + big_page_tokens - 1) // big_page_tokens + 2 + self.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) + self.bounded_att_small_page_allocator = ( + StatePageAllocator(small_page_capacity) if small_page_capacity else None + ) + for state in bounded_states: + state.initialize_page_pools( + big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, + pin_memory=torch.device(device).type == "cuda", + ) + self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 + self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 + + self.cpu_tail_restore = bool(self.get_attention_tail_states()) + # 存储分配可能改变 State 的物理页规格，清除准备阶段的布局及其派生 CPU/PD 页大小缓存。 + self.get_cpu_transfer_layouts.cache_clear() + self.get_cpu_page_bytes.cache_clear() + self.get_pd_full_att_transfer_layouts.cache_clear() + self.get_pd_bounded_att_or_tail_transfer_layouts.cache_clear() + self.get_pd_page_bytes.cache_clear() + + def create_pd_ipc_copy(self) -> AttResourceManager: + """创建 PD 导出副本，保留运行态 Tensor，排除 allocator 和全部大小页池。""" + ipc_copy = copy(self) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + + # PD 只传运行态；先复制 State 对象再移除页池，避免修改推理进程的资源。 + for state in self.get_bounded_attention_states(): + state_copy = copy(state) + state_copy.big_pages_pool = None + state_copy.small_pages_pool = None + setattr(ipc_copy, state.resource_id, state_copy) + return ipc_copy + + def create_dp_ipc_copy(self) -> AttResourceManager: + """创建 DP shared 导出副本，保留运行态和共享大小页池，仅排除 allocator。""" + # DP fetch 初始化已选择共享 CPU 页池；禁止导出私有页池时隐式分配、复制整块存储。 + for state in self.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + assert ( + pool.conv_state.device.type == "cpu" and pool.conv_state.is_shared() + and pool.ssm_state.device.type == "cpu" and pool.ssm_state.is_shared() + ), "DP IPC requires shared CPU page pools" + ipc_copy = copy(self) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + return ipc_copy + + def register_page_pools_pin_memory(self) -> None: + """在当前 GPU 上统一注册所有 bounded State 的大小页池，供 DP 接收端调用。 + + 调用方先选择本进程设备；注册状态不随 IPC 传递，重复调用会跳过已注册存储。 + 本接口只注册本地映射，不申请槽位或改变原进程的 allocator 所有权。 + """ + for state in self.get_bounded_attention_states(): + if state.big_pages_pool is not None: + state.big_pages_pool.register_pin_memory() + if state.small_pages_pool is not None: + state.small_pages_pool.register_pin_memory() + + def get_full_attention_states(self, role: str | None = None) -> tuple[FullAttentionState, ...]: + """返回全部或指定角色的 full-attention State；只缓存对象引用。""" + return self._get_states(FullAttentionState, role) + + def get_bounded_attention_states(self, role: str | None = None) -> tuple[BoundedAttentionState, ...]: + """返回全部或指定角色的 bounded-attention State。""" + return self._get_states(BoundedAttentionState, role) + + def get_attention_tail_states(self, role: str | None = None) -> tuple[AttentionTailState, ...]: + """返回全部或指定角色的 attention-tail State。""" + return self._get_states(AttentionTailState, role) + + def get_full_att_token_bytes(self) -> int: + """累计当前 TP rank 全部 full-attention State 的单 token 存储字节数。""" + total = 0 + for state in self.get_full_attention_states(): + total += state.get_token_bytes() + return total + + @lru_cache(maxsize=None) + def get_cpu_page_bytes(self) -> int: + """从 CPU 联合页布局的最大终点计算物理页大小，包含字段及 State 间的对齐空间。""" + page_bytes = 0 + for layout in self.get_cpu_transfer_layouts().values(): + page_bytes = max(page_bytes, layout.byte_end) + # State 区间按 8 字节对齐，整个 CPU 物理页再按 16 字节对齐。 + return (page_bytes + 15) // 16 * 16 + + @lru_cache(maxsize=None) + def get_pd_page_bytes(self) -> int: + """取 KV 页与完整运行态页占用的最大值，保证两类 PD 任务都能容纳。""" + full_layouts = self.get_pd_full_att_transfer_layouts() + runtime_layouts = self.get_pd_bounded_att_or_tail_transfer_layouts() + full_att_page_bytes = 0 + for layout in full_layouts.values(): + full_att_page_bytes = max(full_att_page_bytes, layout.byte_end) + runtime_page_bytes = 0 + for layout in runtime_layouts.values(): + runtime_page_bytes = max(runtime_page_bytes, layout.byte_end) + + # KV 页不足时扩大物理传输页，不改变 KV 页的 token 数或各 State 的布局。 + if full_layouts and full_att_page_bytes < runtime_page_bytes: + logger.warning( + f"PD KV page size ({full_att_page_bytes} bytes) is smaller than the complete attention runtime " + f"({runtime_page_bytes} bytes); using {runtime_page_bytes} bytes per transfer page" + ) + # 纯 KV 或纯运行态取各自占用；没有任何参与资源时返回 0。 + return max(full_att_page_bytes, runtime_page_bytes) + + # ======================================================================== + # 二、Full attention：token 数据复制 + # ======================================================================== + + # Full attention：在所有参与 State 内复制相同的物理 token 索引，不修改请求映射。 + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制 KV 及随 token 保存的量化参数、indexer 字段等完整 token 数据。""" + full_states = self.get_full_attention_states() + assert full_states, "no full-attention storage in this composition" + for state in full_states: + state.copy_full_att_tokens(src_indices, dst_indices) + + # ======================================================================== + # 三、Bounded attention：运行态重置、大小页保存与恢复、MTP 提交 + # ======================================================================== + + # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页。""" + for state in self.get_bounded_attention_states(): + state.reset_bounded_att_state(req_indices) + + # Bounded attention 状态页：供 radix/CPU cache 复用；未启用对应页池时操作报错。 + # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 + def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """将 prefill 运行态保存到全部参与 State 的同类页池、同一组槽位。""" + bounded_states = self.get_bounded_attention_states() + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator + assert allocator is not None, f"{page_kind} pages are disabled" + assert len(slot_ids) == len(req_indices), "one slot is required per request" + for state in bounded_states: + state.save_bounded_att_pages(req_indices, page_kind=page_kind, slot_ids=slot_ids) + + def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """将调用方选定的页快照恢复到对应请求槽位的运行态。""" + bounded_states = self.get_bounded_attention_states() + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator + assert allocator is not None, f"{page_kind} pages are disabled" + assert len(slot_ids) == len(req_indices), "one slot is required per request" + assert not slot_ids or bounded_states, "state restore requires participating bounded states" + for state in bounded_states: + state.restore_bounded_att_pages(req_indices, page_kind=page_kind, slot_ids=slot_ids) + + # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + *, + role: str, + ) -> None: + """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + assert role in MODEL_ROLES, f"unknown model role {role}" + for state in self.get_bounded_attention_states(role): + state.commit_bounded_att_state(req_indices, accepted_steps) + + # ======================================================================== + # 四、Attention tail：空前缀初始化与无快照恢复 + # ======================================================================== + + # 没有 tail 快照时使用；长度 0 用于空前缀初始化，没有 tail 的模型可为空操作。 + def reset_attention_tail_state( + self, + req_indices: list[int], + ) -> None: + """依次清零指定请求的 tail，恢复边界由模型配置和缓存调用流程保证。 + + 保留已恢复的 KV 和 bounded 运行态，不提供跨 State 的原子性保证。 + """ + for state in self.get_attention_tail_states(): + state.reset_attention_tail_state(req_indices) + + # ======================================================================== + # 五、CPU Cache 数据交换 + # ======================================================================== + + @lru_cache(maxsize=None) + def get_cpu_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: + """聚合各 State 的 CPU 字段区间，返回按 resource_id 索引的联合页布局。 + + 使用全局 cpu_cache_token_page_size，按需缓存；不依赖页大小查询。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + assert page_tokens > 0, "CPU pages require a positive token capacity" + states = (*self.get_full_attention_states(), *self.get_bounded_attention_states()) + assert states, "a CPU page requires participating storage" + + # 按 full-attention、bounded 顺序组织全部角色；各类 State 内部按资源名排序。 + # Indexer tail 沿用前缀边界重置，不加入当前 CPU 联合页。 + layouts, offset = {}, 0 + for state in states: + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_cpu_transfer_regions() + ) + # 以实际字段终点确定 State 占用，保留内部 padding，并对齐下一个 State 的起点。 + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.resource_id] = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + # CPU cache：按 layouts 组织联合页；混合模型的每页都包含全部参与 bounded State 的快照。 + # 纯 full-attention 模型只传 KV；具体实现负责校验格式、长度及各资源的区域映射。 + def write_req_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layouts: Mapping[str, StateTransferLayout], + cpu_page_readies: List[bool] | None = None, + big_page_slot_ids: List[int] | None = None, + small_page_slot_ids: List[int] | None = None, + ) -> None: + """将单个请求本次指定前缀页面的 KV 和 bounded 快照写入 CPU cache。 + + cpu_page_indices 前段对应 big_page_slot_ids，后段对应 small_page_slot_ids， + 各段按列表顺序一一对应；未参与的页池可传 None 或空列表。 + cpu_page_readies 与 CPU 页一一对应，True 跳过整页；None 表示全部写出。 + """ + if cpu_page_readies is not None: + assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" + if all(cpu_page_readies): + return + full_states = self.get_full_attention_states() + bounded_states = self.get_bounded_attention_states() + if bounded_states: + big_page_slot_ids = [] if big_page_slot_ids is None else big_page_slot_ids + small_page_slot_ids = [] if small_page_slot_ids is None else small_page_slot_ids + big_page_count = len(big_page_slot_ids) + assert big_page_count + len(small_page_slot_ids) == len(cpu_page_indices), "one state slot is required per CPU page" + big_cpu_page_indices = cpu_page_indices[:big_page_count] + small_cpu_page_indices = cpu_page_indices[big_page_count:] + big_page_readies = cpu_page_readies[:big_page_count] if cpu_page_readies is not None else None + small_page_readies = cpu_page_readies[big_page_count:] if cpu_page_readies is not None else None + + # 阶段一：统一写出本批 full-attention token 数据。 + for state in full_states: + state.write_to_cpu( + token_indices, cpu_page_buffers, cpu_page_indices, + layout=layouts[state.resource_id], cpu_page_readies=cpu_page_readies, + ) + + # 阶段二：分别从大小页池取快照，写入对应 CPU 页面，不修改请求运行态。 + for state in bounded_states: + layout = layouts[state.resource_id] + if big_page_slot_ids: + state.write_page_state_to_cpu( + page_kind="big", slot_ids=big_page_slot_ids, cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=big_cpu_page_indices, layout=layout, cpu_page_readies=big_page_readies, + ) + if small_page_slot_ids: + state.write_page_state_to_cpu( + page_kind="small", slot_ids=small_page_slot_ids, cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=small_cpu_page_indices, layout=layout, cpu_page_readies=small_page_readies, + ) + + def read_req_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layouts: Mapping[str, StateTransferLayout], + big_page_slot_ids: List[int] | None = None, + req_idx: int | None = None, + ) -> None: + """读回单个请求本次指定的有序前缀页面，并从末页恢复 req_idx 的 bounded 运行态。 + + big_page_slot_ids 只包含完整 CPU 页对应的普通槽位；末尾碎页使用内部预留的 LOAD 大页。 + 存在 bounded 或 tail 时必须提供 req_idx；调用方保证最终前缀位于 tail 可重置边界。 + """ + full_states = self.get_full_attention_states() + bounded_states = self.get_bounded_attention_states() + tail_states = self.get_attention_tail_states() + if bounded_states or tail_states: + assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" + if bounded_states: + # token 索引只用于判断本批末页是否完整，不读取 GPU 上的索引值。 + page_tokens = layouts[bounded_states[0].resource_id].tokens_per_page + has_partial_page = token_indices.numel() % page_tokens != 0 + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), "slots must cover complete CPU pages only" + # 只扩展内部列表，预留槽位不写回调用方，也不作为普通大页发布。 + load_big_page_slot_ids = big_page_slot_ids.copy() + if has_partial_page: + load_big_page_slot_ids.append(self.bounded_att_cpu_load_big_page_slot) + + # 阶段一：恢复 full-attention token 数据；各 State 负责对应布局的搬运约束。 + for state in full_states: + state.read_from_cpu(cpu_page_buffers, cpu_page_indices, token_indices, layout=layouts[state.resource_id]) + + # 阶段二：完整页和末尾碎页统一读入大页池，碎页使用追加的 LOAD 槽位。 + for state in bounded_states: + state.read_page_state_from_cpu( + cpu_page_buffers, cpu_page_indices, big_page_slot_ids=load_big_page_slot_ids, + layout=layouts[state.resource_id], + ) + + # 阶段三：从末页恢复请求运行态；纯碎页也走 LOAD 槽位。 + if bounded_states: + self.restore_bounded_att_pages([req_idx], page_kind="big", slot_ids=[load_big_page_slot_ids[-1]]) + + # 阶段四：清零目标请求的 tail；调用方保证前缀终点满足重置边界。 + for state in tail_states: + state.reset_attention_tail_state([req_idx]) + + # ======================================================================== + # 六、DP 迁移 + # ======================================================================== + + def copy_from_other_dp( + self, + remote_att_manager: AttResourceManager, + src_token_indexes: torch.Tensor, + dst_token_indexes: torch.Tensor, + src_big_page_ids: List[int], + dst_big_page_ids: List[int], + src_req_idx: int, + dst_req_idx: int, + ) -> None: + """复制其他 DP 同一 TP 分片的 token、大页和指定请求的 prefill 基准运行态。 + + token 索引由调用方准备在源、目标各自的 CUDA 设备上，目标索引为 int64。 + 不需要复制 token 或大页时传空 Tensor 或空列表；请求运行态仍按请求索引复制。 + 调用方保证源 MTP 游标为 0、目标前缀位于 tail 可重置边界，以及源页就绪和引用有效。 + """ + assert type(remote_att_manager) is type(self), "DP replicas must use the same model resource manager" + assert src_token_indexes.numel() == dst_token_indexes.numel(), "source and destination token counts differ" + assert len(src_big_page_ids) == len(dst_big_page_ids), "source and destination page counts differ" + if src_token_indexes.numel(): + for state in self.get_full_attention_states(): + state.copy_tokens_from_other_dp( + src_token_indices=src_token_indexes, + dst_token_indices=dst_token_indexes, + remote_dp_full_att_state=getattr(remote_att_manager, state.resource_id), + ) + if src_big_page_ids: + for state in self.get_bounded_attention_states(): + state.copy_big_page_state_from_other_dp( + src_page_ids=src_big_page_ids, + dst_page_ids=dst_big_page_ids, + remote_dp_bounded_state=getattr(remote_att_manager, state.resource_id), + ) + # 运行态不由大页复制隐式恢复；按指定请求复制基准位置，随后重置目标 tail。 + for state in self.get_bounded_attention_states(): + state.copy_run_time_state_from_other_dp( + src_req_indices=[src_req_idx], + dst_req_indices=[dst_req_idx], + remote_dp_bounded_state=getattr(remote_att_manager, state.resource_id), + ) + self.reset_attention_tail_state([dst_req_idx]) + + # ======================================================================== + # 七、PD 分离传输 + # ======================================================================== + + @lru_cache(maxsize=None) + def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: + """按实际字段区间生成 kv 布局，各 State 记录自身在页内的起点与终点。 + + 使用全局 pd_kv_page_size，按需缓存；不依赖页大小查询或运行态布局。 + """ + page_tokens = get_env_start_args().pd_kv_page_size + assert page_tokens > 0, "PD pages require a positive token capacity" + layouts, offset = {}, 0 + for state in self.get_full_attention_states(): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + # 字段保留自身数据区间，State 终点包含留给下一个资源的对齐空间。 + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.resource_id] = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: + """独立生成 att_state 布局，记录各运行态 State 的完整字节区间。 + + 每个 State 保存一份运行态，tokens_per_page 为 None；无 State 时返回空字典。 + """ + layouts, offset = {}, 0 + for state in (*self.get_bounded_attention_states(), *self.get_attention_tail_states()): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.resource_id] = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + # PD 分离：kind="kv" 传输 token 数据，kind="att_state" 传输 bounded 运行态及 tail。 + # pd_page_buffers 第一维为页，pd_page_index 选择物理页；网络收发与同步由传输流程组织。 + def write_to_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + """按任务布局写出 PD 页；含 tail 时传入 prefill 前缀长度 state_token_len。""" + match kind: + case "kv": + # KV 页只写出 full-attention 数据；各 State 负责自身布局与容量约束。 + assert token_indices is not None, "KV PD pages require token indices" + for state in self.get_full_attention_states(): + state.write_to_pd_buffer( + token_indices, pd_page_buffers, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + ) + + case "att_state": + # 提交搬运前确认请求位置；tail 额外使用前缀长度定位逻辑窗口。 + assert req_idx is not None, "state PD pages require a request index" + tail_states = self.get_attention_tail_states() + if tail_states: + assert state_token_len is not None and state_token_len >= 0, "PD tail requires its prefill prefix length" + + # 阶段一：写出 bounded 的 prefill 基准运行态。 + for state in self.get_bounded_attention_states(): + state.write_to_pd_buffer( + req_idx, pd_page_buffers, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + ) + + # 阶段二：按前缀长度写出 tail 逻辑窗口，适配两端各自的 MTP ring。 + for state in tail_states: + state.write_to_pd_buffer( + req_idx, pd_page_buffers, pd_page_index=pd_page_index, state_token_len=state_token_len, + layout=layouts[state.resource_id], + ) + + case _: + raise AssertionError(f"unknown PD page kind {kind}") + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + """按布局恢复 PD 页；含 tail 时使用 state_token_len 适配本地 MTP ring。""" + match kind: + case "kv": + # KV 页只恢复 full-attention 数据；各 State 负责自身布局与容量约束。 + assert token_indices is not None, "KV PD pages require token indices" + for state in self.get_full_attention_states(): + state.read_from_pd_buffer( + pd_page_buffers, token_indices, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + ) + + case "att_state": + # 提交搬运前确认请求位置；tail 额外使用前缀长度定位逻辑窗口。 + assert req_idx is not None, "state PD pages require a request index" + tail_states = self.get_attention_tail_states() + if tail_states: + assert state_token_len is not None and state_token_len >= 0, "PD tail requires its prefill prefix length" + + # 阶段一：恢复 bounded 的 prefill 基准运行态。 + for state in self.get_bounded_attention_states(): + state.read_from_pd_buffer( + pd_page_buffers, req_idx, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + ) + + # 阶段二：按前缀长度恢复 tail 逻辑窗口，适配两端各自的 MTP ring。 + for state in tail_states: + state.read_from_pd_buffer( + pd_page_buffers, req_idx, pd_page_index=pd_page_index, state_token_len=state_token_len, + layout=layouts[state.resource_id], + ) + + case _: + raise AssertionError(f"unknown PD page kind {kind}") + + # ======================================================================== + # 八、私有辅助实现 + # ======================================================================== + + # 私有实现：缓存创建完成后的成员集合，模型逐层计算直接访问具体 State 成员。 + @lru_cache(maxsize=None) + def _get_states( + self, + state_class: type[FullAttentionState] | type[BoundedAttentionState] | type[AttentionTailState], + role: str | None, + ): + """按 State 类型和模型角色遍历规范成员，只缓存对象引用。""" + assert self._states_created, "State creation must finish before querying" + assert role is None or role in MODEL_ROLES, f"unknown model role {role}" + states = ( + state + for state in self.__dict__.values() + if isinstance(state, state_class) and (role is None or state.resource_id.startswith(role + "_")) + ) + return tuple(sorted(states, key=lambda state: state.resource_id)) diff --git a/lightllm/common/att_resources/att_tail/__init__.py b/lightllm/common/att_resources/att_tail/__init__.py new file mode 100644 index 0000000000..05d1c630e1 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/__init__.py @@ -0,0 +1,5 @@ +from .create_kpool_state import create_kpool_state +from .base_state import AttentionTailState +from .kpool_state import KPoolAttentionTailState, KPoolStateConfig + +__all__ = ["create_kpool_state", "AttentionTailState", "KPoolAttentionTailState", "KPoolStateConfig"] diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/att_resources/att_tail/base_state.py new file mode 100644 index 0000000000..1e65380fb6 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/base_state.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any +from abc import ABC, abstractmethod + +from lightllm.common.att_resources.types import ( + StateTransferLayout, + TransferRegion, +) + +if TYPE_CHECKING: + import torch + + +class AttentionTailState(ABC): + """注意力相关模块的尾部状态接口，保存增量计算所需的尾部数据及控制信息。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + @abstractmethod + def get_state_type(self) -> str: + """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + ... + + # ======================================================================== + # 二、Attention tail：计算视图与无快照重置 + # ======================================================================== + + @abstractmethod + def get_layer_view(self, layer_id: int) -> Any: + """按逻辑层号返回尾部状态计算所需的存储视图。""" + ... + + @abstractmethod + def reset_attention_tail_state( + self, + req_indices: list[int], + ) -> None: + """清零指定请求的 tail；调用方保证用于空前缀或可重置的缓存边界，不用于 MTP 回滚。""" + ... + + # ======================================================================== + # 三、PD 分离传输：prefill 逻辑窗口 + # ======================================================================== + + @abstractmethod + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回独立于本地 MTP 容量的 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" + ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, + *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: + """按 state_token_len 对应的 prefill 终点导出逻辑窗口，不导出推测候选历史。""" + ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, + *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: + """按 state_token_len 将逻辑窗口恢复到本地运行态，适配本地 MTP 容量。""" + ... diff --git a/lightllm/common/att_resources/att_tail/create_kpool_state.py b/lightllm/common/att_resources/att_tail/create_kpool_state.py new file mode 100644 index 0000000000..28e87d3235 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/create_kpool_state.py @@ -0,0 +1,18 @@ +"""K-pool attention tail State 的统一创建入口。""" + +from __future__ import annotations + +import torch + +from .kpool_state import KPoolAttentionTailState, KPoolStateConfig + + +def create_kpool_state( + config: KPoolStateConfig, + request_capacity: int, + mtp_step: int = 0, + *, + device: str | torch.device = "cuda", +) -> KPoolAttentionTailState: + """按池化配置和请求容量创建 K-pool 尾部运行态，包含 MTP 额外槽位。""" + return KPoolAttentionTailState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py new file mode 100644 index 0000000000..4106623a67 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import dataclass + +import torch + +from ..types import AxisSlice, TransferRegion +from .base_state import AttentionTailState +from .transfer_utils import att_tail_buffer_copies + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +@dataclass(frozen=True) +class KPoolStateConfig: + # 本 State 管理的逻辑层号；顺序对应 tail_buffer 的 layer 轴，不要求编号连续。 + layer_ids: tuple[int, ...] + # K-pool 的池化步长/压缩比例，单位为原始 token；例如 4 表示每 4 个 token 形成一个池化组。 + # 无快照前缀恢复要求长度对齐该步长；ring 的实际槽位数为 pool_size + State.mtp_step。 + pool_size: int + # 单个原始 indexer key 的特征维度，不是 head 数，也不包含 gate 部分。 + # ring 的 channel 轴拼接 key 和同维度的 gate，共 2 * head_dim 个元素。 + head_dim: int + # 原始 key/gate ring 的存储精度，不代表池化后 packed indexer key 的量化编码类型。 + dtype: torch.dtype + # 当前模型 TP 组内的 rank。K-pool tail 在组内完整复制，传输时仅 rank 0 写出。 + tp_rank: int = 0 + # 当前模型的 TP 组大小，用于约束 rank 范围；不按此数切分 tail 的 head_dim。 + # 外部传输布局只保留一份 pool_size 长的逻辑窗口，所有 rank 可恢复到自己的 ring。 + tp_world_size: int = 1 + + def __post_init__(self) -> None: + """构造配置时检查维度、层号和 TP 范围。""" + assert min(self.pool_size, self.head_dim) > 0, "invalid K-pool dimensions" + assert self.layer_ids and len(set(self.layer_ids)) == len(self.layer_ids), ( + "tail layers must be nonempty and unique" + ) + assert 0 <= self.tp_rank < self.tp_world_size, "invalid TP rank" + + +class KPoolAttentionTailState(AttentionTailState): + """保存 K-pool 原始 key/gate ring，完整保留 pool_size + mtp_step 个槽位。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: KPoolStateConfig, + request_capacity: int, + mtp_step: int = 0, + *, + device: str | torch.device = "cuda", + ) -> None: + """为各层、各请求创建原始 key/gate ring，包含 MTP 所需额外历史。""" + assert request_capacity > 0 and mtp_step >= 0, "invalid request capacity or MTP step" + self.config = config + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step # ring 额外保留的 MTP 历史槽位数。 + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} + self.tail_buffer = torch.zeros( + ( + len(config.layer_ids), + request_capacity, + config.pool_size + mtp_step, + 2 * config.head_dim, + ), + dtype=config.dtype, + device=device, + ) + + def get_state_type(self) -> str: + """返回包含 K-pool 压缩比例的 tail 方案标记。""" + return f"indexer_kpool_c{self.config.pool_size}_tail" + + # ======================================================================== + # 二、Attention tail：计算视图与无快照重置 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> torch.Tensor: + """返回该层所有请求的原生 ring 视图，末轴包含 key 和 gate。""" + return self.tail_buffer[self._layer_mapping[layer_id]] + + def reset_attention_tail_state( + self, + req_indices: list[int], + ) -> None: + """清零指定请求的完整 ring；调用方保证用于空前缀或对齐池化边界，不用于 MTP 回滚。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.tail_buffer[:, req_idx].fill_(0) + + # ======================================================================== + # 三、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """描述最近 pool_size 个逻辑位置的 key/gate，由 TP rank 0 写出；不包含 MTP 额外槽位。""" + config = self.config + window_size = config.pool_size + channels = 2 * config.head_dim + tail_region = TransferRegion( + field="tail", + byte_start=0, + byte_end=len(config.layer_ids) * window_size * channels * config.dtype.itemsize, + dtype=config.dtype, + axes=("layer", "window", "channel"), + shape=(len(config.layer_ids), window_size, channels), + axis_slices=( + AxisSlice(axis="window", src_start=0, dst_start=0, length=window_size), + AxisSlice(axis="channel", src_start=0, dst_start=0, length=channels), + ), + write_enabled=config.tp_rank == 0, + ) + return (tail_region,) + + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, + *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: + """将 prefill 末尾的逻辑窗口写入 PD 页；state_token_len 为已执行的前缀长度。""" + copies = self._pd_buffer_copies(pd_page_buffers, pd_page_index, req_idx, state_token_len, layout, write=True) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, + *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + ) -> None: + """按前缀长度将逻辑窗口放回本地 ring，清零其余候选历史，允许两端 MTP 步数不同。""" + # TODO：PD 恢复可按目标 ring 槽位一次写入有效历史或零，融合清零和逐窗口复制； + # 避免同一 kernel 内先清零再恢复造成竞争，同时清除窗口外的旧候选历史。 + copies = self._pd_buffer_copies(pd_page_buffers, pd_page_index, req_idx, state_token_len, layout, write=False) + self.tail_buffer[:, req_idx].fill_(0) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + # ======================================================================== + # 四、私有辅助实现 + # ======================================================================== + + def _pd_buffer_copies( + self, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + req_idx: int, + state_token_len: int, + layout: StateTransferLayout, + *, + write: bool, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """为 PD prefill 交接生成复制视图，允许 P/D 使用不同的 MTP 步数。 + + 为什么不能直接复制物理 ring： + - 本地 ring 的容量为 P+S，其中 P=pool_size，S=本地 mtp_step；绝对位置 pos + 存在 pos % (P+S) 槽位。P/D 的 S 不同时，形状和同一位置对应的物理槽位都不同。 + - 即使 P 端没有执行 MTP decode，prefill 算子也按实际 ring 容量取模写入， + 最近 P 个位置不一定落在 ring 的前 P 个槽位，因此不能直接截取 [:P] 传输。 + - 外部页固定保存最近 P 个逻辑位置，按时间顺序排列，形状为 [layers, P, 2*head_dim]。 + 这样 PD 布局不依赖两端的 S；D 端的额外候选历史由后续 MTP 执行产生。 + + 位置和复制方向： + - state_token_len=L 是该 State 已执行完 prefill 的 token 数，不是最后一个 token + 的下标，也不包含刚采样但尚未执行的输出 token；来自 PD 任务的 end_kv_index。 + - 外部 window[i] 对应绝对位置 pos=L-P+i。write=True 从本地 ring 的 + pos % (P+S) 读取并写入 window[i]；write=False 按相同规则恢复到本地 ring。 + - 例如 P=4、L=11，窗口保存位置 [7,8,9,10]。S_p=0 时从槽位 [3,0,1,2] + 写出，S_d=2 时恢复到槽位 [1,2,3,4]；直接复制同名物理槽位会恢复错误历史。 + - L

= 0, "PD tail requires a nonnegative prefix length" + # helper 按 L 和本地 ring 实际容量逐位置取视图，不创建索引 Tensor 或读取 GPU 游标。 + return att_tail_buffer_copies( + self.tail_buffer[:, req_idx], pd_page_buffers, layout, + page_index=pd_page_index, state_token_len=state_token_len, write=write, + ) diff --git a/lightllm/common/att_resources/att_tail/transfer_utils.py b/lightllm/common/att_resources/att_tail/transfer_utils.py new file mode 100644 index 0000000000..85eb47fc61 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/transfer_utils.py @@ -0,0 +1,50 @@ +"""Attention tail 的 PD 页传输：在本地 ring 和固定逻辑窗口之间搬运。""" + +import torch + +from ..types import StateTransferLayout + + +def att_tail_buffer_copies( + tail: torch.Tensor, + page_buffers: torch.Tensor, + layout: StateTransferLayout, + *, + page_index: int, + state_token_len: int, + write: bool, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """按绝对位置映射 ring 与逻辑窗口，短前缀写出时将不存在的位置补零。 + + tail 轴为 (layer, ring_slot, channel)，外部页轴为 (layer, window, channel)。 + window 按时间顺序对应 [state_token_len - pool_size, state_token_len)， + 本地物理槽位使用绝对位置对 tail.shape[1] 取模，因此不要求 P/D 的 MTP 步数相同。 + 返回有效位置的 (源视图, 目标视图)，由 State 提交复制；读回前由 State 清零本地 ring。 + """ + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) + copies = [] + for region in layout.regions: + assert region.field == "tail", f"unknown transfer field {region.field}" + external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + external = external.permute(tuple(region.axes.index(axis) for axis in ("layer", "window", "channel"))) + assert tail.dtype == external.dtype, "local and external field dtypes must match" + assert (tail.shape[0], tail.shape[2]) == (external.shape[0], external.shape[2]), ( + "local and external tail layer/channel shapes must match" + ) + if write and not region.write_enabled: + continue + # TODO：用融合传输 kernel 完成逻辑窗口与 ring 的取模映射、写出补零和数据搬运， + # 减少逐窗口槽位 copy_/fill_；必须保留 P/D 不同 MTP 步数及 TP 副本只由 rank 0 写出的语义。 + # 请求级低频搬运；逐槽位取视图，不创建索引 Tensor，也不读取 GPU 上的游标。 + window_start = state_token_len - external.shape[1] + for window_index in range(external.shape[1]): + position = window_start + window_index + external_view = external[:, window_index] + if position < 0: + if write: + external_view.fill_(0) + continue + local_view = tail[:, position % tail.shape[1]] + copies.append((local_view, external_view) if write else (external_view, local_view)) + return copies diff --git a/lightllm/common/att_resources/bounded_att/__init__.py b/lightllm/common/att_resources/bounded_att/__init__.py new file mode 100644 index 0000000000..19f96ce1a3 --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/__init__.py @@ -0,0 +1,19 @@ +from .create_gdn_state import create_gdn_state +from .create_kda_state import create_kda_state +from .conv_ssm_types import ConvSSMStateConfig +from .gdn_pages_pool import GDNPagesPool +from .gdn_state import GDNState +from .kda_pages_pool import KDAPagesPool +from .kda_state import KDAState +from .base_state import BoundedAttentionState + +__all__ = [ + "create_kda_state", + "create_gdn_state", + "ConvSSMStateConfig", + "GDNState", + "KDAState", + "BoundedAttentionState", + "GDNPagesPool", + "KDAPagesPool", +] diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/att_resources/bounded_att/base_state.py new file mode 100644 index 0000000000..461f53b9ab --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/base_state.py @@ -0,0 +1,167 @@ +"""有界 attention 状态接口；具体算法实现运行态与大小页存储。""" + +from __future__ import annotations + +from typing import Any +from abc import ABC, abstractmethod +from typing import List + +import torch + +from lightllm.common.att_resources.types import ( + StateTransferLayout, + PageKind, + TransferRegion, +) + + +class BoundedAttentionState(ABC): + """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + @abstractmethod + def get_state_type(self) -> str: + """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + ... + + @abstractmethod + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建大小页池；容量为零表示不启用该页池。""" + ... + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + @abstractmethod + def get_layer_view(self, layer_id: int) -> Any: + """按逻辑层号返回运行态和算法控制信息的计算视图。""" + ... + + @abstractmethod + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" + ... + + @abstractmethod + def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """将 prefill 完成后的基准运行态保存到指定页槽位，前缀关联由调用方维护。""" + ... + + @abstractmethod + def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """从调用方选定的页槽位恢复运行态和控制信息，保留源快照数据。""" + ... + + @abstractmethod + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + ... + + # ======================================================================== + # 三、CPU Cache 数据交换：读写大小页快照 + # ======================================================================== + + @abstractmethod + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回一份 CPU 快照的字段区间、全局形状与 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 参与 CPU 页时，管理器按字段最大终点计算占用，不乘页内 token 数。 + """ + ... + + @abstractmethod + def write_page_state_to_cpu( + self, + page_kind: PageKind, + slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_readies 中为 True 的页面。""" + ... + + @abstractmethod + def read_page_state_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" + ... + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + @abstractmethod + def copy_big_page_state_from_other_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: BoundedAttentionState, + ) -> None: + """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" + ... + + @abstractmethod + def copy_run_time_state_from_other_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: BoundedAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + # ======================================================================== + # 五、PD 分离传输：读写基准运行态 + # ======================================================================== + + @abstractmethod + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" + ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将指定槽位的 prefill 基准运行态写入 PD 页。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + ... diff --git a/lightllm/common/att_resources/bounded_att/conv_ssm_types.py b/lightllm/common/att_resources/bounded_att/conv_ssm_types.py new file mode 100644 index 0000000000..4858ed11eb --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/conv_ssm_types.py @@ -0,0 +1,53 @@ +"""GDN/KDA 使用的 Conv/SSM 布局配置。""" + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class ConvSSMStateConfig: + """KDA/GDN 的层布局、维度、dtype 和 TP 配置；不包含模型角色或资源身份。""" + + # 本 State 管理的逻辑层编号；顺序对应 Conv/SSM Tensor 的物理层轴，编号不能重复。 + layer_ids: tuple[int, ...] + # TP 切分前的全局 K head 数，不计 TP 副本；Q 使用相同的 head 数和维度。 + global_k_heads: int + # TP 切分前的全局 V head 数；GDN 可与 K 不同，KDA 要求两者相等。 + global_v_heads: int + # 每个 Q/K head 的特征维度，也是 SSM 矩阵的 key 轴长度，不随 TP 切分。 + key_dim: int + # 每个 V head 的特征维度，也是 SSM 矩阵的 value 轴长度；KDA 要求与 key_dim 相等。 + value_dim: int + # Conv 需要保留的历史 token 数，等于卷积核宽度减 1；不含 MTP 额外候选位置。 + conv_width: int + # Conv 运行态及大小页快照的存储 dtype，不决定计算算子的中间精度。 + conv_dtype: torch.dtype + # SSM 运行态及大小页快照的存储 dtype，可与 Conv 不同，不决定算子的累加精度。 + ssm_dtype: torch.dtype + # 当前进程在本模型 TP 组中的编号，范围为 [0, tp_world_size)，用于确定本地 head 分片。 + tp_rank: int = 0 + # 本模型的 TP 并行度；head 足够时均分，少于 TP 并行度时按组复制；1 表示不切分。 + tp_world_size: int = 1 + + @property + def num_k_heads(self) -> int: + """当前 TP rank 的 K head 数；head 少于 TP 时按组复制。""" + return max(1, self.global_k_heads // self.tp_world_size) + + @property + def num_v_heads(self) -> int: + """当前 TP rank 的 V head 数，独立于 K head 的切分计算。""" + return max(1, self.global_v_heads // self.tp_world_size) + + @property + def conv_channels(self) -> int: + """当前 TP rank 的 Conv 通道总数，按 Q、K、V 三段拼接。""" + return 2 * self.num_k_heads * self.key_dim + self.num_v_heads * self.value_dim + + def __post_init__(self) -> None: + assert 0 <= self.tp_rank < self.tp_world_size + for heads in (self.global_k_heads, self.global_v_heads): + assert heads > 0 and ( + heads % self.tp_world_size == 0 or self.tp_world_size % heads == 0 + ), "Conv/SSM heads require complete TP shards or replicas" diff --git a/lightllm/common/att_resources/bounded_att/create_gdn_state.py b/lightllm/common/att_resources/bounded_att/create_gdn_state.py new file mode 100644 index 0000000000..83ed3db150 --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/create_gdn_state.py @@ -0,0 +1,19 @@ +"""GDN State 的统一创建入口。""" + +from __future__ import annotations + +import torch + +from .conv_ssm_types import ConvSSMStateConfig +from .gdn_state import GDNState + + +def create_gdn_state( + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + *, + device: str | torch.device = "cuda", +) -> GDNState: + """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" + return GDNState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) diff --git a/lightllm/common/att_resources/bounded_att/create_kda_state.py b/lightllm/common/att_resources/bounded_att/create_kda_state.py new file mode 100644 index 0000000000..989ecce54c --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/create_kda_state.py @@ -0,0 +1,19 @@ +"""KDA State 的统一创建入口。""" + +from __future__ import annotations + +import torch + +from .conv_ssm_types import ConvSSMStateConfig +from .kda_state import KDAState + + +def create_kda_state( + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + *, + device: str | torch.device = "cuda", +) -> KDAState: + """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" + return KDAState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) diff --git a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py new file mode 100644 index 0000000000..268859db57 --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py @@ -0,0 +1,45 @@ +"""GDN pages with separate Q/K and V dimensions; allocators own the page slots.""" + +from __future__ import annotations + + +import torch + +from .conv_ssm_types import ConvSSMStateConfig +from .utils import create_page_pool_tensor, register_page_pool_tensor + + +class GDNPagesPool: + def __init__( + self, + size: int, + config: ConvSSMStateConfig, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + assert size >= 0, "pool size must be nonnegative" + self.size = size + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + self.conv_state = create_page_pool_tensor( + size=(size, len(config.layer_ids), config.conv_channels, config.conv_width), + dtype=config.conv_dtype, + device=device, + pin_memory=pin_memory, + ) + self.ssm_state = create_page_pool_tensor( + size=(size, len(config.layer_ids), config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + pin_memory=pin_memory, + ) + + # ======================================================================== + # 跨进程映射：在接收进程的当前 GPU 上注册共享页池 + # ======================================================================== + + def register_pin_memory(self) -> None: + """注册 Conv/SSM 共享存储，供本进程 GPU 使用；重复调用不会重复注册。""" + register_page_pool_tensor(tensor=self.conv_state) + register_page_pool_tensor(tensor=self.ssm_state) diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py new file mode 100644 index 0000000000..5c684b794d --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .conv_ssm_types import ConvSSMStateConfig +from .gdn_pages_pool import GDNPagesPool +from .transfer_utils import bounded_att_buffer_copies + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import PageKind + + +logger = init_logger(__name__) + + +class GDNState(BoundedAttentionState): + """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" + + big_pages_pool: GDNPagesPool | None + small_pages_pool: GDNPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, *, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} + num_layers = len(config.layer_ids) + self.conv_state = torch.zeros( + (num_layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + (num_layers, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + def get_state_type(self) -> str: + """返回 GDN 算法标记。""" + return "gdn" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + GDNPagesPool(size=big_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + GDNPagesPool(size=small_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """返回该层 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + layer = self._layer_mapping[layer_id] + return self.conv_state[layer], self.ssm_state[layer] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + start = req_idx * (self.mtp_step + 1) + self.ssm_state[:, start : start + self.mtp_step + 1].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(slot_ids), "request and page counts differ" + pool = self._pool(page_kind) + # TODO：将多请求 Conv/SSM 保存融合为按 req/slot 索引搬运的 kernel； + # GPU 直访 CPU 页池需保证页池为可访问的 pinned/registered 内存，完成通知仍由调用方负责。 + # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 + # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + for req_idx, slot_id in zip(req_indices, slot_ids): + pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) + pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) + + def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(slot_ids), "request and page counts differ" + # TODO：融合多请求 Conv/SSM 恢复及目标游标清零，直接按 slot/req 索引写入基准位置； + # CPU 页池须满足 GPU 可访问条件，并保持未参与请求和其他候选区不变。 + pool = self._pool(page_kind) + for req_idx, slot_id in zip(req_indices, slot_ids): + self.conv_state[:, req_idx, :, : self.config.conv_width].copy_( + pool.conv_state[slot_id], non_blocking=True, + ) + self.ssm_state[:, req_idx * (self.mtp_step + 1)].copy_(pool.ssm_state[slot_id], non_blocking=True) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + if self.mtp_state_index is not None: + self.mtp_state_index[req_indices] = accepted_steps - 1 + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + 起止位置相对本 State，不包含 MTP 候选区、游标或大小页池的其他槽位。 + """ + return self._get_transfer_regions() + + def write_page_state_to_cpu( + self, + page_kind: PageKind, + slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_readies is not None: + assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" + copies = [] + # TODO：批量融合 Conv/SSM 多页写出及 TP 切片，减少逐槽位、逐字段 copy_； + # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): + # ready 页的完整联合数据保持不变，也无需构造快照复制视图。 + if cpu_page_readies is not None and cpu_page_readies[index]: + continue + fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) + copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=True)) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + def read_page_state_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + copies = [] + # TODO：批量融合 Conv/SSM 多页读入及 TP 切片，减少逐槽位、逐字段 copy_； + # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + for page, slot in zip(cpu_page_indices, big_page_slot_ids): + fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) + copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=False)) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_page_state_from_other_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: GDNState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( + "DP copy requires the same algorithm and equivalent TP state configurations" + ) + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # TODO：改用 Triton kernel 按页索引批量复制 Conv/SSM，在 CUDA stream 上异步执行。 + # 两端为 pinned CPU Tensor 时，当前 copy_ 仍执行 CPU 复制，non_blocking=True 不生效； + # kernel 需直接访问本进程已注册且 GPU 可寻址的两端页池，调用方须建立源页就绪依赖， + # 并在复制完成前保留源、目标页的引用；当前 CPU 路径须先确认源页 GPU 写入已完成。 + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + dst_pool.conv_state[dst_page].copy_(src_pool.conv_state[src_page], non_blocking=True) + dst_pool.ssm_state[dst_page].copy_(src_pool.ssm_state[src_page], non_blocking=True) + + def copy_run_time_state_from_other_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: GDNState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( + "DP copy requires the same algorithm and equivalent TP state configurations" + ) + # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； + # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 + # TODO：融合多请求 Conv/SSM 的 DP 复制与目标游标清零；跨设备源存储需可通过 peer/IPC 访问。 + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.conv_state[:, dst_req, :, : self.config.conv_width].copy_( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], non_blocking=True, + ) + self.ssm_state[:, dst_req * (self.mtp_step + 1)].copy_( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + non_blocking=True, + ) + if self.mtp_state_index is not None: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 + copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=True) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 + copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=False) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + layer_count = len(config.layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.resource_id, config.tp_rank, config.tp_world_size, conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_slices=( + AxisSlice( + axis="channel", + src_start=local_channel_start, + dst_start=target_channel_start, + length=local_channel_count, + ), + AxisSlice(axis="window", src_start=0, dst_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", + self.resource_id, name, local_channel_start, target_channel_start, + local_channel_count, replica_count, write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_slices=( + AxisSlice(axis="head", src_start=0, dst_start=head_start, length=config.num_v_heads), + AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=config.key_dim), + AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.resource_id, head_start, config.num_v_heads, write_enabled, ssm_region.byte_start, ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.resource_id}" + return pool + + def _runtime_fields(self, req: int): + """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" + # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 + # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 + return self._fields( + self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] + ) + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py new file mode 100644 index 0000000000..b06bfb1cf6 --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py @@ -0,0 +1,49 @@ +"""KDA pages: Conv history and square recurrent matrices; allocators own the page slots.""" + +from __future__ import annotations + + +import torch + +from .conv_ssm_types import ConvSSMStateConfig +from .utils import create_page_pool_tensor, register_page_pool_tensor + + +class KDAPagesPool: + def __init__( + self, + size: int, + config: ConvSSMStateConfig, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + assert size >= 0, "pool size must be nonnegative" + assert config.num_k_heads == config.num_v_heads and config.key_dim == config.value_dim, ( + "KDA pages require equal Q/K/V heads and dimensions" + ) + self.size = size + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + layers, heads, dim = len(config.layer_ids), config.num_k_heads, config.key_dim + self.conv_state = create_page_pool_tensor( + size=(size, layers, 3 * heads * dim, config.conv_width), + dtype=config.conv_dtype, + device=device, + pin_memory=pin_memory, + ) + self.ssm_state = create_page_pool_tensor( + size=(size, layers, heads, dim, dim), + dtype=config.ssm_dtype, + device=device, + pin_memory=pin_memory, + ) + + # ======================================================================== + # 跨进程映射：在接收进程的当前 GPU 上注册共享页池 + # ======================================================================== + + def register_pin_memory(self) -> None: + """注册 Conv/SSM 共享存储，供本进程 GPU 使用；重复调用不会重复注册。""" + register_page_pool_tensor(tensor=self.conv_state) + register_page_pool_tensor(tensor=self.ssm_state) diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py new file mode 100644 index 0000000000..6e8a5dbbe7 --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .conv_ssm_types import ConvSSMStateConfig +from .kda_pages_pool import KDAPagesPool +from .transfer_utils import bounded_att_buffer_copies + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import PageKind + + +logger = init_logger(__name__) + + +class KDAState(BoundedAttentionState): + """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" + + big_pages_pool: KDAPagesPool | None + small_pages_pool: KDAPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, *, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} + num_layers = len(config.layer_ids) + self.conv_state = torch.zeros( + (num_layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + (num_layers, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + def get_state_type(self) -> str: + """返回 KDA 算法标记。""" + return "kda" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + KDAPagesPool(size=big_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + KDAPagesPool(size=small_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """返回该层 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + layer = self._layer_mapping[layer_id] + return self.conv_state[layer], self.ssm_state[layer] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + start = req_idx * (self.mtp_step + 1) + self.ssm_state[:, start : start + self.mtp_step + 1].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(slot_ids), "request and page counts differ" + pool = self._pool(page_kind) + # TODO：将多请求 Conv/SSM 保存融合为按 req/slot 索引搬运的 kernel； + # GPU 直访 CPU 页池需保证页池为可访问的 pinned/registered 内存，完成通知仍由调用方负责。 + # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 + # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + for req_idx, slot_id in zip(req_indices, slot_ids): + pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) + pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) + + def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" + assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(slot_ids), "request and page counts differ" + # TODO：融合多请求 Conv/SSM 恢复及目标游标清零，直接按 slot/req 索引写入基准位置； + # CPU 页池须满足 GPU 可访问条件，并保持未参与请求和其他候选区不变。 + pool = self._pool(page_kind) + for req_idx, slot_id in zip(req_indices, slot_ids): + self.conv_state[:, req_idx, :, : self.config.conv_width].copy_( + pool.conv_state[slot_id], non_blocking=True, + ) + self.ssm_state[:, req_idx * (self.mtp_step + 1)].copy_(pool.ssm_state[slot_id], non_blocking=True) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + if self.mtp_state_index is not None: + self.mtp_state_index[req_indices] = accepted_steps - 1 + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + 起止位置相对本 State，不包含 MTP 候选区、游标或大小页池的其他槽位。 + """ + return self._get_transfer_regions() + + def write_page_state_to_cpu( + self, + page_kind: PageKind, + slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_readies is not None: + assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" + copies = [] + # TODO：批量融合 Conv/SSM 多页写出及 TP 切片，减少逐槽位、逐字段 copy_； + # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): + # ready 页的完整联合数据保持不变，也无需构造快照复制视图。 + if cpu_page_readies is not None and cpu_page_readies[index]: + continue + fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) + copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=True)) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + def read_page_state_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + copies = [] + # TODO：批量融合 Conv/SSM 多页读入及 TP 切片，减少逐槽位、逐字段 copy_； + # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + for page, slot in zip(cpu_page_indices, big_page_slot_ids): + fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) + copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=False)) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_page_state_from_other_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: KDAState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( + "DP copy requires the same algorithm and equivalent TP state configurations" + ) + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # TODO：改用 Triton kernel 按页索引批量复制 Conv/SSM，在 CUDA stream 上异步执行。 + # 两端为 pinned CPU Tensor 时，当前 copy_ 仍执行 CPU 复制，non_blocking=True 不生效； + # kernel 需直接访问本进程已注册且 GPU 可寻址的两端页池，调用方须建立源页就绪依赖， + # 并在复制完成前保留源、目标页的引用；当前 CPU 路径须先确认源页 GPU 写入已完成。 + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + dst_pool.conv_state[dst_page].copy_(src_pool.conv_state[src_page], non_blocking=True) + dst_pool.ssm_state[dst_page].copy_(src_pool.ssm_state[src_page], non_blocking=True) + + def copy_run_time_state_from_other_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: KDAState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( + "DP copy requires the same algorithm and equivalent TP state configurations" + ) + # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； + # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 + # TODO：融合多请求 Conv/SSM 的 DP 复制与目标游标清零；跨设备源存储需可通过 peer/IPC 访问。 + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.conv_state[:, dst_req, :, : self.config.conv_width].copy_( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], non_blocking=True, + ) + self.ssm_state[:, dst_req * (self.mtp_step + 1)].copy_( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + non_blocking=True, + ) + if self.mtp_state_index is not None: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_buffer( + self, req_idx: int, pd_page_buffers: torch.Tensor, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 + copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=True) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + + def read_from_pd_buffer( + self, pd_page_buffers: torch.Tensor, req_idx: int, + *, pd_page_index: int, layout: StateTransferLayout, + ) -> None: + """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 + copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=False) + for src_tensor, dst_tensor in copies: + dst_tensor.copy_(src_tensor, non_blocking=True) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + layer_count = len(config.layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.resource_id, config.tp_rank, config.tp_world_size, conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_slices=( + AxisSlice( + axis="channel", + src_start=local_channel_start, + dst_start=target_channel_start, + length=local_channel_count, + ), + AxisSlice(axis="window", src_start=0, dst_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", + self.resource_id, name, local_channel_start, target_channel_start, + local_channel_count, replica_count, write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_slices=( + AxisSlice(axis="head", src_start=0, dst_start=head_start, length=config.num_v_heads), + AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=config.key_dim), + AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.resource_id, head_start, config.num_v_heads, write_enabled, ssm_region.byte_start, ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.resource_id}" + return pool + + def _runtime_fields(self, req: int): + """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" + # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 + # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 + return self._fields( + self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] + ) + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/att_resources/bounded_att/transfer_utils.py b/lightllm/common/att_resources/bounded_att/transfer_utils.py new file mode 100644 index 0000000000..61f07a4b1a --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/transfer_utils.py @@ -0,0 +1,73 @@ +"""Bounded attention 的 CPU/PD 页传输：复制一份完整状态快照。""" + +from collections.abc import Mapping + +import torch + +from ..types import StateTransferLayout, TensorField, TransferRegion + + +# ======================================================================== +# 一、CPU/PD 页交换：按状态字段生成复制视图 +# ======================================================================== + +def bounded_att_buffer_copies( + fields: Mapping[str, TensorField], + page_buffers: torch.Tensor, + layout: StateTransferLayout, + *, + page_index: int, + write: bool, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """生成一份 bounded 状态与外部页之间的复制计划，不处理 token 数量。 + + fields 由 State 选取运行态或大小页快照，例如 KDA/GDN 的 Conv、SSM。 + 每份快照完整复制；Conv 的 Q/K/V 区间可以分别映射到同一全局字段。 + page_buffers 至少为二维，第一维统一表示页；不接受没有页维度的单页 Tensor。 + page_index 必须显式传入；CPU/PD 调用方均传入页集合和实际物理页索引。 + 返回 (源视图, 目标视图) 列表，由调用方遍历执行 dst.copy_(src, non_blocking=True)。 + """ + # CPU/PD 统一沿第一维选页，再按字节解释页内存储，不复制数据。 + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + page_buffer = page_buffers[page_index] + page_bytes = page_buffer.view(torch.uint8).view(-1) + copies = [] + for region in layout.regions: + # 阶段一：定位快照字段和外部页中的全局状态区域。 + assert region.field in fields, f"unknown transfer field {region.field}" + field = fields[region.field] + # region 起止位置相对整页，按字段 dtype 和 shape 直接构造外部视图。 + external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + if write and not region.write_enabled: + continue + + # 阶段二：按 channel/head 等 TP 切片生成视图，Conv 窗口及 SSM 状态完整保留。 + local_view, external_view = _copy_pair(field, external, region) + # 比较切片后的实际复制范围，避免 copy_ 隐式广播掩盖布局错误。 + assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" + # 写出只由选定的 TP 副本执行；读回覆盖每个副本自己的分片。 + copies.append((local_view, external_view) if write else (external_view, local_view)) + return copies + + +# ======================================================================== +# 二、私有辅助：本类状态的复制视图 +# ======================================================================== + +def _copy_pair( + field: TensorField, + external: torch.Tensor, + region: TransferRegion, +) -> tuple[torch.Tensor, torch.Tensor]: + """按 bounded 状态的 channel/head 等映射取视图，不处理 token 裁剪。""" + assert field.tensor.dtype == external.dtype, "local and external field dtypes must match" + # 外部布局可采用不同轴顺序，先转换为本地字段的顺序；只创建视图。 + external = external.permute(tuple(region.axes.index(axis) for axis in field.axes)) + local_slices = [slice(None)] * field.tensor.ndim + external_slices = [slice(None)] * external.ndim + # 布局已明确每个分片的两侧起点和长度，其余轴完整保留。 + for axis_slice in region.axis_slices: + axis = field.axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + return field.tensor[tuple(local_slices)], external[tuple(external_slices)] diff --git a/lightllm/common/att_resources/bounded_att/utils.py b/lightllm/common/att_resources/bounded_att/utils.py new file mode 100644 index 0000000000..ca0fc3fb7b --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/utils.py @@ -0,0 +1,114 @@ +"""Bounded attention 页池的存储分配策略。""" + +from __future__ import annotations + +from math import prod +from weakref import finalize + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + + +logger = init_logger(__name__) + +# Portable：本进程内其他 CUDA context 也将该内存识别为 pinned；不代替跨进程注册。 +CUDA_HOST_REGISTER_PORTABLE = 1 +# Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。 +CUDA_HOST_REGISTER_MAPPED = 2 + + +# ======================================================================== +# 一、页池分配：普通存储或 DP fetch 所需的共享 pinned 存储 +# ======================================================================== + +def create_page_pool_tensor( + size: tuple[int, ...], + dtype: torch.dtype, + *, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, +) -> torch.Tensor: + """按 DP fetch 开关创建页池 Tensor，保留调用方指定的设备及 pin_memory 行为。 + + CPU 页池在启用 DP fetch 时直接使用共享 Storage，避免普通 pinned Tensor + 经 Queue 导出时重新分配共享存储、复制整块页池，并丢失原地址的 CUDA 注册。 + 接收进程仍需为自己的共享映射地址注册 CUDA；注册状态不随 Tensor 传递。 + """ + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + shared = device.type == "cpu" and get_env_start_args().enable_dp_prompt_cache_fetch + nbytes = prod(size) * dtype.itemsize + + # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 + if not shared: + tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=pin_memory) + reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" + logger.info( + f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " + f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}" + ) + return tensor + + # 阶段二：直接申请共享字节存储，再绑定无数据区的 Tensor，不经过 share_memory_ 的整块复制。 + # _new_shared 是 PyTorch 私有入口，集中在此使用，便于版本升级时验证或替换。 + storage = torch.UntypedStorage._new_shared(nbytes, device="cpu") + tensor = torch.empty(size=(0,), dtype=dtype, device="cpu").set_( + source=storage, + storage_offset=0, + size=size, + ) + + # 阶段三：在当前进程原地注册；Portable | Mapped 允许 CUDA 异步搬运及设备映射。 + # 不调用 tensor.pin_memory()，否则会另建私有 pinned 存储，破坏共享关系。 + # 空页池没有数据，无需注册；显式 pin_memory=False 时仅创建共享 CPU 存储。 + if pin_memory: + register_page_pool_tensor(tensor=tensor) + logger.info( + f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " + f"device=cpu, shared=True, pin_memory={pin_memory}, reason=DP prompt cache fetch; " + "direct shared allocation avoids a full pool copy during IPC export" + ) + return tensor + + +# ======================================================================== +# 二、本进程注册：首次创建和跨进程接收复用同一入口 +# ======================================================================== + +def register_page_pool_tensor(tensor: torch.Tensor) -> None: + """将共享 CPU 存储注册到当前进程，已 pinned、空存储或 GPU Tensor 无需处理。 + + 接收端必须先选定自己的 CUDA device。按整个 Storage 注册，避免多个视图重复注册 + 重叠区间；注册和注销信息只保存在本进程，不依赖发送方的对象属性或默认 GPU。 + """ + if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned(): + return + assert tensor.is_shared(), "page pool registration requires shared CPU storage" + storage = tensor.untyped_storage() + ptr = storage.data_ptr() + nbytes = storage.nbytes() + device_id = torch.cuda.current_device() + flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED + torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) + # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 + # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 + cleanup = finalize(storage, _unregister_page_pool_storage, ptr, device_id) + cleanup.atexit = False + logger.info( + f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + "flags=Portable|Mapped; registration applies to this process's mapping" + ) + + +# ======================================================================== +# 三、私有辅助：最后一个本地存储引用释放时注销 CUDA 注册 +# ======================================================================== + + +def _unregister_page_pool_storage(ptr: int, device_id: int) -> None: + """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。""" + torch.cuda.synchronize(device=device_id) + torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) + logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}") diff --git a/lightllm/common/att_resources/full_att/__init__.py b/lightllm/common/att_resources/full_att/__init__.py new file mode 100644 index 0000000000..b51fc60359 --- /dev/null +++ b/lightllm/common/att_resources/full_att/__init__.py @@ -0,0 +1,31 @@ +from .create_mla_state import create_mla_state +from .create_gqa_state import create_gqa_state +from .base_state import FullAttentionConfig, FullAttentionState +from .gqa_state import GQAFullAttentionState +from .mla_state import MLAFullAttentionState + +__all__ = [ + "create_gqa_state", + "create_mla_state", + "FullAttentionState", + "GQAFullAttentionState", + "FullAttentionConfig", + "MLAFullAttentionState", +] + +from .group_quantized_state import ( + GroupQuantizedFullAttentionState, + GroupQuantizedAttentionConfig, +) +from .static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig +from .packed_mla_state import PackedMLAFullAttentionState, PackedMLAAttentionConfig, PackedMLALayerView + +__all__ += [ + "GroupQuantizedFullAttentionState", + "GroupQuantizedAttentionConfig", + "StaticFP8FullAttentionState", + "StaticFP8AttentionConfig", + "PackedMLAFullAttentionState", + "PackedMLAAttentionConfig", + "PackedMLALayerView", +] diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py new file mode 100644 index 0000000000..b584e03c14 --- /dev/null +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +from typing import List +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + +import torch + +from lightllm.common.att_resources.types import ( + StateTransferLayout, + TransferRegion, +) + + +@dataclass(frozen=True) +class FullAttentionConfig: + # 当前 State 管理的逻辑层号；元组顺序对应 storage 的 layer 轴顺序，不要求层号连续。 + layer_ids: tuple[int, ...] + # TP 切分前的 KV head 数：GQA 的 K/V 各有这么多个，不把二者相加，也不计 TP 副本。 + # 不是 query head 数；例如全局 8 个 KV head，实际有 8 个 K head 和 8 个 V head。 + # MLA 的压缩 latent 作为一个 KV 字段存储，通常设为 1，并在各 TP rank 复制。 + global_head_num: int + # GQA 中每个 K/V head 的原始向量维度；INT4 的实际编码宽度由具体 State 换算。 + # 普通 MLA 中为 latent 宽度；packed MLA 中为整行元素数，包含 indexer、scale 和 padding。 + head_dim: int + # 具体方案使用的数据类型：普通 GQA/MLA 为 KV 精度，分组 INT4/INT8 为 scale 精度。 + # 静态 FP8 和 FP8 packed MLA 使用 uint8 承载编码字节，实际字段类型由 State 解释。 + dtype: torch.dtype + # 当前 State 所属模型的 TP 组内 rank，不是全局进程 rank 或 DP rank。 + tp_rank: int = 0 + # 当前模型 TP 组大小，用于计算本地 head 数以及传输时的全局 head 范围。 + tp_world_size: int = 1 + # True：每个 rank 都保存全部全局 heads，例如 MLA latent。 + # False：heads 足够时均匀切分；heads 少于 TP rank 数时，每组相邻 ranks 复制一个 head。 + # 例如 global_head_num=2、TP=8 时，rank 0..3 保存 head 0，rank 4..7 保存 head 1。 + replicated: bool = False + + @property + def head_num(self) -> int: + """当前 TP rank 实际存储的 KV head 数。""" + # 仍是单份 K 或 V 的 head 数；GQA 的 K/V 独立分配，各自的 head 轴大小均为 head_num。 + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + + def __post_init__(self) -> None: + assert self.global_head_num > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" + + +class FullAttentionState(ABC): + """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 + + 所有表示 token 索引的参数都须为调用方准备的 CUDA Tensor，位于对应存储设备上； + State 不转换设备、dtype 或连续布局，int32/int64 由具体算子支持，DP 目标索引须为 int64。 + CPU 页位于主机内存不改变 token 索引的 CUDA 契约。 + """ + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + # 具体存储和量化方案标识，用于 main/draft State 成员命名。 + @abstractmethod + def get_state_type(self) -> str: + """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + ... + + @abstractmethod + def __init__( + self, + layout_config: FullAttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + ) -> None: + """只准备 State 配置，不分配 KV 或 HOLD 页；存储由后续统一容量初始化。""" + raise NotImplementedError() + + @abstractmethod + def initialize_storage(self, token_capacity: int) -> None: + """按 token 容量创建独立存储，沿用已设置的页大小、设备及配套参数。""" + raise NotImplementedError() + + @abstractmethod + def get_token_bytes(self) -> int: + """返回当前 TP rank 的单 token 存储字节数，包含动态量化参数和行 padding。""" + raise NotImplementedError() + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + @abstractmethod + def get_layer_view(self, layer_id: int) -> Any: + """按逻辑层号返回计算所需的存储视图，具体字段由该布局定义。""" + raise NotImplementedError() + + @abstractmethod + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: + """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。""" + raise NotImplementedError() + + @abstractmethod + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制全部所辖层的完整 token 数据，包含动态 scale、indexer 字段，保持原存储表示。 + + 源和目标索引为存储设备上的等长一维 CUDA 整数 Tensor，位置不重叠，直接使用原 dtype 和 stride。""" + raise NotImplementedError() + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @abstractmethod + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 cpu_cache_token_page_size 返回 CPU 字段、全局形状和 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 管理器按字段最大终点计算该 State 的 CPU 页占用。 + 本次有效 token 数由读写流程处理，不改变布局。 + """ + ... + + @abstractmethod + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局写出 token 数据及配套字段；cpu_page_readies 中为 True 的页跳过写入。""" + raise NotImplementedError() + + @abstractmethod + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" + raise NotImplementedError() + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + @abstractmethod + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: FullAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，包含量化参数及 indexer 字段。 + + 源、目标索引须已位于各自存储的 CUDA 设备；目标索引为 int64，接口不转换设备或 dtype。 + """ + raise NotImplementedError() + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @abstractmethod + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" + ... + + @abstractmethod + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将 token 数据写入PD 页集合中的指定发送页，网络发送与同步由调用方负责。""" + raise NotImplementedError() + + @abstractmethod + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """从PD 页集合中的指定接收页还原 token 数据，不分配索引或修改请求映射。""" + raise NotImplementedError() diff --git a/lightllm/common/att_resources/full_att/create_gqa_state.py b/lightllm/common/att_resources/full_att/create_gqa_state.py new file mode 100644 index 0000000000..527e1bcf7e --- /dev/null +++ b/lightllm/common/att_resources/full_att/create_gqa_state.py @@ -0,0 +1,67 @@ +"""GQA（含 MHA/MQA）的统一创建入口，选择普通精度或量化 State。""" + +from __future__ import annotations + +from typing import Literal + +import torch + +from .base_state import FullAttentionConfig +from .gqa_state import GQAFullAttentionState +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState + + +def create_gqa_state( + config: FullAttentionConfig, + *, + kv_type: Literal["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"] = "None", + page_size: int = 1, + device: str | torch.device = "cuda", + group_size: int = 8, + scales: torch.Tensor | None = None, + q_scales: torch.Tensor | None = None, + calibration_id: str = "", +) -> GQAFullAttentionState | GroupQuantizedFullAttentionState | StaticFP8FullAttentionState: + """按编码参数创建 GQA（含 MHA/MQA）State，普通 token 容量由后续流程分配。 + + config 描述逻辑 KV 维度和 TP 信息；kv_type 沿用启动参数中的编码名称。 + group_size 仅用于 INT4/INT8；scales、q_scales 和 calibration_id 仅用于静态 FP8。 + 静态标定数据由调用方加载，本入口不读取启动参数、配置文件或绑定模型角色。 + """ + if kv_type == "None": + return GQAFullAttentionState(layout_config=config, page_size=page_size, device=device) + if kv_type in ("int4kv", "int8kv"): + quant_config = GroupQuantizedAttentionConfig( + layer_ids=config.layer_ids, + global_head_num=config.global_head_num, + head_dim=config.head_dim, + dtype=config.dtype, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + replicated=config.replicated, + bits=4 if kv_type == "int4kv" else 8, + group_size=group_size, + ) + return GroupQuantizedFullAttentionState(layout_config=quant_config, page_size=page_size, device=device) + assert kv_type in ("fp8kv_spt", "fp8kv_sph"), f"unsupported GQA encoding: {kv_type}" + assert scales is not None, "static FP8 requires calibration scales" + assert kv_type != "fp8kv_sph" or q_scales is not None, "per-head FP8 requires Q calibration scales" + quant_config = StaticFP8AttentionConfig( + layer_ids=config.layer_ids, + global_head_num=config.global_head_num, + head_dim=config.head_dim, + dtype=torch.uint8, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + replicated=config.replicated, + scale_mode="per_head" if kv_type == "fp8kv_sph" else "per_tensor", + calibration_id=calibration_id, + ) + return StaticFP8FullAttentionState( + layout_config=quant_config, + page_size=page_size, + device=device, + scales=scales, + q_scales=q_scales, + ) diff --git a/lightllm/common/att_resources/full_att/create_mla_state.py b/lightllm/common/att_resources/full_att/create_mla_state.py new file mode 100644 index 0000000000..e87332cbe1 --- /dev/null +++ b/lightllm/common/att_resources/full_att/create_mla_state.py @@ -0,0 +1,55 @@ +"""MLA 的统一创建入口，选择普通、packed 或 packed FP8 State。""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Literal + +import torch + +from .base_state import FullAttentionConfig +from .mla_state import MLAFullAttentionState +from .packed_mla_state import PackedMLAAttentionConfig, PackedMLAFullAttentionState + + +def create_mla_state( + config: FullAttentionConfig, + *, + kv_type: Literal["None", "fp8kv_dsa"] = "None", + packed: bool = False, + indexer_dim: int = 128, + page_size: int = 1, + device: str | torch.device = "cuda", +) -> MLAFullAttentionState | PackedMLAFullAttentionState: + """创建普通或 packed MLA State，统一处理量化编码、TP 副本和物理行宽。 + + config.head_dim 始终是逻辑 KV latent 维度，尚未包含 indexer、scale 和 padding。 + packed=True 时把 indexer 字段存入同一行；fp8kv_dsa 仅适用于 packed MLA。 + """ + if not packed: + assert kv_type == "None", f"unsupported MLA encoding: {kv_type}" + return MLAFullAttentionState( + layout_config=replace(config, replicated=True), + page_size=page_size, + device=device, + ) + assert kv_type in ("None", "fp8kv_dsa"), f"unsupported packed MLA encoding: {kv_type}" + fp8 = kv_type == "fp8kv_dsa" + storage_dtype = torch.uint8 if fp8 else config.dtype + # FP8 latent 包含 512 字节编码、16 字节 scale 和 128 字节 BF16 RoPE。 + # indexer key 使用 INT8 并附带一个 FP32 scale；整行按 16 字节对齐。 + kv_bytes = 656 if fp8 else config.head_dim * config.dtype.itemsize + row_bytes = (kv_bytes + indexer_dim + 4 + 15) // 16 * 16 + packed_config = PackedMLAAttentionConfig( + layer_ids=config.layer_ids, + global_head_num=config.global_head_num, + head_dim=row_bytes // storage_dtype.itemsize, + dtype=storage_dtype, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + replicated=True, + kv_dim=config.head_dim, + indexer_dim=indexer_dim, + fp8_nope_dim=512 if fp8 else 0, + ) + return PackedMLAFullAttentionState(layout_config=packed_config, page_size=page_size, device=device) diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py new file mode 100644 index 0000000000..4a077b6cc1 --- /dev/null +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -0,0 +1,339 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ..utils import dtype_name +from ..types import AxisSlice, TransferRegion +from .base_state import FullAttentionConfig, FullAttentionState +from .gqa_transfer import transfer_gqa_pages + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +class GQAFullAttentionState(FullAttentionState): + """GQA 的 full-attention KV 存储，MHA 和 MQA 均为其特例。 + + TP 切分前，H_q 个 query head 分为 H_kv 组，每组共享一对 K/V head， + 要求 H_q 能被 H_kv 整除。MHA 对应 H_kv == H_q；MQA 对应 H_kv == 1； + 其余 GQA 对应 1 < H_kv < H_q。 + + config.head_num 表示 TP 切分或复制后的本地 KV head 数。 + K/V 分别为 [layers, token_slots, local_kv_heads, head_dim]，两者独立分配。 + query 与 KV 的分组计算由 attention backend 处理，因此三种情况共用此 State。 + """ + + # 保留既有 K/V 编码的稳定传输布局标识。 + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: FullAttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + ) -> None: + """只记录布局配置，KV 与 HOLD 页由后续统一容量初始化分配。""" + self.config = layout_config + self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化方案显式设置，区别于 scale 精度。 + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 GQA KV 精度对应的存储方案标记。""" + return "full_att_" + dtype_name(dtype=self.config.dtype) + + def initialize_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + + def get_token_bytes(self) -> int: + """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" + config = self.config + return ( + len(config.layer_ids) * 2 * config.head_num * config.head_dim + * self.storage_dtype.itemsize + ) + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" + layer = self._layer_mapping[layer_id] + return self.k_storage[layer], self.v_storage[layer] + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + layer = self._layer_mapping[layer_id] + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 + destindex_copy_kv(K=kv[:, :self.config.head_num], DestLoc=token_indices, Out=self.k_storage[layer]) + destindex_copy_kv(K=kv[:, self.config.head_num:], DestLoc=token_indices, Out=self.v_storage[layer]) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( + "token copies require equal one-dimensional index tensors" + ) + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( + "token indices must be integer tensors" + ) + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局将 token 数据写到指定 CPU 页，覆盖本 State 负责的字段和 TP 分片。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_readies, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: GQAFullAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" + assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( + "DP copy requires the same encoding and equivalent TP state configurations" + ) + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； + # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 + # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 + src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) + dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) + self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) + src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) + dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) + self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按 token 索引将 K/V 直接写入指定 PD 页，网络发送与同步由调用方负责。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将指定 PD 页直接恢复到 K/V 物理槽位，不分配中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", self.storage_dtype, config.head_dim, "dim"), + ("v", self.storage_dtype, config.head_dim, "dim"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_slices = tuple( + AxisSlice( + axis=axis, + src_start=0, + dst_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_slices=axis_slices, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.resource_id, field_name, head_start, local_head_count, replica_count, + write_enabled, field_region.byte_start, field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.layer_ids), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/att_resources/full_att/gqa_transfer.py b/lightllm/common/att_resources/full_att/gqa_transfer.py new file mode 100644 index 0000000000..55cd40f05e --- /dev/null +++ b/lightllm/common/att_resources/full_att/gqa_transfer.py @@ -0,0 +1,238 @@ +"""GQA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ..types import StateTransferLayout + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_gqa_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + *, + layout: StateTransferLayout, + write: bool, + grid_num: int = 36, + page_readies: list[bool] | None = None, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + else: + assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + axes = ("layer", "token", "head", "dim") + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + elif region.field == "v": + local = v_storage + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, dim]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_slice in region.axis_slices: + # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], external.shape[3], external.shape[4] + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 K/V 字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_gqa_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=256, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_gqa_pages( + # 数据指针：一次处理 K 或 V；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, dim]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 每个 head 参与拷贝的元素数。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in range(program_start, element_count, grid_stride): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + layer = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 K/V 字段的大小。 + local_offset = ( + layer * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + layer * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py new file mode 100644 index 0000000000..ed857731df --- /dev/null +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -0,0 +1,408 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ..utils import dtype_name +from ..types import AxisSlice, TransferRegion +from .base_state import FullAttentionConfig, FullAttentionState +from .group_quantized_transfer import transfer_group_quantized_pages + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class GroupQuantizedAttentionConfig(FullAttentionConfig): + # 每个 KV 元素的量化位数，仅支持 4/8；INT4 将两个编码值打包到一个字节。 + bits: int = 8 + # 沿每个 head 的 dim 轴分组，每个 token、每组共用一个动态 scale；须整除 head_dim。 + # K/V 分别保存 scale，精度由继承的 dtype 指定。 + group_size: int = 8 + + def __post_init__(self) -> None: + """创建配置时校验通用 TP 约束及量化分组要求。""" + super().__post_init__() + assert self.bits in (4, 8) and self.group_size >= 2 and not ( + self.group_size & (self.group_size - 1) + ), "INT4/INT8 quantization requires a power-of-two group size" + assert self.head_dim % self.group_size == 0 and self.dtype in ( + torch.float16, torch.bfloat16, torch.float32 + ), "group scales require a floating dtype and complete groups" + + +class GroupQuantizedFullAttentionState(FullAttentionState): + """管理 INT4/INT8 KV 与逐 token 分组 scale，复制和传输始终覆盖二者。""" + + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: GroupQuantizedAttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + ) -> None: + """按分组量化配置准备 KV 与动态 scale 存储。""" + self.config = layout_config + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + self.storage_dtype = torch.int8 + self.k_scales: torch.Tensor | None = None + self.v_scales: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" + return f"full_att_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.dtype)}_scale" + + def initialize_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + shape = ( + len(self.config.layer_ids), token_capacity + self.page_size, + self.config.head_num, self.config.head_dim // self.config.group_size, + ) + self.k_scales = torch.empty(size=shape, dtype=self.config.dtype, device=self.device) + self.v_scales = torch.empty(size=shape, dtype=self.config.dtype, device=self.device) + + def get_token_bytes(self) -> int: + """累计量化 KV 编码与逐 token 分组 scale 的开销，无需分配存储。""" + config = self.config + head_bytes = config.head_dim * config.bits // 8 + head_bytes += config.head_dim // config.group_size * config.dtype.itemsize + return len(config.layer_ids) * 2 * config.head_num * head_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + """返回该层 K/V 编码数据及各自的分组 scale 视图。""" + layer = self._layer_mapping[layer_id] + return (self.k_storage[layer], self.k_scales[layer]), (self.v_storage[layer], self.v_scales[layer]) + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """量化输入浮点 K/V，并将编码数据和分组 scale 写入相同 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + layer = self._layer_mapping[layer_id] + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + # TODO：融合 K/V 两次量化写入，同时输出各自编码和动态 scale，保持分组及舍入语义。 + # K/V 分别写入独立存储，各自携带同一 token 的动态 scale。 + for offset, storage, scales_buffer in ( + (0, self.k_storage, self.k_scales), + (self.config.head_num, self.v_storage, self.v_scales), + ): + values = kv[:, offset:offset + self.config.head_num] + if self.config.bits == 4: + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv + + destindex_copy_int4kv( + KV=values, + DestLoc=token_indices, + KV_buffer=storage[layer], + KV_scale_buffer=scales_buffer[layer], + quant_group_size=self.config.group_size, + ) + else: + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import destindex_copy_quantize_kv + + destindex_copy_quantize_kv( + K=values, + DestLoc=token_indices, + Out=storage[layer], + Out_scale=scales_buffer[layer], + quant_group_dim=self.config.group_size, + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( + "token copies require equal one-dimensional index tensors" + ) + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( + "token indices must be integer tensors" + ) + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 K/V 编码及各自的动态 scale 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_scales, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_scales, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_readies, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: GroupQuantizedFullAttentionState, + ) -> None: + """校验量化位数和分组大小，再复制 KV 与动态 scale。""" + assert ( + remote_dp_full_att_state.config.bits, + remote_dp_full_att_state.config.group_size, + ) == ( + self.config.bits, self.config.group_size + ), "incompatible group quantization" + assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( + "DP copy requires the same encoding and equivalent TP state configurations" + ) + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； + # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 + # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 + src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) + dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) + self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) + src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) + dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) + self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) + src_k_scale = remote_dp_full_att_state.k_scales.index_select(dim=1, index=src_token_indices) + dst_k_scale = src_k_scale.to(device=self.k_scales.device, non_blocking=True) + self.k_scales.index_copy_(dim=1, index=dst_token_indices, source=dst_k_scale) + src_v_scale = remote_dp_full_att_state.v_scales.index_select(dim=1, index=src_token_indices) + dst_v_scale = src_v_scale.to(device=self.v_scales.device, non_blocking=True) + self.v_scales.index_copy_(dim=1, index=dst_token_indices, source=dst_v_scale) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", torch.int8, config.head_dim * config.bits // 8, "dim"), + ("v", torch.int8, config.head_dim * config.bits // 8, "dim"), + ("k_scale", config.dtype, config.head_dim // config.group_size, "group"), + ("v_scale", config.dtype, config.head_dim // config.group_size, "group"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_slices = tuple( + AxisSlice( + axis=axis, + src_start=0, + dst_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_slices=axis_slices, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.resource_id, field_name, head_start, local_head_count, replica_count, + write_enabled, field_region.byte_start, field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.layer_ids), self.capacity + self.page_size, + self.config.head_num, self.config.head_dim * self.config.bits // 8, + ) diff --git a/lightllm/common/att_resources/full_att/group_quantized_transfer.py b/lightllm/common/att_resources/full_att/group_quantized_transfer.py new file mode 100644 index 0000000000..9b47b269fd --- /dev/null +++ b/lightllm/common/att_resources/full_att/group_quantized_transfer.py @@ -0,0 +1,248 @@ +"""INT4/INT8 分组量化 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V、动态 scale 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ..types import StateTransferLayout + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_group_quantized_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + k_scales: torch.Tensor, + v_scales: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + *, + layout: StateTransferLayout, + write: bool, + grid_num: int = 36, + page_readies: list[bool] | None = None, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + else: + assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "v": + local = v_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "k_scale": + # 动态 scale 与 K/V 使用相同的 token/page 索引。 + local = k_scales + axes = ("layer", "token", "head", "group") + elif region.field == "v_scale": + local = v_scales + axes = ("layer", "token", "head", "group") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_slice in region.axis_slices: + # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], external.shape[3], external.shape[4] + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 及动态 scale 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_group_quantized_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=256, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_group_quantized_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in range(program_start, element_count, grid_stride): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + layer = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + layer * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + layer * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py new file mode 100644 index 0000000000..62e27b058f --- /dev/null +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ..utils import dtype_name +from ..types import AxisSlice, TransferRegion +from .base_state import FullAttentionConfig, FullAttentionState +from .mla_transfer import transfer_mla_pages + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +class MLAFullAttentionState(FullAttentionState): + """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" + + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: FullAttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + ) -> None: + """按 MLA 布局配置准备压缩 KV 存储。""" + self.config = layout_config + self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 MLA 压缩 KV 精度对应的存储方案标记。""" + return "mla_" + dtype_name(dtype=self.config.dtype) + + def initialize_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + + def get_token_bytes(self) -> int: + """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" + config = self.config + return ( + len(config.layer_ids) * config.head_num * config.head_dim + * self.storage_dtype.itemsize + ) + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> torch.Tensor: + """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" + return self.storage[self._layer_mapping[layer_id]] + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """通过 CUDA 算子将形状匹配的 MLA latent 数据写入指定物理 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + layer = self._layer_mapping[layer_id] + if not len(token_indices): + return + assert self.storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[layer]) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( + "token copies require equal one-dimensional index tensors" + ) + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( + "token indices must be integer tensors" + ) + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_readies, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: MLAFullAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" + assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( + "DP copy requires the same encoding and equivalent TP state configurations" + ) + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； + # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 + # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 + src_kv = remote_dp_full_att_state.storage.index_select(dim=1, index=src_token_indices) + dst_kv = src_kv.to(device=self.storage.device, non_blocking=True) + self.storage.index_copy_(dim=1, index=dst_token_indices, source=dst_kv) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = (("kv", self.storage_dtype, config.head_dim, "dim"),) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # field 已按 K/V 拆开：这里读取的是单个字段的本地 head 数，不是合并 buffer 的 2H。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_slices = tuple( + AxisSlice( + axis=axis, + src_start=0, + dst_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_slices=axis_slices, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.resource_id, field_name, head_start, local_head_count, replica_count, + write_enabled, field_region.byte_start, field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.layer_ids), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/att_resources/full_att/mla_transfer.py b/lightllm/common/att_resources/full_att/mla_transfer.py new file mode 100644 index 0000000000..cc763e570a --- /dev/null +++ b/lightllm/common/att_resources/full_att/mla_transfer.py @@ -0,0 +1,233 @@ +"""MLA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 latent 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ..types import StateTransferLayout + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_mla_pages( + kv_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + *, + layout: StateTransferLayout, + write: bool, + grid_num: int = 36, + page_readies: list[bool] | None = None, +) -> None: + """直接搬运 latent；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + """ + assert grid_num > 0, "grid_num must be positive" + assert kv_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + else: + assert page_buffers.device == kv_storage.device, "CUDA pages must reside on the State device" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 latent。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=kv_storage.device, + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=kv_storage.device, + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 latent 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "kv": + local = kv_storage + axes = ("layer", "token", "head", "dim") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_slice in region.axis_slices: + # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], external.shape[3], external.shape[4] + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + for local, external in copy_views: + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_mla_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=256, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_mla_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in range(program_start, element_count, grid_stride): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + layer = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + layer * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + layer * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py new file mode 100644 index 0000000000..adf5767e41 --- /dev/null +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -0,0 +1,413 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ..utils import dtype_name +from ..types import AxisSlice, TransferRegion +from .base_state import FullAttentionConfig, FullAttentionState +from .packed_mla_transfer import transfer_packed_mla_pages + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class PackedMLAAttentionConfig(FullAttentionConfig): + # MLA KV latent 的有效维度，包含需要保存的 RoPE 部分，不含 indexer、scale 或行 padding。 + kv_dim: int = 576 + # 每个 token 的 indexer key 维度，与 MLA latent 一起存放在同一物理行。 + indexer_dim: int = 128 + # 0 表示不使用 FP8 latent 编码；非零表示需要 FP8 量化的非 RoPE 部分维度。 + # 当前 FP8 实现要求 512 维非 RoPE latent，剩余 64 维 RoPE 以 BF16 保存。 + fp8_nope_dim: int = 0 + # FP8 非 RoPE latent 的分组量化宽度，当前 kernel 要求每 128 个元素对应一个 scale。 + quant_group_size: int = 128 + + def __post_init__(self) -> None: + """创建配置时校验通用 TP 约束、MLA 编码及 packed 行宽。""" + super().__post_init__() + assert self.head_num == 1 and self.indexer_dim > 0, ( + "packed MLA requires one replicated latent and a positive indexer dimension" + ) + if self.fp8_nope_dim: + assert (self.fp8_nope_dim, self.kv_dim, self.quant_group_size, self.dtype) == ( + 512, + 576, + 128, + torch.uint8, + ), "FP8 MLA kernel requires 512 latent dimensions, 64 RoPE dimensions and groups of 128" + payload_bytes = 656 + self.indexer_dim + 4 + else: + assert self.dtype in (torch.float16, torch.bfloat16), "unquantized packed MLA requires FP16 or BF16" + payload_bytes = self.kv_dim * self.dtype.itemsize + self.indexer_dim + 4 + assert self.head_dim * self.dtype.itemsize == (payload_bytes + 15) // 16 * 16, ( + "packed token width must include aligned KV and indexer storage" + ) + + +@dataclass(frozen=True) +class PackedMLALayerView: + kv: torch.Tensor + indexer_k: torch.Tensor + + +class PackedMLAFullAttentionState(FullAttentionState): + """在同一物理 token 行中存储 MLA latent、量化参数及稀疏 indexer key/scale。""" + + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: PackedMLAAttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + ) -> None: + """按 packed MLA 配置准备 KV 与 indexer 存储。""" + self.config = layout_config + self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" + encoding = "fp8" if self.config.fp8_nope_dim else dtype_name(dtype=self.config.dtype) + return "mla_packed_" + encoding + + def initialize_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + + def get_token_bytes(self) -> int: + """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" + config = self.config + return ( + len(config.layer_ids) * config.head_num * config.head_dim + * self.storage_dtype.itemsize + ) + + @property + def kv_nbytes(self) -> int: + """返回一行中 KV 区域的字节数，包含该编码需要的 scale 与 RoPE 数据。""" + return 656 if self.config.fp8_nope_dim else self.config.kv_dim * self.config.dtype.itemsize + + @property + def indexer_offset(self) -> int: + """返回一行中 indexer key 区域的字节偏移。""" + if self.config.fp8_nope_dim: + return self.kv_nbytes + return self.config.head_dim * self.config.dtype.itemsize - self.config.indexer_dim - 4 + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> PackedMLALayerView: + """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" + row = self.storage[self._layer_mapping[layer_id]] + raw = row.view(torch.uint8) + kv = raw[..., : self.kv_nbytes] if self.config.fp8_nope_dim else row[..., : self.config.kv_dim] + indexer = raw[..., self.indexer_offset : self.indexer_offset + self.config.indexer_dim + 4] + return PackedMLALayerView(kv=kv, indexer_k=indexer) + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """写入 MLA KV 及其量化参数，保留同一 token 行中的 indexer 区域。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + layer = self._layer_mapping[layer_id] + if not len(token_indices): + return + assert self.storage.is_cuda, "KV writes require CUDA storage" + if self.config.fp8_nope_dim: + from lightllm.common.basemodel.triton_kernel.kv_copy.destindex_copy_kv_flashmla_fp8 import ( + destindex_copy_kv_flashmla_fp8, + ) + + raw = self.storage[layer] + destindex_copy_kv_flashmla_fp8( + KV_nope=kv[..., :512], + KV_rope=kv[..., 512:], + DestLoc=token_indices, + O_nope=raw[..., :512].view(torch.float8_e4m3fn), + O_scale=raw[..., 512:528].view(torch.float32), + O_rope=raw[..., 528:656].view(torch.bfloat16), + ) + else: + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[layer, ..., : self.config.kv_dim]) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( + "token copies require equal one-dimensional index tensors" + ) + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( + "token indices must be integer tensors" + ) + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + kv_dim=self.config.kv_dim, + indexer_offset=self.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=True, + page_readies=cpu_page_readies, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + kv_dim=self.config.kv_dim, + indexer_offset=self.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: PackedMLAFullAttentionState, + ) -> None: + """校验 MLA 编码与区域尺寸，再复制完整 KV 和 indexer 数据。""" + assert ( + remote_dp_full_att_state.config.fp8_nope_dim, + remote_dp_full_att_state.config.kv_dim, + remote_dp_full_att_state.config.indexer_dim, + ) == ( + self.config.fp8_nope_dim, + self.config.kv_dim, + self.config.indexer_dim, + ), "incompatible packed MLA format" + assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( + "DP copy requires the same encoding and equivalent TP state configurations" + ) + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； + # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 + # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 + # 整行包含 KV、动态 scale、RoPE、indexer 及 padding，直接复制可保留全部有效字段。 + src_kv = remote_dp_full_att_state.storage.index_select(dim=1, index=src_token_indices) + dst_kv = src_kv.to(device=self.storage.device, non_blocking=True) + self.storage.index_copy_(dim=1, index=dst_token_indices, source=dst_kv) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + kv_dim=self.config.kv_dim, + indexer_offset=self.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=True, + ) + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + kv_dim=self.config.kv_dim, + indexer_offset=self.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + if config.fp8_nope_dim: + field_specs = [ + ("kv_nope", torch.uint8, 512, "dim"), + ("kv_scale", torch.float32, 4, "group"), + ("kv_rope", torch.bfloat16, 64, "dim"), + ] + else: + field_specs = [("kv", config.dtype, config.kv_dim, "dim")] + field_specs.extend(( + ("indexer_k", torch.uint8, config.indexer_dim, "dim"), + ("indexer_scale", torch.float32, 1, "scale"), + )) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # field 已按 K/V 拆开：这里读取的是单个字段的本地 head 数，不是合并 buffer 的 2H。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_slices = tuple( + AxisSlice( + axis=axis, + src_start=0, + dst_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_slices=axis_slices, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.resource_id, field_name, head_start, local_head_count, replica_count, + write_enabled, field_region.byte_start, field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.layer_ids), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/att_resources/full_att/packed_mla_transfer.py b/lightllm/common/att_resources/full_att/packed_mla_transfer.py new file mode 100644 index 0000000000..b7ef3640ee --- /dev/null +++ b/lightllm/common/att_resources/full_att/packed_mla_transfer.py @@ -0,0 +1,257 @@ +"""packed MLA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 latent 及配套字段 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ..types import StateTransferLayout + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_packed_mla_pages( + kv_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + *, + layout: StateTransferLayout, + kv_dim: int, + indexer_offset: int, + indexer_dim: int, + write: bool, + grid_num: int = 36, + page_readies: list[bool] | None = None, +) -> None: + """直接搬运 latent 及配套字段；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + """ + assert grid_num > 0, "grid_num must be positive" + assert kv_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + else: + assert page_buffers.device == kv_storage.device, "CUDA pages must reside on the State device" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 latent 及配套字段。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=kv_storage.device, + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=kv_storage.device, + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 latent 及配套字段 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + # 只创建原始字节视图；每行保留真实步长，传输仅覆盖已声明字段，不包含 padding。 + raw = kv_storage.view(torch.uint8) + indexer_end = indexer_offset + indexer_dim + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "kv": + local = kv_storage[..., :kv_dim] + axes = ("layer", "token", "head", "dim") + elif region.field == "kv_nope": + # FP8 编码、FP32 scale 和 BF16 RoPE 分别位于固定字节区间。 + local = raw[..., :512] + axes = ("layer", "token", "head", "dim") + elif region.field == "kv_scale": + local = raw[..., 512:528].view(torch.float32) + axes = ("layer", "token", "head", "group") + elif region.field == "kv_rope": + local = raw[..., 528:656].view(torch.bfloat16) + axes = ("layer", "token", "head", "dim") + elif region.field == "indexer_k": + local = raw[..., indexer_offset:indexer_end] + axes = ("layer", "token", "head", "dim") + elif region.field == "indexer_scale": + local = raw[..., indexer_end:indexer_end + 4].view(torch.float32) + axes = ("layer", "token", "head", "scale") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_slice in region.axis_slices: + # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], external.shape[3], external.shape[4] + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 KV、indexer 及各量化字段 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_packed_mla_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=256, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_packed_mla_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in range(program_start, element_count, grid_stride): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + layer = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + layer * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + layer * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py new file mode 100644 index 0000000000..dc36d675ab --- /dev/null +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -0,0 +1,381 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ..types import AxisSlice, TransferRegion +from .base_state import FullAttentionConfig, FullAttentionState +from .static_fp8_transfer import transfer_static_fp8_pages + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class StaticFP8AttentionConfig(FullAttentionConfig): + # per_tensor：每层 K/V 各一个固定 scale；per_head：每层、每个本地 K/V head 各一个。 + # 标定参数不随 token 增长，也不作为动态 KV 字段写入 CPU/PD 页。 + scale_mode: str = "per_tensor" + # 标定数据的身份标记，用于 DP 复制时检查编码是否兼容；实际 scale Tensor 由 State 持有。 + calibration_id: str = "" + + def __post_init__(self) -> None: + """创建配置时校验通用 TP 约束、标定方式及身份。""" + super().__post_init__() + assert self.scale_mode in ("per_tensor", "per_head") and self.calibration_id, ( + "static FP8 requires a scale mode and calibration identity" + ) + + +class StaticFP8FullAttentionState(FullAttentionState): + """使用静态标定参数的 FP8 KV 存储，固定 scale 随 State 初始化，不逐 token 搬运。""" + + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: StaticFP8AttentionConfig, + page_size: int = 1, + *, + device: str | torch.device = "cuda", + scales: torch.Tensor, + q_scales: torch.Tensor | None = None, + ) -> None: + """按静态 FP8 配置准备 KV 存储，并保留传入的标定参数。""" + # 标定 Tensor 由 State 持有，在接收构造参数时校验，不重复放入 config。 + heads = 1 if layout_config.scale_mode == "per_tensor" else layout_config.head_num + assert scales is not None and scales.shape == (len(layout_config.layer_ids), 2 * heads), ( + "calibration scales do not match KV layers and heads" + ) + assert scales.dtype == torch.float32 and bool((torch.isfinite(input=scales) & (scales > 0)).all()), ( + "calibration scales must be finite positive float32 values" + ) + if q_scales is not None: + assert q_scales.shape == (len(layout_config.layer_ids), layout_config.head_num) and ( + q_scales.dtype == torch.float32 + ), "Q calibration does not match layers and heads" + self.scales = scales + self.q_scales = q_scales + self.cpu_scales = None # 实际存储初始化时准备计算 backend 使用的 CPU 标定值。 + self.config = layout_config + self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + self.storage_dtype = torch.uint8 + + def get_state_type(self) -> str: + """返回包含 per-tensor 或 per-head 标定方式的 FP8 方案标记。""" + return "full_att_fp8_static_" + self.config.scale_mode + + def initialize_storage(self, token_capacity: int) -> None: + """绑定 KV 与真实标定参数，并预备 FlashInfer 使用的 CPU scale。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + # meta 准备阶段保留真实标定值，实际分配时再迁移到 KV 所在设备。 + if self.k_storage.device.type != "meta": + self.scales = self.scales.to(device=self.k_storage.device) + if self.q_scales is not None: + self.q_scales = self.q_scales.to(device=self.k_storage.device) + # 提前准备 CPU scale，避免计算热路径同步设备。 + self.cpu_scales = self.scales.detach().cpu().numpy() + + def get_token_bytes(self) -> int: + """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" + config = self.config + return ( + len(config.layer_ids) * 2 * config.head_num * config.head_dim + * self.storage_dtype.itemsize + ) + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" + layer = self._layer_mapping[layer_id] + return self.k_storage[layer], self.v_storage[layer] + + def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """按固定 per-tensor 或 per-head scale 量化输入 K/V，写入 FP8 编码。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + layer = self._layer_mapping[layer_id] + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 + + destindex_copy_kv_fp8( + KV=kv, + DestLoc=token_indices, + scales=self.scales[layer], + Out=self.k_storage[layer].view(torch.float8_e4m3fn), + is_per_tensor_quant=self.config.scale_mode == "per_tensor", + Out_v=self.v_storage[layer].view(torch.float8_e4m3fn), + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( + "token copies require equal one-dimensional index tensors" + ) + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( + "token indices must be integer tensors" + ) + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 FP8 K/V 原始编码 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + *, + layout: StateTransferLayout, + cpu_page_readies: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_readies, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + *, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_other_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: StaticFP8FullAttentionState, + ) -> None: + """先校验标定方式与标定身份，再复制兼容的 FP8 KV 编码。""" + assert ( + remote_dp_full_att_state.config.scale_mode, + remote_dp_full_att_state.config.calibration_id, + ) == ( + self.config.scale_mode, + self.config.calibration_id, + ), "incompatible static FP8 calibration" + assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( + "DP copy requires the same encoding and equivalent TP state configurations" + ) + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； + # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 + # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 + src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) + dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) + self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) + src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) + dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) + self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_buffer( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_buffer( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + *, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", self.storage_dtype, config.head_dim, "dim"), + ("v", self.storage_dtype, config.head_dim, "dim"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_slices = tuple( + AxisSlice( + axis=axis, + src_start=0, + dst_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_slices=axis_slices, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.resource_id, field_name, head_start, local_head_count, replica_count, + write_enabled, field_region.byte_start, field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.layer_ids), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/att_resources/full_att/static_fp8_transfer.py b/lightllm/common/att_resources/full_att/static_fp8_transfer.py new file mode 100644 index 0000000000..107ec9bda7 --- /dev/null +++ b/lightllm/common/att_resources/full_att/static_fp8_transfer.py @@ -0,0 +1,240 @@ +"""静态 FP8 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ..types import StateTransferLayout + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_static_fp8_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + *, + layout: StateTransferLayout, + write: bool, + grid_num: int = 36, + page_readies: list[bool] | None = None, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + else: + assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=k_storage.device, + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + # 静态 KV/Q 标定参数属于模型配置，不随 token 页面传输。 + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "v": + local = v_storage + axes = ("layer", "token", "head", "dim") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_slice in region.axis_slices: + # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], external.shape[3], external.shape[4] + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 FP8 K/V 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_static_fp8_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=256, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_static_fp8_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in range(program_start, element_count, grid_stride): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + layer = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + layer * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + layer * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/att_resources/types.py b/lightllm/common/att_resources/types.py new file mode 100644 index 0000000000..3972edd4f5 --- /dev/null +++ b/lightllm/common/att_resources/types.py @@ -0,0 +1,216 @@ +"""注意力存储字段与传输位置的只读描述，不执行数据搬运。 + +TensorField 描述本地 Tensor；AxisSlice 描述某个轴的源、目标范围。 +TransferRegion / StateTransferLayout 用于 State 与 CPU/PD 页之间的搬运； +DP 在对应 TP 分片间直接复制，不需要字段映射类型。 + +byte_start、byte_end 使用字节单位，表示左闭右开的区间；shape、轴切片起点和 length 使用元素单位。 +层维度完整、同序对应，token/request 索引由每次搬运的调用参数提供。 +frozen=True 只禁止重新赋值描述属性，不冻结被引用 Tensor 的数据。 +KV 示例统一使用 head_dim=128。 + +PageKind 使用示例，选择 bounded attention 大页池： + >>> page_kind: PageKind = "big" + >>> assert page_kind == "big" +""" + +from dataclasses import dataclass +from typing import Literal + +import torch + +# ======================================================================== +# 一、状态页类别 +# ======================================================================== + +# 选择 bounded attention 的页池及共享分配域，不表示 KV token 页或传输 buffer 的大小。 +# "big" 对应 big_pages_pool 与 bounded_att_big_page_allocator； +# "small" 对应 small_pages_pool 与 bounded_att_small_page_allocator。 +# 两个域中的同一个 slot_id 指向不同状态页。调用方先在对应 allocator 分配槽位， +# 再将类别传给 save_bounded_att_pages / restore_bounded_att_pages 等操作，例如： +# manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[3]) +# 上例假定请求槽位 0 的已接受长度为 128，且大页槽位 3 已由调用方分配。 +PageKind = Literal["big", "small"] + + +# ======================================================================== +# 二、本地存储字段与轴范围 +# ======================================================================== + + +@dataclass(frozen=True) +class TensorField: + """为本地 Tensor 补充轴语义，供 State 内部布局计算及搬运使用。 + + 创建描述只引用已有 Tensor，不分配或复制其底层存储。具体 State 按固定布局组织 + k、v、动态量化 scale、Conv、SSM 等搬运字段，字段名称由所在字典的 key 提供； + packed buffer 的字段可使用切片视图。 + + 示例：本地 K 数据有 2 层、8 个 token、2 个 heads，每个 head 有 128 个元素。 + >>> k = torch.empty((2, 8, 2, 128), dtype=torch.float16) + >>> field = TensorField( + ... tensor=k, axes=("layer", "token", "head", "dim") + ... ) + >>> assert field.tensor is k + """ + + # 实际存储或其视图，提供 dtype、shape、stride、device 等物理信息。 + # 可引用非连续的 packed/分层存储视图，搬运代码沿用其真实 stride。 + tensor: torch.Tensor + + # 按 Tensor 维度顺序列出语义名称，数量等于 tensor.ndim，名称不能重复。 + # 例如 token/head 用于 KV 选取和 TP 切片，request/page 用于运行态或页池描述； + # 选择某个请求或状态页后，搬运视图会移除已选定的 request/page 轴。 + axes: tuple[str, ...] + + +@dataclass(frozen=True) +class AxisSlice: + """某一轴上的连续范围映射，起点与长度均以元素为单位。 + + TransferRegion 中的方向固定为本地 State 到外部页，读取时反向应用。 + + 示例：本地 head[0:2] 写到全局传输区域 head[2:4]，读回时反向复制。 + >>> head_slice = AxisSlice(axis="head", src_start=0, dst_start=2, length=2) + >>> assert head_slice.dst_start + head_slice.length == 4 + """ + + # 要切片的轴名，必须同时存在于源、目标字段中，例如 head、channel、dim。 + # layer 固定完整复制，token 由本次传输数量控制，这两个轴不放入 AxisSlice。 + axis: str + + # 源轴的起始元素索引，包含该位置且必须非负。 + # TransferRegion 中本地字段的位置。 + src_start: int + + # 目标轴的起始元素索引，包含该位置且必须非负。 + # TransferRegion 中外部页区域的位置。 + # 用于将本地 TP 分片放到全局 head/channel 的对应位置。 + dst_start: int + + # 搬运的元素数量，必须为正；源和目标的半开区间均为 [start, start + length)。 + # 它是该轴上的长度，不是字节数，也不是整个 Tensor 的元素总数。 + length: int + + +# ======================================================================== +# 三、CPU/PD 外部页布局 +# ======================================================================== + + +@dataclass(frozen=True) +class TransferRegion: + """描述 State 内一个字段与 CPU/PD 外部页区域之间的拷贝映射。 + + 由 StateTransferLayout.regions 组织，同一字段可以对应多个区域映射。 + 区域的 shape 覆盖外部完整数据，axis_slices 选择当前 rank 的子范围。 + 本地和外部区域的层数必须一致，并按相同存储顺序完整复制。 + + 示例：外部 K 区域有 4 个全局 heads，head_dim=128，当前 rank 持有后 2 个 heads。 + 整个区域占 2 * 8 * 4 * 128 * 2 = 16384 字节（16 KiB），位于页内 [16384, 32768)。 + >>> k_region = TransferRegion( + ... field="k", + ... byte_start=16384, byte_end=32768, + ... dtype=torch.float16, + ... axes=("layer", "token", "head", "dim"), + ... shape=(2, 8, 4, 128), + ... axis_slices=( + ... AxisSlice(axis="head", src_start=0, dst_start=2, length=2), + ... AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + ... ), + ... write_enabled=True, + ... ) + """ + + # 本地搬运字段字典的 key，用来定位 k/v/scale/conv/ssm 等 Tensor。 + # 同一字段可出现在多个 region 中，例如 Conv 的 Q/K/V 使用不同 channel 切片。 + field: str + + # 字段区域起点，包含该字节，必须按 dtype 的元素大小对齐。 + # State 的 get_cpu/pd_transfer_regions 返回相对本 State 起点的区间； + # 管理器将起止位置同时加上 State 的页内起点，最终区间统一相对整页。 + # CPU/PD 都从页集合的第一维选择物理页；字段区间不包含物理页索引。 + byte_start: int + + # 字段区域终点，不包含该字节，与 byte_start 使用相同的坐标基准。 + # byte_end - byte_start 必须等于 prod(shape) * dtype.itemsize， + # 仅包含字段数据；字段之间的对齐空隙不计入该区间。 + byte_end: int + + # 外部区域按此类型解释原始字节，需与本地搬运字段的 dtype 一致。 + # packed 字段可分别使用 uint8、float16 等实际类型；此描述不执行类型转换。 + dtype: torch.dtype + + # 外部区域各维度的语义顺序，与 shape 一一对应且名称唯一。 + # 可与本地 Tensor 的轴顺序不同，搬运时按名称对齐；本地视图需具有相同轴集合。 + axes: tuple[str, ...] + + # 完整外部区域的各轴元素数量，不包含外部页池的 page 轴，也不是本 rank 的局部形状。 + # 例如全局 4 heads、本地 2 heads 时，这里的 head 维仍为 4。 + # 区域连续排列，stride 由 shape 和 dtype 推导；区域末尾必须位于当前页内。 + shape: tuple[int, ...] + + # 描述本地字段到外部区域的轴范围，每个非 layer/token 轴必须恰好出现一次。 + # 无分片的轴也显式填写完整范围，例如 dim: 0 -> 0、length=head_dim。 + # 读回时交换源、目标起点；同一组切片应用于全部层和本次有效 token。 + axis_slices: tuple[AxisSlice, ...] + + # 当前 TP rank 是否负责向外部区域写入；False 不影响从该区域读回。 + # KV head 有 TP 副本时，每组只选一个 rank 写入；复制的 indexer tail 只由 rank 0 写。 + # 标记属于 region，因为同一 State 内各字段的 TP 分片/复制规则可能不同。 + write_enabled: bool = True + + +@dataclass(frozen=True) +class StateTransferLayout: + """描述一个 State 的 CPU/PD 传输布局，按字段组织 TransferRegion。 + + 每个 layout 通过 resource_id 绑定一个 State；regions 描述该 State 各字段 + 与外部页之间的映射，具体轴范围由各 region 的 AxisSlice 描述。 + byte_start / byte_end 记录本 State 在联合页中的完整预留区间，包含自身对齐空间。 + CPU/PD 根据页集合的行选择物理页；PD 通过显式 pd_page_index 指定目标页。 + 实际页容量由 manager 统一计算，可以大于当前 State 的终点。 + 示例：一张 64 KiB 的页中，本 State 的 K/V 各占 16 KiB， + 分别位于 [16 KiB, 32 KiB) 和 [32 KiB, 48 KiB)，前后空间可属于其他资源。 + >>> from dataclasses import replace + >>> k_region = TransferRegion( + ... field="k", byte_start=16384, byte_end=32768, dtype=torch.float16, + ... axes=("layer", "token", "head", "dim"), shape=(2, 8, 4, 128), + ... axis_slices=( + ... AxisSlice(axis="head", src_start=0, dst_start=2, length=2), + ... AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + ... ), + ... ) + >>> layout = StateTransferLayout( + ... resource_id="main_full_att_fp16_state", + ... tokens_per_page=8, + ... byte_start=16384, byte_end=49152, + ... regions=(k_region, replace(k_region, field="v", byte_start=32768, byte_end=49152)), + ... ) + >>> assert layout.regions[1].byte_end == layout.byte_end == 49152 + """ + + # 当前 State 的规范身份，与管理器成员名、state.resource_id 一致。 + # 用于从联合页布局字典中选择对应资源，区分 main/draft 及不同存储方案。 + resource_id: str + + # 整张外部页可容纳的 token 数，不是本次实际传输数量；含 KV 的页使用正整数。 + # CPU/PD 分别使用各自的全局页配置，页容量不随本次传输长度改变。 + # CPU 的 KV + bounded 联合页中，bounded layout 也沿用该 token 容量。 + # 纯运行态页（如 PD 的 att_state 页）为 None，表示按快照而非 token 序列搬运。 + tokens_per_page: int | None + + # 本 State 在外部页中的起始字节，包含该位置，始终相对整页起点。 + # 多个 State 依次排列，后一 State 的起点为前一 State 的终点。 + # regions 已使用页内坐标，搬运时不再重复加上此起点。 + byte_start: int + + # 本 State 在外部页中的结束字节，不包含该位置；涵盖全部字段及本 State 的对齐、预留空间。 + # byte_end - byte_start 是本 State 的占用，不是整张物理页的跨度。 + # 同类布局的最大 byte_end 表示整组 State 的占用；CPU 物理页再按 16 字节对齐， + # PD 物理页由 manager.get_pd_page_bytes() 取 KV 与完整运行态布局占用的最大值。 + byte_end: int + + # 当前 State 参与传输的字段区域；其他 State 的区域放在它们各自的 layout 中。 + # 包含所需的 KV、动态量化参数或运行态字段，所有偏移均相对同一张完整外部页。 + regions: tuple[TransferRegion, ...] diff --git a/lightllm/common/att_resources/utils.py b/lightllm/common/att_resources/utils.py new file mode 100644 index 0000000000..669b40256a --- /dev/null +++ b/lightllm/common/att_resources/utils.py @@ -0,0 +1,16 @@ +"""注意力资源的 State 命名工具。""" + +from __future__ import annotations + +import torch + +MAX_DRAFT_MODELS = 36 +MODEL_ROLES = ("main", *(f"draft{i}" for i in range(MAX_DRAFT_MODELS))) + + +# ======================================================================== +# 一、State 命名与成员管理 +# ======================================================================== + +def dtype_name(dtype: torch.dtype): + return {torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32"}[dtype] diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index ee8d06ed9a..938a294f6b 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -4,12 +4,13 @@ import torch from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Optional, TYPE_CHECKING, Tuple, Union, Dict +from typing import Optional, TYPE_CHECKING, Tuple, Dict from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: + from lightllm.common.att_resources.full_att import FullAttentionState from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -76,14 +77,13 @@ def uses_causal_attention(self) -> bool: return not is_parallel_block_draft def _find_layer_index( - self, k: torch.Tensor, v: torch.Tensor, att_state: Union["BasePrefillAttState", "BaseDecodeAttState"] + self, k: torch.Tensor, v: torch.Tensor, state: "FullAttentionState" ) -> int: - kv_buffer = att_state.infer_state.mem_manager.kv_buffer - layer_count = len(kv_buffer) - find_dict = {kv_buffer[i].data_ptr(): i for i in range(layer_count)} - key = min(k.data_ptr(), v.data_ptr()) - assert key in find_dict - return find_dict[key] + for layer_index, layer_id in enumerate(state.config.layer_ids): + layer_k, layer_v = state.get_layer_view(layer_id) + if k.data_ptr() == layer_k.data_ptr() and v.data_ptr() == layer_v.data_ptr(): + return layer_index + raise AssertionError("K/V views do not belong to the same State layer") @dataclass diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index 01633e294b..c82de810b8 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -39,10 +39,13 @@ class Fp8Fa3PrefillAttState(Fa3PrefillAttState): def init_state(self): super().init_state() batch_size = self.infer_state.batch_size - mem_manager = self.backend.model.mem_manager + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + ) - offline_scales: torch.Tensor = mem_manager.scales - head_num = mem_manager.head_num + offline_scales: torch.Tensor = state.scales + head_num = state.config.head_num # 为了减少推理计算量，在推理外部初始化k_descale和v_descale self.k_descale = ( offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num) @@ -82,8 +85,12 @@ def _fp8_prefill_att( k_head_dim = k.shape[2] cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self) - static_q_scales = self.backend.model.mem_manager.q_scales[layer_index] + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + ) + layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, state=state) + static_q_scales = state.q_scales[layer_index] q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales) q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) o = flash_attn_with_kvcache( @@ -114,12 +121,15 @@ class Fp8Fa3DecodeAttState(Fa3DecodeAttState): def init_state(self): self.backend: Fp8Fa3AttBackend = self.backend - mem_manager = self.backend.model.mem_manager + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + ) super().init_state() att_batch_size = self.b_att_seq_len.shape[0] - offline_scales: torch.Tensor = mem_manager.scales - head_num = mem_manager.head_num + offline_scales: torch.Tensor = state.scales + head_num = state.config.head_num # 为了减少推理计算量，在推理外部初始化k_descale和v_descale self.k_descale = ( @@ -167,11 +177,15 @@ def _fp8_decode_att( cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self) + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + ) + layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, state=state) q_head_num = q.shape[1] att_batch_size = self.b_att_seq_len.shape[0] - static_q_scales = self.backend.model.mem_manager.q_scales + static_q_scales = state.q_scales q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales[layer_index]) q_scale = static_q_scales[layer_index].view(1, k_head_num).expand(att_batch_size, k_head_num) o = flash_attn_with_kvcache( diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 189f8daddb..28a22733cd 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -66,9 +66,13 @@ def _fp8_prefill_att( o_tensor = alloc_func(q.shape, q.dtype, device="cuda") k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self) - k_descale = self.infer_state.mem_manager.cpu_scales[layer_index][0] - v_descale = self.infer_state.mem_manager.cpu_scales[layer_index][1] + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", + ) + layer_index = self.backend._find_layer_index(k=k, v=v, state=state) + k_descale = state.cpu_scales[layer_index][0] + v_descale = state.cpu_scales[layer_index][1] self.prefill_wrapper.run( q, (k, v), @@ -118,10 +122,14 @@ def _fp8_decode_att( k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self) + state = getattr( + self.infer_state.att_resource_manager, + f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", + ) + layer_index = self.backend._find_layer_index(k=k, v=v, state=state) - k_descale = self.infer_state.mem_manager.cpu_scales[layer_index][0] - v_descale = self.infer_state.mem_manager.cpu_scales[layer_index][1] + k_descale = state.cpu_scales[layer_index][0] + v_descale = state.cpu_scales[layer_index][1] self.decode_wrapper.run( q, (k, v), diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 2cde561013..1d487e05b9 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -28,6 +28,8 @@ class LinearAttBackend(BaseAttBackend, ABC): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) + role = model.att_resource_role + self.linear_state = getattr(model.att_resource_manager, f"{role}_gdn_state") self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) self.prefill_kernel = self.get_prefill_kernel() @@ -139,7 +141,7 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: LinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -227,7 +229,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.backend.linear_state.mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -249,7 +251,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = self.backend.linear_state.mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -281,7 +283,7 @@ def decode_att( backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 3d0f0cf5f4..5d359d9f8f 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -37,6 +37,8 @@ class KDALinearAttBackend(BaseAttBackend): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) + role = model.att_resource_role + self.linear_state = getattr(model.att_resource_manager, f"{role}_kda_state") config = model.config["linear_attn_config"] self.num_heads = config["num_heads"] self.head_dim = config["head_dim"] @@ -83,7 +85,7 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -151,7 +153,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.backend.linear_state.mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -173,7 +175,7 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = self.backend.linear_state.mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -204,7 +206,7 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) conv_update = causal_conv1d_update if backend.mtp_step > 0: diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py index 29c1d0fb56..10ad62d6b8 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py @@ -84,11 +84,16 @@ def _nsa_prefill_att( # 当前推理生成的token kv部分从 prefill_cache_kv 中获取，历史 # 部分kv 从 packed_kv 中获取, 并进行反量化，这样可以避免 prefill_cache_kv # 部分的数据进行重复的反量化操作，提升整体的性能。 - kv, topk_indices = self.infer_state.mem_manager.get_prefill_kv_cache_and_remap_indices( + 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_indices=topk_mem_indices, + topk_mem_indices=topk_mem_indices, prefill_mem_index=self.infer_state.mem_index, prefill_cache_kv=prefill_cache_kv, + prefill_dtype=prefill_cache_kv.dtype, ) else: kv = prefill_cache_kv diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..df4730d69d 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -13,8 +13,6 @@ from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.kv_cache_mem_manager import MemoryManager -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.req_manager import ReqManager from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( @@ -27,12 +25,10 @@ from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed from lightllm.utils.log_utils import init_logger -from lightllm.utils.dist_utils import get_dp_world_size -from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory +from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp from lightllm.utils.envs_utils import ( get_env_start_args, get_llm_data_type, - get_added_mtp_kv_layer_num, ) from lightllm.distributed.communication_op import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput @@ -41,7 +37,6 @@ ) from lightllm.common.basemodel.mtp_manager import MtpManager from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch -from lightllm.utils.envs_utils import set_model_init_status from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.infer_utils import post_empty_cache from lightllm.utils.torch_memory_saver_utils import ( @@ -70,6 +65,7 @@ class TpPartBaseModel: # infer state class infer_state_class = InferStateInfo + is_hybrid_att_model = False def __init__(self, kvargs): self.args = get_env_start_args() @@ -97,6 +93,7 @@ def __init__(self, kvargs): self.expert_dtype = kvargs.get("expert_dtype", None) self.mem_fraction = kvargs.get("mem_fraction", 0.9) self.tp_world_size_ = get_dp_world_size() + self.tp_rank_ = get_current_rank_in_dp() self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode self.graph_max_batch_size = self._align_decode_batch_size(self.graph_max_batch_size) @@ -109,41 +106,63 @@ def __init__(self, kvargs): self._init_quant() enable_weight_cpu_backup = self.args.enable_weight_cpu_backup - with profile_mtp_weight_memory(self): - with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): - self._init_weights() + with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): + self._init_weights() with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): self._init_req_manager() - self._init_mem_manager() - - # Qwen3.5 等 linear attention 模型会在 req_manager 中保存大量运行时 state。先初始化 - # req_manager、再初始化 mem_manager，可以让 KV cache 显存评估包含这些 state 的实际占用； - # 因此 req_manager 创建时暂不传入 mem_manager，需要在两者初始化完成后再进行绑定。 - self.req_manager.bind_mem_manager(self.mem_manager) - self._check_mem_size() self._init_infer_layer() self._init_some_value() self._init_custom() self.load_weights(self.weight_dict) + # The backend calls each model's runtime stages after all weights are loaded. + self.graph = None + return + + @property + def att_resource_role(self) -> str: + """当前模型的资源前缀；每个 draft 独立拥有 State。""" + from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + + if not self.is_mtp_draft_model: + return "main" + draft_index = len(self.mtp_previous_draft_models) + assert draft_index < MAX_DRAFT_MODELS, "at most 36 draft models are supported" + return f"draft{draft_index}" + + def prepare_att_resources(self) -> None: + """Prepare this model's attention States from its own configuration.""" + assert getattr(self, "att_resource_manager", None) is None, "resources already initialized" + if self.is_mtp_draft_model: + self.att_resource_manager = self.main_model.att_resource_manager + assert not self.att_resource_manager._states_created, "State creation is already finished" + assert self.max_req_num + 1 == self.att_resource_manager._request_capacity, "models must share the request domain" + with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + self._init_att_resource_manager() + def init_runtime(self) -> None: + """Initialize this model's backends, warmup and graphs after shared allocation.""" + self.req_manager.bind_hold_page(self.att_resource_manager.full_att_allocator.hold_token_indices) + self._check_mem_size() self._init_att_backend() self._init_att_backend1() - - logger.info(f"use prefill att backend: {self.prefill_att_backend.__class__.__name__}") - logger.info(f"use decode att backend: {self.decode_att_backend.__class__.__name__}") + self._init_hidden_collector() + role = self.att_resource_role + logger.info( + f"{role} prefill/decode attention: {type(self.prefill_att_backend).__name__} / " + f"{type(self.decode_att_backend).__name__}" + ) if self.prefill_att_backend1 is not None: - logger.info(f"use prefill att backend1: {self.prefill_att_backend1.__class__.__name__}") - logger.info(f"use decode att backend1: {self.decode_att_backend1.__class__.__name__}") + logger.info( + f"{role} prefill/decode attention1: {type(self.prefill_att_backend1).__name__} / " + f"{type(self.decode_att_backend1).__name__}" + ) - self._init_hidden_collector() self._autotune_warmup() self._init_padded_req() self._init_cudagraph() self._init_prefill_cuda_graph() - self._check_max_len_infer() - torch.cuda.empty_cache() - set_model_init_status(True) - return + if not self.is_mtp_draft_model: + self._check_max_len_infer() def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: @@ -196,20 +215,12 @@ def load_weights(self, weight_dict: dict): [weight.verify_load() for weight in self.trans_layers_weight] return - def _init_mem_manager(self): - assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 - self.mem_manager: MemoryManager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_attention_heads"] // self.tp_world_size_, - head_dim=self.config["n_embed"] // self.config["num_attention_heads"], - layer_num=self.config["n_layer"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + def _init_att_resource_manager(self) -> None: + """Prepare all owned attention States using model-specific configuration.""" + raise NotImplementedError def _check_mem_size(self): - self.max_total_token_num = self.mem_manager.size + self.max_total_token_num = self.att_resource_manager.full_att_allocator.size assert ( self.max_total_token_num > self.batch_max_tokens @@ -230,6 +241,9 @@ def _check_mem_size(self): return def _init_req_manager(self): + if self.is_mtp_draft_model: + self.req_manager = self.main_model.req_manager + return create_max_seq_len = 0 if self.batch_max_tokens is not None: @@ -237,7 +251,7 @@ def _init_req_manager(self): if self.max_seq_length is not None: create_max_seq_len = max(create_max_seq_len, self.max_seq_length) - self.req_manager = ReqManager(self.max_req_num, create_max_seq_len, None) + self.req_manager = ReqManager(self.max_req_num, create_max_seq_len) return def _init_infer_layer(self, start_layer_index=0): @@ -260,6 +274,11 @@ def _init_some_value(self): return def _init_att_backend(self): + role = self.att_resource_role + if not self.att_resource_manager.get_full_attention_states(role): + # FFN-only and linear-only models do not need a full-attention backend. + self.prefill_att_backend = self.decode_att_backend = None + return self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0)(model=self) self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) return @@ -352,6 +371,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.is_draft_model = self.is_mtp_draft_model + infer_state.att_resource_role = self.att_resource_role infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num @@ -374,8 +394,8 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.multimodal_params = model_input.multimodal_params - infer_state.mem_manager = self.mem_manager infer_state.req_manager = self.req_manager + infer_state.att_resource_manager = self.att_resource_manager infer_state.mem_index = self._select_mem_indexes(model_input) infer_state.microbatch_index = microbatch_index @@ -385,13 +405,21 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens if infer_state.is_prefill: - infer_state.prefill_att_state = self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) + infer_state.prefill_att_state = ( + self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) + if self.prefill_att_backend is not None + else None + ) if self.prefill_att_backend1 is not None: infer_state.prefill_att_state1 = self.prefill_att_backend1.create_att_prefill_state( infer_state=infer_state ) else: - infer_state.decode_att_state = self.decode_att_backend.create_att_decode_state(infer_state=infer_state) + infer_state.decode_att_state = ( + self.decode_att_backend.create_att_decode_state(infer_state=infer_state) + if self.decode_att_backend is not None + else None + ) if self.decode_att_backend1 is not None: infer_state.decode_att_state1 = self.decode_att_backend1.create_att_decode_state( infer_state=infer_state @@ -988,9 +1016,9 @@ def _check_max_len_infer(self): logger.info("begin check max_len infer") dummy_input_ids = torch.ones(self.batch_max_tokens, dtype=torch.int64, device="cuda") b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.mem_manager.page_size + page_size = self.att_resource_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() + mem_indexes = self.att_resource_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = self.batch_max_tokens @@ -1024,7 +1052,7 @@ def _check_max_len_infer(self): torch.argmax(prob_out, dim=1, keepdim=True) prob_out = None self.req_manager.free_all() - self.mem_manager.free_all() + self.att_resource_manager.full_att_allocator.free_all() logger.info(f"check max_len {self.batch_max_tokens} infer ok") except (RuntimeError, torch.OutOfMemoryError) as e: logger.exception(str(e)) @@ -1069,9 +1097,9 @@ def _autotune_warmup(self): 0, 10000, (input_len,), dtype=torch.int64, device="cuda", generator=rand_gen ) b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.mem_manager.page_size + page_size = self.att_resource_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() + mem_indexes = self.att_resource_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = input_len @@ -1105,14 +1133,14 @@ def _autotune_warmup(self): ) del model_output self.req_manager.free_all() - self.mem_manager.free_all() + self.att_resource_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() except Exception as e: logger.warning(f"autotune warmup for length {input_len} failed: {str(e)}") logger.exception(str(e)) self.req_manager.free_all() - self.mem_manager.free_all() + self.att_resource_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py new file mode 100644 index 0000000000..c7af388b98 --- /dev/null +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -0,0 +1,153 @@ +"""Shared helpers for model resource definitions used by startup and CPU caches.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from lightllm.common.build_utils import repair_config + +if TYPE_CHECKING: + from collections.abc import Sequence + from easydict import EasyDict + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.server.core.objs.start_args_type import StartArgs + + +def get_resource_config(model_config: dict[str, object], *, model_class: type[TpPartBaseModel] | None = None): + """Normalize resource config, applying a model-specific override when present.""" + build_config = getattr(model_class, "_resource_config", None) + if build_config is not None: + return build_config(model_config) + config = model_config.get("text_config", model_config).copy() + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + return config + + +def get_full_attention_options( + config: dict[str, object], + args: StartArgs | EasyDict, + *, + model_class: type[TpPartBaseModel] | None = None, + model_layer_ids: Sequence[int] | None = None, +): + """Return default KV geometry or the model's specialized storage options.""" + build_options = getattr(model_class, "_full_attention_options", None) + if build_options is not None: + options = build_options(config, args) + else: + heads = config.get("num_key_value_heads", config["num_attention_heads"]) + dim = config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]) + layers = config["num_hidden_layers"] + options = dict(heads=heads, dim=dim, layer_ids=tuple(range(layers))) + if model_layer_ids is not None: + options["layer_ids"] = tuple(model_layer_ids[i] for i in options["layer_ids"] if i < len(model_layer_ids)) + return options + + +def create_full_attention_definition( + model_class: type[TpPartBaseModel], + model_config: dict[str, object], + args: StartArgs | EasyDict, + *, + tp_rank: int = 0, + tp_world_size: int = 1, + model_layer_ids: Sequence[int] | None = None, + calibration_layer_start: int | None = None, +): + """Build a meta State using the model's configuration and storage hooks.""" + from lightllm.utils.att_resource_config import build_full_attention_state + + config = get_resource_config(model_config, model_class=model_class) + options = get_full_attention_options(config, args, model_class=model_class, model_layer_ids=model_layer_ids) + if not options["layer_ids"]: + return None + state = build_full_attention_state( + args, + **options, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + calibration_layer_start=calibration_layer_start, + ) + return state + + +def create_bounded_attention_definition( + model_class: type[TpPartBaseModel], + model_config: dict[str, object], + args: StartArgs | EasyDict, + *, + model_layer_ids: Sequence[int], + tp_rank: int = 0, + tp_world_size: int = 1, +): + """Use the model's optional bounded-attention definition, including inherited implementations.""" + create_definition = getattr(model_class, "create_bounded_attention_definition", None) + if create_definition is None: + return None + return create_definition( + model_config, args, model_layer_ids=model_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size + ) + + +def get_model_calibration_layer_start(model: TpPartBaseModel) -> int: + """Count preceding models' full-attention layers without initializing their States.""" + if not model.is_mtp_draft_model: + return 0 + layer_start = 0 + for previous in (model.main_model, *model.mtp_previous_draft_models): + if previous.is_mtp_draft_model and previous.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + continue + config = get_resource_config(previous.config, model_class=type(previous)) + options = get_full_attention_options( + config, + previous.args, + model_class=type(previous), + model_layer_ids=tuple(layer.layer_num_ for layer in previous.layers_infer), + ) + layer_start += len(options["layer_ids"]) + return layer_start + + +def init_model_full_attention(model: TpPartBaseModel) -> None: + """按当前模型的配置和层号创建 full-attention meta State，直接挂到共享管理器。""" + if model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + return + definition = create_full_attention_definition( + type(model), + model.config, + model.args, + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + calibration_layer_start=get_model_calibration_layer_start(model), + ) + if definition is None: + return + model.att_resource_manager.set_att_state(model.att_resource_role, definition) + # 层和 indexer 只记录规范成员名，推理时直接访问管理器中已分配的 State。 + name = definition.resource_id + for layer in model.layers_infer: + layer_id = layer.kv_share_target_layer_ if getattr(layer, "is_kv_shared_", False) else layer.layer_num_ + if layer_id in definition.config.layer_ids: + layer.full_att_state_name = name + if hasattr(layer, "indexer"): + layer.indexer.full_att_state_name = name + + +def init_model_bounded_attention(model: TpPartBaseModel) -> None: + """按当前模型的配置创建请求运行态，并直接挂到共享管理器；大小页稍后分配。""" + definition = create_bounded_attention_definition( + type(model), model.config, model.args, + model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, tp_world_size=model.tp_world_size_, + ) + if definition is None: + return + from lightllm.common.att_resources.bounded_att import create_gdn_state, create_kda_state + + create_state = {"gdn": create_gdn_state, "kda": create_kda_state}[definition.get_state_type()] + state = create_state( + config=definition.config, request_capacity=model.max_req_num + 1, mtp_step=model.args.mtp_step, device="cuda" + ) + model.att_resource_manager.set_att_state(model.att_resource_role, state) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ee888c1a00..d79ef0e245 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -292,7 +292,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.mem_manager.free_all() + model.att_resource_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -354,7 +354,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1]) - model.mem_manager.free_all() + model.att_resource_manager.full_att_allocator.free_all() model.req_manager.free_all() del decode_batches diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index bf25149450..08af85b649 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,7 +1,7 @@ import torch import triton import collections -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.att_resources.att_resource_manager import AttResourceManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup from typing import Tuple, Any, Optional, List @@ -47,7 +47,8 @@ def __init__(self): self.is_prefill: bool = None self.is_draft_model: bool = False - self.mem_manager: MemoryManager = None + self.att_resource_manager: AttResourceManager = None + self.att_resource_role: str = "main" # main 或 draft0 .. draft35，区分共享管理器中的 State。 self.req_manager: ReqManager = None self.mem_index: torch.Tensor = None diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index f0cc129c09..cf3df4b0aa 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -33,10 +33,10 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[to raise Exception("need to impl") def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight): - mem_manager = infer_state.mem_manager - mem_manager.operator.copy_kv_to_mem_manager( - layer_index=self.layer_num_, - mem_index=infer_state.mem_index, + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state.write_full_att_kv( + layer_id=self.layer_num_, + token_indices=infer_state.mem_index, kv=cache_kv, ) return diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 28d44fefbf..8ffffb657e 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -226,7 +226,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.mem_manager.free_all() + model.att_resource_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -289,7 +289,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_prefill(prefill_batches[0], prefill_batches[1]) - model.mem_manager.free_all() + model.att_resource_manager.full_att_allocator.free_all() model.req_manager.free_all() del prefill_batches diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py index be5eb6d8a9..23be8798e5 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py @@ -9,6 +9,7 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( KV, Dest_loc, Out, + Out_v, scale, stride_kv_bs, stride_kv_h, @@ -22,6 +23,7 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( FP8_MIN: tl.constexpr, FP8_MAX: tl.constexpr, IS_PER_TENSOR_QUANT: tl.constexpr, + SPLIT_KV: tl.constexpr, ): cur_index = tl.program_id(0) offs_h = tl.arange(0, BLOCK_HEAD) @@ -30,7 +32,13 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) kv_ptrs = KV + cur_index * stride_kv_bs + stride_kv_h * offs_h[:, None] + stride_kv_d * offs_d[None, :] - o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] + if SPLIT_KV: + # 两块独立分配的 K/V 具有相同形状和 stride，输入仍按 K heads、V heads 排列。 + k_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] + v_ptrs = Out_v + dest_index * stride_o_bs + stride_o_h * (offs_h[:, None] - head_num // 2) + stride_o_d * offs_d[None, :] + o_ptrs = tl.where(offs_h[:, None] < head_num // 2, k_ptrs, v_ptrs) + else: + o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] # to fp8 if not IS_PER_TENSOR_QUANT: @@ -50,15 +58,20 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( @torch.no_grad() -def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): +def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False, Out_v=None): """ - 当 is_per_tensor_quant 为 False 时，为 为 per_head 量化。 + 当 is_per_tensor_quant 为 False 时，为 per_head 量化。 + Out_v 非空时，Out/Out_v 分别为 K/V 存储；否则沿用合并输出布局。 """ seq_len = DestLoc.shape[0] head_num = KV.shape[1] assert head_num % 2 == 0 head_dim = KV.shape[2] - assert KV.shape[1] == Out.shape[1] and KV.shape[2] == Out.shape[2] + if Out_v is None: + assert KV.shape[1] == Out.shape[1] and KV.shape[2] == Out.shape[2] + else: + assert KV.shape[1] == 2 * Out.shape[1] and KV.shape[2] == Out.shape[2] + assert Out_v.shape == Out.shape and Out_v.stride() == Out.stride() BLOCK_HEAD = triton.next_power_of_2(head_num) grid = (seq_len,) num_warps = 1 @@ -67,6 +80,7 @@ def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): KV, DestLoc, Out, + Out if Out_v is None else Out_v, scales, KV.stride(0), KV.stride(1), @@ -80,6 +94,7 @@ def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): FP8_MIN=torch.finfo(torch.float8_e4m3fn).min, FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, IS_PER_TENSOR_QUANT=is_per_tensor_quant, + SPLIT_KV=Out_v is not None, num_warps=num_warps, num_stages=1, ) diff --git a/lightllm/models/deepseek3_2/triton_kernel/destindex_copy_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py similarity index 82% rename from lightllm/models/deepseek3_2/triton_kernel/destindex_copy_kv_flashmla_fp8.py rename to lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py index d5859c6867..eca30ec090 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/destindex_copy_kv_flashmla_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py @@ -4,77 +4,6 @@ import triton.language as tl -@triton.jit -def _quant_scale(max_nope, fp8_max): - return tl.exp2(tl.ceil(tl.log2(tl.maximum(max_nope / fp8_max, 1e-4)))) - - -@triton.jit -def _fwd_kernel_destindex_copy_kv_flashmla_fp8( - KV_nope, - KV_rope, - Dest_loc, - O_nope, - O_scale, - O_rope, - stride_kv_nope_bs, - stride_kv_nope_h, - stride_kv_nope_d, - stride_kv_rope_bs, - stride_kv_rope_h, - stride_kv_rope_d, - stride_o_nope_bs, - stride_o_nope_h, - stride_o_nope_d, - stride_o_scale_bs, - stride_o_scale_h, - stride_o_scale_d, - stride_o_rope_bs, - stride_o_rope_h, - stride_o_rope_d, - FP8_MIN: tl.constexpr, - FP8_MAX: tl.constexpr, - BLOCK_DMODEL_NOPE: tl.constexpr, - BLOCK_DMODEL_ROPE: tl.constexpr, - GROUP_SIZE: tl.constexpr, -): - cur_index = tl.program_id(0) - dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) - - offs_rope = tl.arange(0, BLOCK_DMODEL_ROPE) - - # This kernel is only used by the DeepSeek-V3.2 DSA FP8 path, which - # stores a single MQA-style KV head per token. Keep all accesses 1-D so - # Triton treats per-tile scales as scalars instead of 1-element blocks. - kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_rope - - kv_rope = tl.load(kv_rope_ptrs) - - o_rope_ptrs = O_rope + dest_index * stride_o_rope_bs + stride_o_rope_d * offs_rope - tl.store(o_rope_ptrs, kv_rope) - - num_tiles = BLOCK_DMODEL_NOPE // GROUP_SIZE - for tile_idx in range(0, num_tiles): - offs_tile = tile_idx * GROUP_SIZE + tl.arange(0, GROUP_SIZE) - kv_nope_tile_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_tile - kv_nope_tile = tl.load(kv_nope_tile_ptrs) - max_nope = tl.max(tl.abs(kv_nope_tile), axis=0) - kv_scale = _quant_scale(max_nope, FP8_MAX) - kv_nope_fp8 = tl.clamp(kv_nope_tile / kv_scale, min=FP8_MIN, max=FP8_MAX).to(tl.float8e4nv) - - o_nope_ptrs = ( - O_nope - + dest_index * stride_o_nope_bs - + (tile_idx * GROUP_SIZE) * stride_o_nope_d - + tl.arange(0, GROUP_SIZE) * stride_o_nope_d - ) - tl.store(o_nope_ptrs, kv_nope_fp8) - - o_scale_ptrs = O_scale + dest_index * stride_o_scale_bs + tile_idx * stride_o_scale_d - tl.store(o_scale_ptrs, kv_scale.to(tl.float32)) - return - - @torch.no_grad() def destindex_copy_kv_flashmla_fp8( KV_nope: torch.Tensor, @@ -125,3 +54,74 @@ def destindex_copy_kv_flashmla_fp8( num_stages=1, ) return + + +@triton.jit +def _quant_scale(max_nope: tl.tensor, fp8_max: float): + return tl.exp2(tl.ceil(tl.log2(tl.maximum(max_nope / fp8_max, 1e-4)))) + + +@triton.jit +def _fwd_kernel_destindex_copy_kv_flashmla_fp8( + KV_nope: tl.tensor, + KV_rope: tl.tensor, + Dest_loc: tl.tensor, + O_nope: tl.tensor, + O_scale: tl.tensor, + O_rope: tl.tensor, + stride_kv_nope_bs: tl.tensor, + stride_kv_nope_h: tl.tensor, + stride_kv_nope_d: tl.tensor, + stride_kv_rope_bs: tl.tensor, + stride_kv_rope_h: tl.tensor, + stride_kv_rope_d: tl.tensor, + stride_o_nope_bs: tl.tensor, + stride_o_nope_h: tl.tensor, + stride_o_nope_d: tl.tensor, + stride_o_scale_bs: tl.tensor, + stride_o_scale_h: tl.tensor, + stride_o_scale_d: tl.tensor, + stride_o_rope_bs: tl.tensor, + stride_o_rope_h: tl.tensor, + stride_o_rope_d: tl.tensor, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + BLOCK_DMODEL_NOPE: tl.constexpr, + BLOCK_DMODEL_ROPE: tl.constexpr, + GROUP_SIZE: tl.constexpr, +): + cur_index = tl.program_id(0) + dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) + + offs_rope = tl.arange(0, BLOCK_DMODEL_ROPE) + + # This kernel is only used by the DeepSeek-V3.2 DSA FP8 path, which + # stores a single MQA-style KV head per token. Keep all accesses 1-D so + # Triton treats per-tile scales as scalars instead of 1-element blocks. + kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_rope + + kv_rope = tl.load(kv_rope_ptrs) + + o_rope_ptrs = O_rope + dest_index * stride_o_rope_bs + stride_o_rope_d * offs_rope + tl.store(o_rope_ptrs, kv_rope) + + num_tiles = BLOCK_DMODEL_NOPE // GROUP_SIZE + for tile_idx in range(0, num_tiles): + offs_tile = tile_idx * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + kv_nope_tile_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_tile + kv_nope_tile = tl.load(kv_nope_tile_ptrs) + max_nope = tl.max(tl.abs(kv_nope_tile), axis=0) + kv_scale = _quant_scale(max_nope, FP8_MAX) + kv_nope_fp8 = tl.clamp(kv_nope_tile / kv_scale, min=FP8_MIN, max=FP8_MAX).to(tl.float8e4nv) + + o_nope_ptrs = ( + O_nope + + dest_index * stride_o_nope_bs + + (tile_idx * GROUP_SIZE) * stride_o_nope_d + + tl.arange(0, GROUP_SIZE) * stride_o_nope_d + ) + tl.store(o_nope_ptrs, kv_nope_fp8) + + o_scale_ptrs = O_scale + dest_index * stride_o_scale_bs + tile_idx * stride_o_scale_d + tl.store(o_scale_ptrs, kv_scale.to(tl.float32)) + return diff --git a/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py similarity index 87% rename from lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py rename to lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py index 0d0447fec8..bef84c054b 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py @@ -1,13 +1,86 @@ +"""Compact packed FP8 MLA history and newly computed prefill KV for attention.""" + import torch import triton import triton.language as tl +@torch.no_grad() +def get_prefill_kv_cache_and_remap_indices_triton( + packed_kv: torch.Tensor, + topk_mem_indices: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_cache_kv: torch.Tensor, + prefill_dtype: torch.dtype, +): + squeeze_h_kv = topk_mem_indices.ndim == 2 + if squeeze_h_kv: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + + original_shape = topk_mem_indices.shape + flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) + valid_mask = flat_topk != -1 + valid_topk = flat_topk[valid_mask] + table_size = packed_kv.shape[0] + + prefill_row_table = torch.full((table_size,), -1, dtype=torch.int32, device=packed_kv.device) + _build_prefill_row_table_kernel[(prefill_mem_index.numel(),)]( + prefill_mem_index_ptr=prefill_mem_index.to(torch.int32).contiguous(), + row_table_ptr=prefill_row_table, + prefill_token_num=prefill_mem_index.numel(), + num_warps=4, + ) + + unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) + unique_mem_index = unique_mem_index.to(torch.int32) + unique_count = unique_mem_index.numel() + remapped_flat = torch.full_like(flat_topk, -1) + remapped_flat[valid_mask] = inverse.to(torch.int32) + + compact_kv = torch.empty((unique_count, 1, 576), dtype=prefill_dtype, device=packed_kv.device) + packed_nope = packed_kv[:, :, :512].view(torch.float8_e4m3fn).view(-1, 512) + packed_scale = packed_kv[:, :, 512:528].view(torch.float32).view(-1, 4) + packed_rope = packed_kv[:, :, 528:].view(torch.bfloat16).view(-1, 64) + prefill_kv_2d = prefill_cache_kv.view(-1, 576) + compact_kv_2d = compact_kv.view(-1, 576) + + _fill_compact_kv_kernel[(unique_count, 5)]( + packed_nope_ptr=packed_nope, + packed_scale_ptr=packed_scale, + packed_rope_ptr=packed_rope, + unique_mem_index_ptr=unique_mem_index, + prefill_row_table_ptr=prefill_row_table, + prefill_kv_ptr=prefill_kv_2d, + compact_kv_ptr=compact_kv_2d, + packed_nope_stride_s=packed_nope.stride(0), + packed_nope_stride_d=packed_nope.stride(1), + packed_scale_stride_s=packed_scale.stride(0), + packed_scale_stride_d=packed_scale.stride(1), + packed_rope_stride_s=packed_rope.stride(0), + packed_rope_stride_d=packed_rope.stride(1), + prefill_kv_stride_s=prefill_kv_2d.stride(0), + prefill_kv_stride_d=prefill_kv_2d.stride(1), + compact_kv_stride_s=compact_kv_2d.stride(0), + compact_kv_stride_d=compact_kv_2d.stride(1), + unique_num=unique_count, + KV_NOPE_DIM=512, + KV_ROPE_DIM=64, + GROUP_SIZE=128, + BLOCK_D=128, + num_warps=4, + ) + + remapped = remapped_flat.view(original_shape) + if squeeze_h_kv: + remapped = remapped.squeeze(1) + return compact_kv, remapped + + @triton.jit def _build_prefill_row_table_kernel( - prefill_mem_index_ptr, - row_table_ptr, - prefill_token_num, + prefill_mem_index_ptr: tl.tensor, + row_table_ptr: tl.tensor, + prefill_token_num: tl.tensor, ): pid = tl.program_id(0) if pid < prefill_token_num: @@ -17,24 +90,24 @@ def _build_prefill_row_table_kernel( @triton.jit def _fill_compact_kv_kernel( - packed_nope_ptr, - packed_scale_ptr, - packed_rope_ptr, - unique_mem_index_ptr, - prefill_row_table_ptr, - prefill_kv_ptr, - compact_kv_ptr, - packed_nope_stride_s, - packed_nope_stride_d, - packed_scale_stride_s, - packed_scale_stride_d, - packed_rope_stride_s, - packed_rope_stride_d, - prefill_kv_stride_s, - prefill_kv_stride_d, - compact_kv_stride_s, - compact_kv_stride_d, - unique_num, + packed_nope_ptr: tl.tensor, + packed_scale_ptr: tl.tensor, + packed_rope_ptr: tl.tensor, + unique_mem_index_ptr: tl.tensor, + prefill_row_table_ptr: tl.tensor, + prefill_kv_ptr: tl.tensor, + compact_kv_ptr: tl.tensor, + packed_nope_stride_s: tl.tensor, + packed_nope_stride_d: tl.tensor, + packed_scale_stride_s: tl.tensor, + packed_scale_stride_d: tl.tensor, + packed_rope_stride_s: tl.tensor, + packed_rope_stride_d: tl.tensor, + prefill_kv_stride_s: tl.tensor, + prefill_kv_stride_d: tl.tensor, + compact_kv_stride_s: tl.tensor, + compact_kv_stride_d: tl.tensor, + unique_num: tl.tensor, KV_NOPE_DIM: tl.constexpr, KV_ROPE_DIM: tl.constexpr, GROUP_SIZE: tl.constexpr, @@ -102,74 +175,3 @@ def _fill_compact_kv_kernel( value, mask=mask, ) - - -@torch.no_grad() -def get_prefill_kv_cache_and_remap_indices_triton( - packed_kv: torch.Tensor, - topk_mem_indices: torch.Tensor, - prefill_mem_index: torch.Tensor, - prefill_cache_kv: torch.Tensor, - prefill_dtype: torch.dtype, -): - squeeze_h_kv = topk_mem_indices.ndim == 2 - if squeeze_h_kv: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - - original_shape = topk_mem_indices.shape - flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) - valid_mask = flat_topk != -1 - valid_topk = flat_topk[valid_mask] - table_size = packed_kv.shape[0] - - prefill_row_table = torch.full((table_size,), -1, dtype=torch.int32, device=packed_kv.device) - _build_prefill_row_table_kernel[(prefill_mem_index.numel(),)]( - prefill_mem_index_ptr=prefill_mem_index.to(torch.int32).contiguous(), - row_table_ptr=prefill_row_table, - prefill_token_num=prefill_mem_index.numel(), - num_warps=4, - ) - - unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) - unique_mem_index = unique_mem_index.to(torch.int32) - unique_count = unique_mem_index.numel() - remapped_flat = torch.full_like(flat_topk, -1) - remapped_flat[valid_mask] = inverse.to(torch.int32) - - compact_kv = torch.empty((unique_count, 1, 576), dtype=prefill_dtype, device=packed_kv.device) - packed_nope = packed_kv[:, :, :512].view(torch.float8_e4m3fn).view(-1, 512) - packed_scale = packed_kv[:, :, 512:528].view(torch.float32).view(-1, 4) - packed_rope = packed_kv[:, :, 528:].view(torch.bfloat16).view(-1, 64) - prefill_kv_2d = prefill_cache_kv.view(-1, 576) - compact_kv_2d = compact_kv.view(-1, 576) - - _fill_compact_kv_kernel[(unique_count, 5)]( - packed_nope_ptr=packed_nope, - packed_scale_ptr=packed_scale, - packed_rope_ptr=packed_rope, - unique_mem_index_ptr=unique_mem_index, - prefill_row_table_ptr=prefill_row_table, - prefill_kv_ptr=prefill_kv_2d, - compact_kv_ptr=compact_kv_2d, - packed_nope_stride_s=packed_nope.stride(0), - packed_nope_stride_d=packed_nope.stride(1), - packed_scale_stride_s=packed_scale.stride(0), - packed_scale_stride_d=packed_scale.stride(1), - packed_rope_stride_s=packed_rope.stride(0), - packed_rope_stride_d=packed_rope.stride(1), - prefill_kv_stride_s=prefill_kv_2d.stride(0), - prefill_kv_stride_d=prefill_kv_2d.stride(1), - compact_kv_stride_s=compact_kv_2d.stride(0), - compact_kv_stride_d=compact_kv_2d.stride(1), - unique_num=unique_count, - KV_NOPE_DIM=512, - KV_ROPE_DIM=64, - GROUP_SIZE=128, - BLOCK_D=128, - num_warps=4, - ) - - remapped = remapped_flat.view(original_shape) - if squeeze_h_kv: - remapped = remapped.squeeze(1) - return compact_kv, remapped diff --git a/lightllm/common/basemodel/triton_kernel/kv_move.py b/lightllm/common/basemodel/triton_kernel/kv_move.py index d2ce0764b7..ce99a007bd 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_move.py +++ b/lightllm/common/basemodel/triton_kernel/kv_move.py @@ -8,6 +8,8 @@ def _copy_kv_buffer_to_kv_buffer( mem_num, src_mem_index, dst_mem_index, + src_index_stride: tl.constexpr, + dst_index_stride: tl.constexpr, kv_buffer, kv_buffer_stride_l, kv_buffer_stride_s, @@ -24,8 +26,8 @@ def _copy_kv_buffer_to_kv_buffer( kv_buffer_stride_d = tl.cast(kv_buffer_stride_d, dtype=tl.int64) for i in range(start_index, mem_num, grid_num): - src_mem = tl.load(src_mem_index + i) - dst_mem = tl.load(dst_mem_index + i) + src_mem = tl.load(src_mem_index + i * src_index_stride) + dst_mem = tl.load(dst_mem_index + i * dst_index_stride) for j in range(tl.cdiv(kv_buffer_tail_dim, BLOCK)): offs = j * BLOCK + tl.arange(0, BLOCK) mask = offs < kv_buffer_tail_dim @@ -58,6 +60,8 @@ def copy_kv_buffer_to_kv_buffer( mem_num=len(src_mem_index), src_mem_index=src_mem_index, dst_mem_index=dst_mem_index, + src_index_stride=src_mem_index.stride(0), + dst_index_stride=dst_mem_index.stride(0), kv_buffer=kv_buffer, kv_buffer_stride_l=kv_buffer.stride(0), kv_buffer_stride_s=kv_buffer.stride(1), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py deleted file mode 100644 index b0c9a60649..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py +++ /dev/null @@ -1,531 +0,0 @@ -import torch -import triton -import triton.language as tl -from lightllm.common.state_cache_manager import LinearAttCacheConfig - - -@triton.jit -def _copy_kv_buffer_to_cpu_cache( - page_num, - mem_indexes_ptr, # [move_token_num] - page_indexes_ptr, # [page_num], - page_readies_ptr, # [page_num], - big_page_buffer_ids, # [page_num] - cpu_cache_full_att, # [all_page_num, head, xdim] - cpu_cache_full_att_stride_p, - cpu_cache_full_att_stride_h, - cpu_cache_full_att_stride_d, - cpu_cache_conv, # [all_page_num, tp_world_size, xxdim] - cpu_cache_conv_stride_p, - cpu_cache_conv_stride_t, - cpu_cache_conv_stride_d, - cpu_cache_ssm, # [all_page_num, tp_world_size, xxxdim] - cpu_cache_ssm_stride_p, - cpu_cache_ssm_stride_t, - cpu_cache_ssm_stride_d, - gpu_kv_full_att_state, # [token_size, full_att_layer_num, xdim] - gpu_kv_full_att_stride_s, - gpu_kv_full_att_stride_l, - gpu_kv_full_att_stride_d, - cpu_kv_conv_state, # [buffer_count, xxxxxdim] - cpu_kv_conv_stride_s, - cpu_kv_conv_stride_d, - cpu_kv_ssm_state, # [buffer_count, xxxxxxxdim] - cpu_kv_ssm_stride_s, - cpu_kv_ssm_stride_d, - gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim, - tp_rank, - full_att_layer_num, - big_page_token_num, - head_scale_size, - BLOCK: tl.constexpr, -): - split_index_start = tl.program_id(0) - grid_num = tl.num_programs(0) - # 将 所有stride 切成 tl.int64 - cpu_cache_full_att_stride_p = tl.cast(cpu_cache_full_att_stride_p, tl.int64) - cpu_cache_full_att_stride_h = tl.cast(cpu_cache_full_att_stride_h, tl.int64) - cpu_cache_full_att_stride_d = tl.cast(cpu_cache_full_att_stride_d, tl.int64) - cpu_cache_conv_stride_p = tl.cast(cpu_cache_conv_stride_p, tl.int64) - cpu_cache_conv_stride_t = tl.cast(cpu_cache_conv_stride_t, tl.int64) - cpu_cache_conv_stride_d = tl.cast(cpu_cache_conv_stride_d, tl.int64) - cpu_cache_ssm_stride_p = tl.cast(cpu_cache_ssm_stride_p, tl.int64) - cpu_cache_ssm_stride_t = tl.cast(cpu_cache_ssm_stride_t, tl.int64) - cpu_cache_ssm_stride_d = tl.cast(cpu_cache_ssm_stride_d, tl.int64) - gpu_kv_full_att_stride_s = tl.cast(gpu_kv_full_att_stride_s, tl.int64) - gpu_kv_full_att_stride_l = tl.cast(gpu_kv_full_att_stride_l, tl.int64) - gpu_kv_full_att_stride_d = tl.cast(gpu_kv_full_att_stride_d, tl.int64) - cpu_kv_conv_stride_s = tl.cast(cpu_kv_conv_stride_s, tl.int64) - cpu_kv_conv_stride_d = tl.cast(cpu_kv_conv_stride_d, tl.int64) - cpu_kv_ssm_stride_s = tl.cast(cpu_kv_ssm_stride_s, tl.int64) - cpu_kv_ssm_stride_d = tl.cast(cpu_kv_ssm_stride_d, tl.int64) - - for block_index in range(page_num): - cpu_page_index = tl.load(page_indexes_ptr + block_index).to(tl.int64) - run_flag = 1 - if cpu_page_index == -1: - run_flag = 0 - ready_state = tl.load(page_readies_ptr + block_index) - if ready_state: - run_flag = 0 - if tp_rank % head_scale_size == 0: - head_flag = 1 - else: - head_flag = 0 - - mem_start_ptr = mem_indexes_ptr + big_page_token_num * block_index - for i in range(split_index_start, tl.cdiv(gpu_full_att_tail_dim, BLOCK) * run_flag * head_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < gpu_full_att_tail_dim - per_token_size = gpu_full_att_tail_dim // big_page_token_num - per_layer_size = per_token_size // full_att_layer_num - mem_offs = gpu_start_i // (per_token_size) - mem_index = tl.load(mem_start_ptr + mem_offs, mask=mask, other=-1) - layer_index = (gpu_start_i // (per_layer_size)) % full_att_layer_num - dim_index = gpu_start_i % per_layer_size - gpu_full_att_data = tl.load( - gpu_kv_full_att_state - + mem_index * gpu_kv_full_att_stride_s - + layer_index * gpu_kv_full_att_stride_l - + dim_index * gpu_kv_full_att_stride_d, - mask=mask & (mem_index != -1), - other=0, - ) - dest_cpu_cache_full_att_ptr = ( - cpu_cache_full_att - + cpu_page_index * cpu_cache_full_att_stride_p - + (tp_rank // head_scale_size) * cpu_cache_full_att_stride_h - + gpu_start_i - ) - tl.store(dest_cpu_cache_full_att_ptr, gpu_full_att_data, mask=mask & (mem_index != -1)) - - big_page_idx = tl.load(big_page_buffer_ids + block_index) - - for i in range(split_index_start, tl.cdiv(cpu_kv_conv_tail_dim, BLOCK) * run_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_conv_tail_dim - cpu_kv_conv_data = tl.load( - cpu_kv_conv_state + big_page_idx * cpu_kv_conv_stride_s + gpu_start_i, - mask=mask, - other=0, - ) - dest_cpu_cache_conv_ptr = ( - cpu_cache_conv - + cpu_page_index * cpu_cache_conv_stride_p - + tp_rank * cpu_cache_conv_stride_t - + gpu_start_i - ) - tl.store(dest_cpu_cache_conv_ptr, cpu_kv_conv_data, mask=mask) - - for i in range(split_index_start, tl.cdiv(cpu_kv_ssm_tail_dim, BLOCK) * run_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_ssm_tail_dim - - cpu_kv_ssm_data = tl.load( - cpu_kv_ssm_state + big_page_idx * cpu_kv_ssm_stride_s + gpu_start_i, - mask=mask, - other=0, - ) - dest_cpu_cache_ssm_ptr = ( - cpu_cache_ssm + cpu_page_index * cpu_cache_ssm_stride_p + tp_rank * cpu_cache_ssm_stride_t + gpu_start_i - ) - tl.store(dest_cpu_cache_ssm_ptr, cpu_kv_ssm_data, mask=mask) - - return - - -def copy_kv_buffer_to_cpu_cache( - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - big_page_buffer_ids: torch.Tensor, - gpu_kv_full_att_state: torch.Tensor, # [full_att_layer_num, s, head_num, head_dim] - cpu_kv_conv_state: torch.Tensor, # [s, linear_layer_num, dim] - cpu_kv_ssm_state: torch.Tensor, # [s, linear_layer_num, xdim] - cpu_cache_tensor: torch.Tensor, # [page_num, 1, 1, 1, xxdim] - tp_rank: int, - tp_world_size: int, - big_page_token_num: int, - linear_config: LinearAttCacheConfig, - grid_num: int = 12, -): - assert len(page_indexes) == len(page_readies) == len(big_page_buffer_ids) - assert len(mem_indexes) % len(page_indexes) == 0 - - BLOCK = 4096 - if linear_config.full_att_all_num_kv_heads % tp_world_size == 0: - # tp world size 不比 kv 的 head 多时 - head_scale_size = 1 - else: - head_scale_size = tp_world_size // linear_config.full_att_all_num_kv_heads - - cpu_page_num = cpu_cache_tensor.shape[0] - cpu_cache_tensor = cpu_cache_tensor.view(cpu_page_num, -1).view(dtype=torch.uint8) - a = linear_config.get_cpu_cache_full_att_bytes() - b = linear_config.get_cpu_cache_conv_bytes() - c = linear_config.get_cpu_cache_ssm_bytes() - - if head_scale_size == 1: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, tp_world_size, -1) - else: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, linear_config.full_att_all_num_kv_heads, -1) - - cpu_cache_full_att = cpu_cache_full_att.view(dtype=torch.uint64) - - cpu_cache_conv = cpu_cache_tensor[:, a : (a + b)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - cpu_cache_ssm = ( - cpu_cache_tensor[:, (a + b) : (a + b + c)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - ) - - gpu_kv_full_att_state = gpu_kv_full_att_state.view( - gpu_kv_full_att_state.shape[0], gpu_kv_full_att_state.shape[1], -1 - ).view(dtype=torch.uint64) - - gpu_kv_full_att_state = gpu_kv_full_att_state.permute(1, 0, 2) # [s, layer_num, xxdim] - - cpu_kv_conv_state = cpu_kv_conv_state.view(cpu_kv_conv_state.shape[0], -1).view(dtype=torch.uint64) - cpu_kv_ssm_state = cpu_kv_ssm_state.view(cpu_kv_ssm_state.shape[0], -1).view(dtype=torch.uint64) - - gpu_full_att_tail_dim = gpu_kv_full_att_state.shape[-1] * gpu_kv_full_att_state.shape[-2] * big_page_token_num - cpu_kv_conv_tail_dim = cpu_kv_conv_state.shape[-1] - cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1] - full_att_layer_num = gpu_kv_full_att_state.shape[-2] - - assert full_att_layer_num == linear_config.get_full_att_kv_layer_num_with_draft_model() - assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1] - assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1] - assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1] - assert gpu_kv_full_att_state.stride(2) == 1 - assert ( - gpu_full_att_tail_dim % big_page_token_num == 0 - and (gpu_full_att_tail_dim // big_page_token_num) % full_att_layer_num == 0 - ) - assert (tp_rank // head_scale_size) < linear_config.full_att_all_num_kv_heads - - grid = (grid_num,) - _copy_kv_buffer_to_cpu_cache[grid]( - page_num=len(page_indexes), - mem_indexes_ptr=mem_indexes, - page_indexes_ptr=page_indexes, - page_readies_ptr=page_readies, - big_page_buffer_ids=big_page_buffer_ids, - cpu_cache_full_att=cpu_cache_full_att, - cpu_cache_full_att_stride_p=cpu_cache_full_att.stride(0), - cpu_cache_full_att_stride_h=cpu_cache_full_att.stride(1), - cpu_cache_full_att_stride_d=cpu_cache_full_att.stride(2), - cpu_cache_conv=cpu_cache_conv, - cpu_cache_conv_stride_p=cpu_cache_conv.stride(0), - cpu_cache_conv_stride_t=cpu_cache_conv.stride(1), - cpu_cache_conv_stride_d=cpu_cache_conv.stride(2), - cpu_cache_ssm=cpu_cache_ssm, - cpu_cache_ssm_stride_p=cpu_cache_ssm.stride(0), - cpu_cache_ssm_stride_t=cpu_cache_ssm.stride(1), - cpu_cache_ssm_stride_d=cpu_cache_ssm.stride(2), - gpu_kv_full_att_state=gpu_kv_full_att_state, - gpu_kv_full_att_stride_s=gpu_kv_full_att_state.stride(0), - gpu_kv_full_att_stride_l=gpu_kv_full_att_state.stride(1), - gpu_kv_full_att_stride_d=gpu_kv_full_att_state.stride(2), - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_conv_stride_s=cpu_kv_conv_state.stride(0), - cpu_kv_conv_stride_d=cpu_kv_conv_state.stride(1), - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_kv_ssm_stride_s=cpu_kv_ssm_state.stride(0), - cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(1), - gpu_full_att_tail_dim=gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim=cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim=cpu_kv_ssm_tail_dim, - tp_rank=tp_rank, - full_att_layer_num=full_att_layer_num, - big_page_token_num=big_page_token_num, - head_scale_size=head_scale_size, - BLOCK=BLOCK, - ) - - -@triton.jit -def _copy_cpu_cache_to_kv_buffer( - page_num, - mem_indexes_ptr, # [move_token_num] - page_indexes_ptr, # [page_num], - big_page_buffer_ids, # [page_num] - cpu_cache_full_att, # [all_page_num, head, xdim] - cpu_cache_full_att_stride_p, - cpu_cache_full_att_stride_h, - cpu_cache_full_att_stride_d, - cpu_cache_conv, # [all_page_num, tp_world_size, xxdim] - cpu_cache_conv_stride_p, - cpu_cache_conv_stride_t, - cpu_cache_conv_stride_d, - cpu_cache_ssm, # [all_page_num, tp_world_size, xxxdim] - cpu_cache_ssm_stride_p, - cpu_cache_ssm_stride_t, - cpu_cache_ssm_stride_d, - gpu_kv_full_att_state, # [token_size, full_att_layer_num, xdim] - gpu_kv_full_att_stride_s, - gpu_kv_full_att_stride_l, - gpu_kv_full_att_stride_d, - cpu_kv_conv_state, # [buffer_count, xxxxxdim] - cpu_kv_conv_stride_s, - cpu_kv_conv_stride_d, - cpu_kv_ssm_state, # [buffer_count, xxxxxxxdim] - cpu_kv_ssm_stride_s, - cpu_kv_ssm_stride_d, - gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim, - tp_rank, - full_att_layer_num, - big_page_token_num, - head_scale_size, - BLOCK: tl.constexpr, -): - split_index_start = tl.program_id(0) - grid_num = tl.num_programs(0) - # 将 所有stride 切成 tl.int64 - cpu_cache_full_att_stride_p = tl.cast(cpu_cache_full_att_stride_p, tl.int64) - cpu_cache_full_att_stride_h = tl.cast(cpu_cache_full_att_stride_h, tl.int64) - cpu_cache_full_att_stride_d = tl.cast(cpu_cache_full_att_stride_d, tl.int64) - cpu_cache_conv_stride_p = tl.cast(cpu_cache_conv_stride_p, tl.int64) - cpu_cache_conv_stride_t = tl.cast(cpu_cache_conv_stride_t, tl.int64) - cpu_cache_conv_stride_d = tl.cast(cpu_cache_conv_stride_d, tl.int64) - cpu_cache_ssm_stride_p = tl.cast(cpu_cache_ssm_stride_p, tl.int64) - cpu_cache_ssm_stride_t = tl.cast(cpu_cache_ssm_stride_t, tl.int64) - cpu_cache_ssm_stride_d = tl.cast(cpu_cache_ssm_stride_d, tl.int64) - gpu_kv_full_att_stride_s = tl.cast(gpu_kv_full_att_stride_s, tl.int64) - gpu_kv_full_att_stride_l = tl.cast(gpu_kv_full_att_stride_l, tl.int64) - gpu_kv_full_att_stride_d = tl.cast(gpu_kv_full_att_stride_d, tl.int64) - cpu_kv_conv_stride_s = tl.cast(cpu_kv_conv_stride_s, tl.int64) - cpu_kv_conv_stride_d = tl.cast(cpu_kv_conv_stride_d, tl.int64) - cpu_kv_ssm_stride_s = tl.cast(cpu_kv_ssm_stride_s, tl.int64) - cpu_kv_ssm_stride_d = tl.cast(cpu_kv_ssm_stride_d, tl.int64) - - for block_index in range(page_num): - cpu_page_index = tl.load(page_indexes_ptr + block_index).to(tl.int64) - - mem_start_ptr = mem_indexes_ptr + big_page_token_num * block_index - for i in range(split_index_start, tl.cdiv(gpu_full_att_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < gpu_full_att_tail_dim - per_token_size = gpu_full_att_tail_dim // big_page_token_num - per_layer_size = per_token_size // full_att_layer_num - mem_offs = gpu_start_i // (per_token_size) - mem_index = tl.load(mem_start_ptr + mem_offs, mask=mask, other=-1) - layer_index = (gpu_start_i // (per_layer_size)) % full_att_layer_num - dim_index = gpu_start_i % per_layer_size - - src_cpu_cache_full_att_ptr = ( - cpu_cache_full_att - + cpu_page_index * cpu_cache_full_att_stride_p - + (tp_rank // head_scale_size) * cpu_cache_full_att_stride_h - + gpu_start_i - ) - cpu_full_att_data = tl.load(src_cpu_cache_full_att_ptr, mask=mask & (mem_index != -1), other=0) - - tl.store( - gpu_kv_full_att_state - + mem_index * gpu_kv_full_att_stride_s - + layer_index * gpu_kv_full_att_stride_l - + dim_index * gpu_kv_full_att_stride_d, - cpu_full_att_data, - mask=mask & (mem_index != -1), - ) - - big_page_idx = tl.load(big_page_buffer_ids + block_index) - - for i in range(split_index_start, tl.cdiv(cpu_kv_conv_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_conv_tail_dim - - src_cpu_cache_conv_ptr = ( - cpu_cache_conv - + cpu_page_index * cpu_cache_conv_stride_p - + tp_rank * cpu_cache_conv_stride_t - + gpu_start_i - ) - cpu_kv_conv_data = tl.load(src_cpu_cache_conv_ptr, mask=mask, other=0) - - tl.store( - cpu_kv_conv_state + big_page_idx * cpu_kv_conv_stride_s + gpu_start_i, - cpu_kv_conv_data, - mask=mask, - ) - - for i in range(split_index_start, tl.cdiv(cpu_kv_ssm_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_ssm_tail_dim - - src_cpu_cache_ssm_ptr = ( - cpu_cache_ssm + cpu_page_index * cpu_cache_ssm_stride_p + tp_rank * cpu_cache_ssm_stride_t + gpu_start_i - ) - cpu_kv_ssm_data = tl.load(src_cpu_cache_ssm_ptr, mask=mask, other=0) - - tl.store( - cpu_kv_ssm_state + big_page_idx * cpu_kv_ssm_stride_s + gpu_start_i, - cpu_kv_ssm_data, - mask=mask, - ) - - return - - -def copy_cpu_cache_to_kv_buffer( - mem_indexes: torch.Tensor, - big_page_buffer_ids: torch.Tensor, - page_indexes: torch.Tensor, - gpu_full_att_kv_state: torch.Tensor, # [layer_num, s, head_num, head_dim] - cpu_kv_conv_state: torch.Tensor, # [layer_num, s, dim] - cpu_kv_ssm_state: torch.Tensor, # [layer_num, s, xdim] - cpu_cache_tensor: torch.Tensor, # [page_num, 1, 1, tp_world_size, xxdim] - tp_rank: int, - tp_world_size: int, - big_page_token_num: int, - linear_config: LinearAttCacheConfig, - grid_num: int = 12, -): - assert len(mem_indexes) % len(page_indexes) == 0 - - BLOCK = 4096 - if linear_config.full_att_all_num_kv_heads % tp_world_size == 0: - head_scale_size = 1 - else: - head_scale_size = tp_world_size // linear_config.full_att_all_num_kv_heads - - cpu_page_num = cpu_cache_tensor.shape[0] - cpu_cache_tensor = cpu_cache_tensor.view(cpu_page_num, -1).view(dtype=torch.uint8) - a = linear_config.get_cpu_cache_full_att_bytes() - b = linear_config.get_cpu_cache_conv_bytes() - c = linear_config.get_cpu_cache_ssm_bytes() - - if head_scale_size == 1: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, tp_world_size, -1) - else: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, linear_config.full_att_all_num_kv_heads, -1) - - cpu_cache_full_att = cpu_cache_full_att.view(dtype=torch.uint64) - - cpu_cache_conv = cpu_cache_tensor[:, a : (a + b)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - cpu_cache_ssm = ( - cpu_cache_tensor[:, (a + b) : (a + b + c)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - ) - - gpu_full_att_kv_state = gpu_full_att_kv_state.view( - gpu_full_att_kv_state.shape[0], gpu_full_att_kv_state.shape[1], -1 - ).view(dtype=torch.uint64) - gpu_full_att_kv_state = gpu_full_att_kv_state.permute(1, 0, 2) # [s, layer_num, xxdim] - - cpu_kv_conv_state = cpu_kv_conv_state.view(cpu_kv_conv_state.shape[0], -1).view(dtype=torch.uint64) - cpu_kv_ssm_state = cpu_kv_ssm_state.view(cpu_kv_ssm_state.shape[0], -1).view(dtype=torch.uint64) - - gpu_full_att_tail_dim = gpu_full_att_kv_state.shape[-1] * gpu_full_att_kv_state.shape[-2] * big_page_token_num - cpu_kv_conv_tail_dim = cpu_kv_conv_state.shape[-1] - cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1] - full_att_layer_num = gpu_full_att_kv_state.shape[-2] - - assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1] - assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1] - assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1] - assert gpu_full_att_kv_state.stride(2) == 1 - - assert (tp_rank // head_scale_size) < linear_config.full_att_all_num_kv_heads - - grid = (grid_num,) - _copy_cpu_cache_to_kv_buffer[grid]( - page_num=len(page_indexes), - mem_indexes_ptr=mem_indexes, - page_indexes_ptr=page_indexes, - big_page_buffer_ids=big_page_buffer_ids, - cpu_cache_full_att=cpu_cache_full_att, - cpu_cache_full_att_stride_p=cpu_cache_full_att.stride(0), - cpu_cache_full_att_stride_h=cpu_cache_full_att.stride(1), - cpu_cache_full_att_stride_d=cpu_cache_full_att.stride(2), - cpu_cache_conv=cpu_cache_conv, - cpu_cache_conv_stride_p=cpu_cache_conv.stride(0), - cpu_cache_conv_stride_t=cpu_cache_conv.stride(1), - cpu_cache_conv_stride_d=cpu_cache_conv.stride(2), - cpu_cache_ssm=cpu_cache_ssm, - cpu_cache_ssm_stride_p=cpu_cache_ssm.stride(0), - cpu_cache_ssm_stride_t=cpu_cache_ssm.stride(1), - cpu_cache_ssm_stride_d=cpu_cache_ssm.stride(2), - gpu_kv_full_att_state=gpu_full_att_kv_state, - gpu_kv_full_att_stride_s=gpu_full_att_kv_state.stride(0), - gpu_kv_full_att_stride_l=gpu_full_att_kv_state.stride(1), - gpu_kv_full_att_stride_d=gpu_full_att_kv_state.stride(2), - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_conv_stride_s=cpu_kv_conv_state.stride(0), - cpu_kv_conv_stride_d=cpu_kv_conv_state.stride(1), - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_kv_ssm_stride_s=cpu_kv_ssm_state.stride(0), - cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(1), - gpu_full_att_tail_dim=gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim=cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim=cpu_kv_ssm_tail_dim, - tp_rank=tp_rank, - full_att_layer_num=full_att_layer_num, - big_page_token_num=big_page_token_num, - head_scale_size=head_scale_size, - BLOCK=BLOCK, - ) - - -@triton.jit -def _copy_linear_att_state_to_linear_att_state( - src_conv_state, - dst_conv_state, - src_ssm_state, - dst_ssm_state, - conv_size, - ssm_size, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - grid_num = tl.num_programs(0) - - # copy conv state - num_conv_blocks = tl.cdiv(conv_size, BLOCK) - for i in range(pid, num_conv_blocks, grid_num): - start = i * BLOCK + tl.arange(0, BLOCK) - mask = start < conv_size - data = tl.load(src_conv_state + start, mask=mask, other=0) - tl.store(dst_conv_state + start, data, mask=mask) - - # copy ssm state - num_ssm_blocks = tl.cdiv(ssm_size, BLOCK) - for i in range(pid, num_ssm_blocks, grid_num): - start = i * BLOCK + tl.arange(0, BLOCK) - mask = start < ssm_size - data = tl.load(src_ssm_state + start, mask=mask, other=0) - tl.store(dst_ssm_state + start, data, mask=mask) - - -def copy_linear_att_state_to_linear_att_state( - src_conv_state: torch.Tensor, - src_ssm_state: torch.Tensor, - dst_conv_state: torch.Tensor, - dst_ssm_state: torch.Tensor, - grid_num: int = 16, -): - assert src_conv_state.shape == dst_conv_state.shape - assert src_ssm_state.shape == dst_ssm_state.shape - - BLOCK = 4096 - - src_conv_flat = src_conv_state.view(-1).view(dtype=torch.uint8) - dst_conv_flat = dst_conv_state.view(-1).view(dtype=torch.uint8) - src_ssm_flat = src_ssm_state.view(-1).view(dtype=torch.uint8) - dst_ssm_flat = dst_ssm_state.view(-1).view(dtype=torch.uint8) - - conv_size = src_conv_flat.shape[0] - ssm_size = src_ssm_flat.shape[0] - - grid = (grid_num,) - _copy_linear_att_state_to_linear_att_state[grid]( - src_conv_state=src_conv_flat, - dst_conv_state=dst_conv_flat, - src_ssm_state=src_ssm_flat, - dst_ssm_state=dst_ssm_flat, - conv_size=conv_size, - ssm_size=ssm_size, - BLOCK=BLOCK, - ) diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py deleted file mode 100644 index 3cd6a211d7..0000000000 --- a/lightllm/common/kv_cache_mem_manager/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -from .allocator import KvCacheAllocator -from .mem_manager import MemoryManager, ReadOnlyStaticsMemoryManager -from .ppl_int8kv_mem_manager import PPLINT8KVMemoryManager -from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from .deepseek2_mem_manager import Deepseek2MemoryManager -from .deepseek3_2mem_manager import Deepseek3_2MemoryManager -from .deepseek4_mem_manager import DeepseekV4MemoryManager -from .fp8_per_token_group_quant_deepseek3_2mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager -from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager -from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager -from .qwen3next_mem_manager import Qwen3NextMemManager -from .glm5_next_mem_manager import Glm5NextMemManager -from .fp8_glm5_next_mem_manager import FP8Glm5NextMemManager - -__all__ = [ - "KvCacheAllocator", - "MemoryManager", - "ReadOnlyStaticsMemoryManager", - "PPLINT4KVMemoryManager", - "PPLINT8KVMemoryManager", - "Deepseek2MemoryManager", - "Deepseek3_2MemoryManager", - "DeepseekV4MemoryManager", - "FP8PerTokenGroupQuantDeepseek3_2MemoryManager", - "FP8StaticPerHeadQuantMemManager", - "FP8StaticPerTensorQuantMemManager", - "Qwen3NextMemManager", - "Glm5NextMemManager", - "FP8Glm5NextMemManager", -] diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py deleted file mode 100644 index 1331943b1a..0000000000 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ /dev/null @@ -1,103 +0,0 @@ -import torch -from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt -from lightllm.utils.dist_utils import get_current_rank_in_node -from lightllm.utils.log_utils import init_logger -from typing import Union, List - -logger = init_logger(__name__) - - -class KvCacheAllocator: - def __init__(self, size: int) -> None: - self.size = size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - - rank_in_node = get_current_rank_in_node() - # 用共享内存进行共享，router 模块读取进行精确的调度估计；基础层会统一添加服务前缀以防止实例冲突。 - self.shared_can_use_token_num = SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return - - def alloc(self, need_size) -> torch.Tensor: - if need_size > self.mark_end - self.mark_start: - logger.error(f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}") - assert False, "error alloc state" - - start = self.mark_start - end = self.mark_start + need_size - self.mark_start += need_size - - self.can_use_mem_size -= need_size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - # 利用缓冲区返回，避免异步情况下的内存竞争 - if self._return_start + need_size > self._mem_state_return.shape[0]: - self._return_start = 0 - ans = self._mem_state_return[self._return_start : self._return_start + need_size] - ans.copy_(self.mem_state[start:end]) - self._return_start += need_size - return ans - - def free(self, free_index: Union[torch.Tensor, List[int]]): - """_summary_ - - Args: - free_index (torch.Tensor): _description_ - """ - - end = self.mark_start - start = self.mark_start - len(free_index) - assert start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}" - - if isinstance(free_index, list): - self.mem_state.numpy()[start:end] = free_index - else: - # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 - self.mem_state[start:end] = free_index - - self.mark_start -= len(free_index) - - self.can_use_mem_size += len(free_index) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - if self.can_use_mem_size == len(self.mem_state): - logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") - return - - def free_all(self): - self.mem_state.numpy()[:] = list(range(0, len(self.mem_state))) - self.mark_start = 0 - self.mark_end = len(self.mem_state) - self.can_use_mem_size = len(self.mem_state) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return - - def resize(self, new_size: int) -> None: - """ - just for test code - """ - self.size = new_size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) diff --git a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py deleted file mode 100644 index 4be820b753..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -import os -import torch.distributed as dist -from .mem_manager import MemoryManager -from typing import List, Union, Any -from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io -from .operator import Deepseek2MemOperator - - -logger = init_logger(__name__) - - -class Deepseek2MemoryManager(MemoryManager): - - operator_class = Deepseek2MemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - kv = self.kv_buffer[layer_index] - return kv - - def get_cell_size(self): - return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - self.kv_buffer = torch.empty((layer_num, size + self.page_size, head_num, head_dim), dtype=dtype, device="cuda") - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - # DeepSeek MLA's single compressed KV latent is replicated across TP ranks, - # so the transfer buffer only needs one copy and does not multiply head_num by TP size. - return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - mla_page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[0].kv_buffer, - mode="write", - ) - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for mem in dp_mems: - mla_page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=mem.kv_buffer, - mode="read", - ) diff --git a/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py deleted file mode 100644 index 4034bcc8fa..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from typing import Any -from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager -from .operator import Deepseek3_2MemOperator - - -class Deepseek3_2MemoryManager(Deepseek2MemoryManager): - - operator_class = Deepseek3_2MemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - assert dtype in [torch.bfloat16, torch.float16] - # 因为V3.2 使用了NSA 稀疏的缘故，所以其head_dim 会比原始的kv 多 128 + 4 = 132 个字节 (128 fp8 + 4byte float32 scale)， - # 但是为了让整个数组具备16字节对齐，满足一些算子的约束，修改为添加 128 + 16 = 144 个字节, 这 144个字节中，后面132个字节用于 - # 存储真实数据，剩下12个，浪费了，只是占位。 - # 所以在子类中定制为其pad上，对外使用的接口，需要进行重载区别。 - super().__init__(size, dtype, head_num, head_dim + (144 // 2), layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - kv = self.kv_buffer[layer_index][:, :, : (self.head_dim - (144 // 2))] - return kv - - def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: - return self.kv_buffer[layer_index].view(dtype=torch.uint8)[:, :, -132:] diff --git a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py deleted file mode 100644 index b72587545f..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py +++ /dev/null @@ -1,66 +0,0 @@ -import torch -from typing import Any - -from .deepseek2_mem_manager import Deepseek2MemoryManager -from .operator import FP8PerTokenGroupQuantDeepseek3_2MemOperator - - -class FP8PerTokenGroupQuantDeepseek3_2MemoryManager(Deepseek2MemoryManager): - - operator_class = FP8PerTokenGroupQuantDeepseek3_2MemOperator - - kv_nope_dim = 512 - kv_rope_dim = 64 - # 576 = 512 + 64 - kv_head_dim = kv_nope_dim + kv_rope_dim - - quant_group_size = 128 - # 4 = 512 / 128 - quant_group_num = kv_nope_dim // quant_group_size - # 4 * 4 = quant_group_num * fp32 - # 64 * 2 = kv_rope_dim * bfloat16 - # 656 bytes = 512 + (4 * 4) + (64 * 2) - flashmla_bytes_per_token = kv_nope_dim + quant_group_num * 4 + kv_rope_dim * 2 - - indexer_head_dim = 128 - # 128 + 4 = indexer_head_dim + fp32 - # 132 bytes = 128 + 4 - indexer_bytes_per_token = indexer_head_dim + 4 - - # 16-byte 对齐，满足FlashMLA的对齐要求 - alignment = 16 - total_bytes_per_token = ( - (flashmla_bytes_per_token + indexer_bytes_per_token + alignment - 1) // alignment * alignment - ) - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - assert head_num == 1, "DeepSeek-V3.2 DSA FP8 path expects MQA-style head_num == 1" - self.prefill_dtype = dtype - super().__init__(size, torch.uint8, head_num, self.total_bytes_per_token, layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - return self.kv_buffer[layer_index][:, :, : self.flashmla_bytes_per_token] - - def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: - begin = self.flashmla_bytes_per_token - end = begin + self.indexer_bytes_per_token - return self.kv_buffer[layer_index][:, :, begin:end] - - def get_prefill_kv_cache_and_remap_indices( - self, - packed_kv: torch.Tensor, - topk_indices: torch.Tensor, - prefill_mem_index: torch.Tensor, - prefill_cache_kv: torch.Tensor, - ): - from lightllm.models.deepseek3_2.triton_kernel.prefill_compact_kv_flashmla_fp8 import ( - get_prefill_kv_cache_and_remap_indices_triton, - ) - - return get_prefill_kv_cache_and_remap_indices_triton( - packed_kv=packed_kv, - topk_mem_indices=topk_indices, - prefill_mem_index=prefill_mem_index, - prefill_cache_kv=prefill_cache_kv, - prefill_dtype=self.prefill_dtype, - ) diff --git a/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py deleted file mode 100755 index 2d5e8a1e5e..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py +++ /dev/null @@ -1,98 +0,0 @@ -import os -import json -import torch -import torch.distributed as dist -from typing import Tuple, Any -from lightllm.utils.config_utils import get_model_architectures -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp -from .mem_manager import MemoryManager -from .operator import FP8StaticPerHeadQuantMemOperator - -logger = init_logger(__name__) - - -class FP8StaticPerHeadQuantMemManager(MemoryManager): - operator_class = FP8StaticPerHeadQuantMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - # 这里用uint8存储量化后的kv，方便兼容各种torch算子。fp8量化目前采用离线方案，kv_buffer不存储scale - super().__init__(size, torch.uint8, head_num, head_dim, layer_num, always_copy, mem_fraction) - - self.qmax = torch.finfo(torch.float8_e4m3fn).max - self.qmin = torch.finfo(torch.float8_e4m3fn).min - self.scales = None - - cfg, all_scales = self._load_and_check_config() - q_cfg = cfg["q_calibration"] - - all_head_num = cfg["num_head"] - # A joint target+draft config uses its leading target scale rows on a target-only server. - all_scales = all_scales[: self.layer_num].to(device="cuda") - - factor = (get_dp_world_size() * head_num) // all_head_num - all_scales = torch.repeat_interleave(input=all_scales, repeats=factor, dim=-1) - rank_in_dp = get_current_rank_in_dp() - - v_offset = all_scales.shape[1] // 2 - start_head = rank_in_dp * head_num - end_head = start_head + head_num - k_scales = all_scales[:, start_head:end_head].contiguous() - v_scales = all_scales[:, v_offset + start_head : v_offset + end_head].contiguous() - self.scales = torch.cat((k_scales, v_scales), dim=-1) - - all_q_scales = torch.tensor(q_cfg["scales"], dtype=torch.float32, device="cuda").view(q_cfg["scales_shape"])[ - : self.layer_num - ] - q_factor = (get_dp_world_size() * head_num) // q_cfg["num_head"] - all_q_scales = torch.repeat_interleave(input=all_q_scales, repeats=q_factor, dim=-1) - self.q_scales = all_q_scales[:, start_head:end_head].contiguous() - return - - def _load_and_check_config(self): - config_path = get_env_start_args().kv_quant_calibration_config_path - if config_path is None: - raise ValueError("fp8kv_sph requires kv_quant_calibration_config_path with q_calibration") - - logger.info(f"kv_quant_calibration_config_path {config_path} is set, will load kv quant calibration config") - if not os.path.exists(config_path): - raise FileNotFoundError(f"kv_quant_calibration_config {config_path} not found") - - with open(config_path, "r") as f: - cfg = json.load(f) - - num_layers, num_heads = cfg["num_layers"], cfg["num_head"] - scales = torch.tensor(cfg["scales"], dtype=torch.float32) - expected_shape = [num_layers, 2 * num_heads] - if cfg["scales_shape"] != expected_shape or list(scales.shape) != expected_shape: - raise ValueError(f"calibration scales and scales_shape must match {expected_shape}") - - runtime_draft_layers = get_added_mtp_kv_layer_num() - runtime_target_layers = self.layer_num - runtime_draft_layers - if "num_target_layers" in cfg or "num_draft_layers" in cfg: - target_layers, draft_layers = cfg["num_target_layers"], cfg["num_draft_layers"] - if ( - draft_layers < 0 - or target_layers + draft_layers != num_layers - or target_layers != runtime_target_layers - or (runtime_draft_layers > 0 and draft_layers != runtime_draft_layers) - ): - raise ValueError( - f"invalid calibration layer layout: target={target_layers}, draft={draft_layers}, " - f"runtime target={runtime_target_layers}, draft={runtime_draft_layers}" - ) - elif num_layers != self.layer_num: - raise ValueError(f"legacy calibration num_layers {num_layers} do not match runtime {self.layer_num}") - - q_cfg = cfg.get("q_calibration") - if not q_cfg: - raise ValueError("fp8kv_sph calibration config requires q_calibration") - assert (get_dp_world_size() * self.head_num) % num_heads == 0 - assert (get_dp_world_size() * self.head_num) % q_cfg["num_head"] == 0 - return cfg, scales - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v diff --git a/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py deleted file mode 100755 index 06338808bb..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py +++ /dev/null @@ -1,70 +0,0 @@ -import os -import json -import torch -import torch.distributed as dist -from typing import Tuple, Any -from lightllm.utils.config_utils import get_model_architectures -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from .mem_manager import MemoryManager -from .operator import FP8StaticPerTensorQuantMemOperator - -logger = init_logger(__name__) - - -class FP8StaticPerTensorQuantMemManager(MemoryManager): - - operator_class = FP8StaticPerTensorQuantMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - # 这里用uint8存储量化后的kv，方便兼容各种torch算子。fp8量化目前采用离线方案，kv_buffer不存储scale - super().__init__(size, torch.uint8, head_num, head_dim, layer_num, always_copy, mem_fraction) - - self.qmax = torch.finfo(torch.float8_e4m3fn).max - self.qmin = torch.finfo(torch.float8_e4m3fn).min - self.scales = None - - if get_env_start_args().kv_quant_calibration_config_path is not None: - logger.info( - f"kv_quant_calibration_config_path {get_env_start_args().kv_quant_calibration_config_path} is set, " - "will load kv quant calibration config" - ) - cfg = self._load_and_check_config() - self.scales = torch.tensor(cfg["scales"], dtype=torch.float32, device="cuda").view(cfg["scales_shape"]) - else: - self.scales = torch.ones((self.kv_buffer.shape[0], 2), dtype=torch.float32, device="cuda") - - self.cpu_scales = self.scales.detach().cpu().numpy() - return - - def _load_and_check_config(self): - if os.path.exists(get_env_start_args().kv_quant_calibration_config_path): - with open(get_env_start_args().kv_quant_calibration_config_path, "r") as f: - cfg = json.load(f) - - if cfg["qmin"] != self.qmin: - raise ValueError(f"qmin {cfg['qmin']} in config not match torch.float8_e4m3fn.min {self.qmin}") - if cfg["qmax"] != self.qmax: - raise ValueError(f"qmax {cfg['qmax']} in config not match torch.float8_e4m3fn.max {self.qmax}") - model_arch = get_model_architectures(get_env_start_args().model_dir) - if cfg["architectures"] != model_arch: - raise ValueError( - f"architectures {cfg['architectures']} in config " f"not match current model_arch {model_arch}" - ) - if cfg["num_layers"] != self.layer_num: - raise ValueError( - f"num_layers {cfg['num_layers']} in config " f"not match current layer_num {self.layer_num}" - ) - assert ( - cfg["quant_type"] == "per_tensor" - ), f"quant type {cfg['quant_type']} in config not match per-tensor backend" - return cfg - else: - raise FileNotFoundError( - f"kv_quant_calibration_config {get_env_start_args().kv_quant_calibration_config_path} not found" - ) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py deleted file mode 100644 index 74f0f26810..0000000000 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ /dev/null @@ -1,122 +0,0 @@ -import torch - -from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( - Qwen3NextMemManager, - Qwen3NextLinearAttPageHelper, -) -from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv -from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io - - -class Glm5NextMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): - output = self.mem_manager.get_att_input_params(layer_index)[:, :, : kv.shape[-1]] - destindex_copy_kv(kv, mem_index, output) - - -class Glm5NextMemManager(Qwen3NextMemManager): - """One packed token buffer; KDA uses the standard big/small page pools.""" - - operator_class = Glm5NextMemOperator - - def __init__(self, *args, mla_head_dim=512, **kwargs): - self.mla_head_dim = mla_head_dim - super().__init__(*args, **kwargs) - - def get_cell_size(self): - return self.head_dim * self.dtype.itemsize * self.layer_num - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert head_num == 1 - self.kv_buffer = torch.empty((layer_num, size + self.page_size, 1, head_dim), dtype=dtype, device="cuda") - self._init_linear_att_buffers() - - def _layer_buffer(self, layer_index): - return self.kv_buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] - - def get_att_input_params(self, layer_index): - return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] - - def get_indexer_k_buffer(self, layer_index): - return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] - - def write_to_shm(self, req_manager): - self.req_to_indexer_tail = req_manager.req_to_indexer_tail - return super().write_to_shm(req_manager) - - def _create_att_state_page_helper(self): - return Glm5NextAttStatePageHelper(self) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - # Packed MLA/index KV is replicated across TP ranks. - return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) - - def write_mem_to_page_kv_move_buffer( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - return super().write_mem_to_page_kv_move_buffer( - mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx - ) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - mla_page_io( - mem_indexes=pin_indexes.cuda(non_blocking=True), - page_tensor=self.kv_move_buffer[page_index], - kv_buffer=mem_managers[dp_index * dp_world_size].kv_buffer, - mode="write", - ) - - def read_page_kv_move_buffer_to_mem( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - return super().read_page_kv_move_buffer_to_mem( - mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx - ) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - indexes = pin_indexes.cuda(non_blocking=True) - for mem in mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size]: - mla_page_io( - mem_indexes=indexes, - page_tensor=self.kv_move_buffer[page_index], - kv_buffer=mem.kv_buffer, - mode="read", - ) - - -class Glm5NextAttStatePageHelper(Qwen3NextLinearAttPageHelper): - """Append the replicated K-pool tail to the global Conv/SSM state page.""" - - def __init__(self, mem_manager): - super().__init__(mem_manager) - self.tail_dtype = mem_manager.req_to_indexer_tail.buffer.dtype - self.tail_shape = ( - self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - *mem_manager.req_to_indexer_tail.buffer.shape[2:], - ) - self.tail_offset = ((self.state_nbytes + 15) // 16) * 16 - self.tail_nbytes = self.tail_shape[0] * self.tail_shape[1] * self.tail_shape[2] * self.tail_dtype.itemsize - self.state_nbytes = self.tail_offset + self.tail_nbytes - - def _view_page_to_tail(self, page_index): - page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - return ( - page_bytes[self.tail_offset : self.tail_offset + self.tail_nbytes] - .view(self.tail_dtype) - .view(self.tail_shape) - ) - - def write_req_to_page(self, page_index, req_idx, dp_mems): - super().write_req_to_page(page_index, req_idx, dp_mems) - # The ring uses absolute token positions; sequence length already travels - # with the request, so restoring it needs no additional PD metadata. - self._view_page_to_tail(page_index).copy_(dp_mems[0].req_to_indexer_tail.buffer[:, req_idx], non_blocking=True) - - def read_page_to_req(self, page_index, req_idx, dp_mems): - super().read_page_to_req(page_index, req_idx, dp_mems) - tail = self._view_page_to_tail(page_index) - for mem in dp_mems: - mem.req_to_indexer_tail.buffer[:, req_idx].copy_(tail, non_blocking=True) diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py deleted file mode 100755 index 7108f39d79..0000000000 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ /dev/null @@ -1,308 +0,0 @@ -import math -import os -import re -import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from typing import List, Tuple, Any, Union -from lightllm.utils.log_utils import init_logger -from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt -from .allocator import KvCacheAllocator -from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory -from lightllm.utils.dist_utils import ( - get_current_device_id, - get_current_rank_in_node, - get_node_world_size, -) -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_utils import get_service_shm_name -from lightllm.utils.config_utils import get_num_key_value_heads -from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io -from lightllm.utils.device_utils import kv_trans_use_p2p -from lightllm.utils.shm_utils import create_or_link_shm -from multiprocessing.reduction import ForkingPickler -from filelock import FileLock -from .operator import BaseMemManagerOperator, NormalMemOperator - -logger = init_logger(__name__) - - -class MemoryManager: - - operator_class = NormalMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - self.size = size - self.head_num = head_num - self.head_dim = head_dim - self.layer_num = layer_num - self.always_copy = always_copy - self.dtype = dtype - # profile the max total token num if the size is None - self.profile_size(mem_fraction) - - # A physical KV page must never straddle the allocator boundary. The - # unused remainder is intentionally dropped so every managed page has - # a stable ``mem_index // page_size`` id. - self.page_size = get_env_start_args().page_size - self.size = self.size // self.page_size * self.page_size - - self.allocator = KvCacheAllocator(self.size) - - self._init_buffers( - self.size, - dtype, - head_num, - head_dim, - layer_num, - ) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - - # 构建对外的操作类接口 - self.operator: BaseMemManagerOperator = self.operator_class(self) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - num_kv_head = get_num_key_value_heads(get_env_start_args().model_dir) - return ( - page_num, - page_size, - self.layer_num, - 2 * num_kv_head, - self.head_dim, - ) - - def get_pd_kv_move_buffer_size(self): - args = get_env_start_args() - if args.run_mode not in ["prefill", "decode"]: - return 0 - shape = self.get_paged_kv_move_buffer_shape( - page_num=args.pd_kv_page_num, - page_size=args.pd_kv_page_size, - ) - return math.prod(shape) * torch._utils._element_size(self.dtype) - - def profile_size(self, mem_fraction): - if self.size is not None: - return - - torch.cuda.empty_cache() - world_size = dist.get_world_size() - available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) - cell_size = self.get_cell_size() - pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size - self.size = int(available_memory_bytes / cell_size) - if world_size > 1: - tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") - dist.all_reduce(tensor, op=dist.ReduceOp.MIN) - self.size = tensor.item() - logger.info( - f"{str(available_memory)} GB space is available after load the model weight\n" - f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" - f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" - f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" - ) - return - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - # 在初始化 kv_buffer 时，每层额外保留一整页；这部分空间不会被 allocator 对外分配。 - # 这一页用于多 DP、overlap microbatch 等模式下的 padding 请求， - # 页内索引存储在 HOLD_TOKEN_MEMINDEXES 中，与 req_manager 的 HOLD_REQUEST_ID 作用类似。 - self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=dtype, device="cuda" - ) - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - self.kv_move_buffer = torch.empty( - self.get_paged_kv_move_buffer_shape(page_num, page_size), dtype=self.dtype, device="cuda" - ) - self._buffer_mem_indexes_tensors = [ - torch.empty((page_size,), dtype=torch.int64, device="cpu", pin_memory=True) for _ in range(page_num) - ] - return self.kv_move_buffer - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - repeat_count = dp_world_size * self.kv_buffer.shape[2] // self.kv_move_buffer.shape[3] - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for tp_index in range(dp_world_size): - if tp_index % repeat_count == 0: - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="write", - ) - # keep for debug - # logger.info(f"src token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") - # logger.info(f"src page token tensor {cur_page[0, :, 0, 0]}") - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - mem_indexes_gpu = torch.tensor(mem_indexes, dtype=torch.int64, device="cpu", pin_memory=True).cuda( - non_blocking=True - ) - for tp_index in range(dp_world_size): - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="read", - ) - # keep for debug - # logger.info(f"dst token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") - # logger.info(f"dst page token tensor {cur_page[0, :, 0, 0]}") - - def _free_buffers(self): - self.kv_buffer = None - - def alloc(self, need_size) -> torch.Tensor: - assert ( - need_size % self.page_size == 0 - ), f"KV cache allocation size {need_size} must be a multiple of page_size {self.page_size}" - return self.allocator.alloc(need_size) - - def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: - self.allocator.free(free_index) - - def free_all(self): - self.allocator.free_all() - - def resize_mem(self, new_size): - """ - just for test code - """ - size = new_size // self.page_size * self.page_size - dtype = self.dtype - head_num = self.head_num - head_dim = self.head_dim - layer_num = self.layer_num - - self.size = size - self.allocator.resize(size) - self._free_buffers() - self._init_buffers(size, dtype, head_num, head_dim, layer_num) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - return - - def get_index_kv_buffer(self, index): - return {"kv_buffer": self.kv_buffer[:, index]} - - def load_index_kv_buffer(self, index, load_tensor_dict): - self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - - def write_to_shm(self, req_manager): - """ - 将 mem manager 写入到 shm中，方便pd分离等特性直接从中读取，不依赖进程间队列。 - """ - if kv_trans_use_p2p(): - from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor - - mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ - - from lightllm.common.req_manager import ReqManager - - req_manager: ReqManager = req_manager - - # 这个地方是一个不太优雅的设计，但是暂时这么做，可以让dp shared kv swap模块直接访问 req_manager 中的 req_to_token_indexs - # 避免过多无用的数据复制和传输开销。 - self.req_to_token_indexs: torch.Tensor = req_manager.req_to_token_indexs - - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - with lock: - node_world_size = get_node_world_size() - shm_name = f"mem_manager_{get_current_rank_in_node()}" - obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] - # Each consumer needs its own IPC ref-counter; its pickle size can vary. - obj_size = max(len(obj_bytes) for obj_bytes in obj_bytes_array) - shm = create_or_link_shm( - name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create" - ) - logger.info(f"create shm {shm.name} size {shm.size} for mem manger shared buffer") - shm.buf[0:4] = (node_world_size * 2).to_bytes(4, "little") - shm.buf[4:8] = obj_size.to_bytes(4, "little") - start_index = 8 - for obj_bytes in obj_bytes_array: - shm.buf[start_index : start_index + obj_size] = obj_bytes.ljust(obj_size, b"\0") - start_index += obj_size - - @staticmethod - def loads_from_shm(rank_in_node: int) -> "MemoryManager": - shm_name = f"mem_manager_{rank_in_node}" - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - logger.info(f"get memmanager from shm {shm_name}") - with lock: - shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link") - left_num = int.from_bytes(shm.buf[0:4], "little") - obj_size = int.from_bytes(shm.buf[4:8], "little") - assert left_num > 0 - end_index = 8 + left_num * obj_size - start_index = 8 + (left_num - 1) * obj_size - obj_bytes = shm.buf[start_index:end_index].tobytes() - shm.buf[0:4] = (left_num - 1).to_bytes(4, byteorder="little") - shm.close() - return ForkingPickler.loads(obj_bytes) - - -class ReadOnlyStaticsMemoryManager: - """ - 读取一些统计信息 - """ - - def __init__(self) -> None: - args = get_env_start_args() - self.global_world_size = args.tp - self.node_world_size = args.tp // args.nnodes - self.dp_world_size = self.global_world_size // args.dp - # 兼容多机 dp size=1 纯 tp 模式的情况 - self.is_multinode_tp = args.dp == 1 and args.nnodes > 1 - self.shared_tp_infos = [ - SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - for rank_in_node in range(0, self.node_world_size, self.dp_world_size) - ] - - def get_unrefed_token_num(self, dp_rank_in_node: int): - if self.is_multinode_tp: - return self.shared_tp_infos[0].get_value() - return self.shared_tp_infos[dp_rank_in_node].get_value() diff --git a/lightllm/common/kv_cache_mem_manager/mem_utils.py b/lightllm/common/kv_cache_mem_manager/mem_utils.py deleted file mode 100644 index 398f407e36..0000000000 --- a/lightllm/common/kv_cache_mem_manager/mem_utils.py +++ /dev/null @@ -1,76 +0,0 @@ -from . import ( - MemoryManager, - PPLINT8KVMemoryManager, - PPLINT4KVMemoryManager, - Deepseek2MemoryManager, - Deepseek3_2MemoryManager, - FP8StaticPerHeadQuantMemManager, - FP8StaticPerTensorQuantMemManager, -) -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.llm_utils import get_llm_model_class -from functools import lru_cache - -logger = init_logger(__name__) - - -@lru_cache(maxsize=None) -def select_mem_manager_class(): - # case 1 - # 先判断是否是 deepseek 系列的模型 - model_class = get_llm_model_class() - from lightllm.models import DeepseekV4TpPartModel - - if issubclass(model_class, DeepseekV4TpPartModel): - from . import DeepseekV4MemoryManager - - mem_class = DeepseekV4MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_class - - from lightllm.models import Deepseek3_2TpPartModel - - if issubclass(model_class, Deepseek3_2TpPartModel): - if get_env_start_args().llm_kv_type == "fp8kv_dsa": - from . import FP8PerTokenGroupQuantDeepseek3_2MemoryManager - - mem_class = FP8PerTokenGroupQuantDeepseek3_2MemoryManager - else: - mem_class = Deepseek3_2MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_class - - from lightllm.models import Deepseek2TpPartModel - - if issubclass(model_class, Deepseek2TpPartModel): - mem_class = Deepseek2MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_class - - # case normal - logger.info(f"mode setting params: {get_env_start_args().llm_kv_type}") - if get_env_start_args().llm_kv_type == "int8kv": - memory_manager_class = PPLINT8KVMemoryManager - elif get_env_start_args().llm_kv_type == "int4kv": - memory_manager_class = PPLINT4KVMemoryManager - elif get_env_start_args().llm_kv_type == "fp8kv_sph": - memory_manager_class = FP8StaticPerHeadQuantMemManager - elif get_env_start_args().llm_kv_type == "fp8kv_spt": - memory_manager_class = FP8StaticPerTensorQuantMemManager - elif get_env_start_args().llm_kv_type == "None": - memory_manager_class = MemoryManager - - logger.info(f"Model kv cache using mem_manager class: {memory_manager_class}") - return memory_manager_class - - -@lru_cache(maxsize=None) -def used_mem_manager_has_scale() -> bool: - mem_class = select_mem_manager_class() - return mem_class in [ - PPLINT8KVMemoryManager, - PPLINT4KVMemoryManager, - FP8StaticPerHeadQuantMemManager, - FP8StaticPerTensorQuantMemManager, - ] diff --git a/lightllm/common/kv_cache_mem_manager/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py deleted file mode 100644 index 442c2e300e..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -from .base import BaseMemManagerOperator -from .normal import NormalMemOperator -from .quant import QuantScaleMemOperator, PPLInt4KVMemOperator, PPLInt8KVMemOperator -from .linear_att import LinearAttMemOperator -from .deepseek import ( - Deepseek2MemOperator, - Deepseek3_2MemOperator, - DeepseekV4MemOperator, - FP8PerTokenGroupQuantDeepseek3_2MemOperator, -) -from .fp8_quant import ( - FP8StaticPerHeadQuantMemOperator, - FP8StaticPerTensorQuantMemOperator, -) diff --git a/lightllm/common/kv_cache_mem_manager/operator/base.py b/lightllm/common/kv_cache_mem_manager/operator/base.py deleted file mode 100644 index 682b8a5d60..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/base.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from abc import ABC, abstractmethod -from typing import List, TYPE_CHECKING - -if TYPE_CHECKING: - from ..mem_manager import MemoryManager - from lightllm.server.router.model_infer.infer_batch import InferReq - -# 定义一个抽象基类 -class BaseMemManagerOperator(ABC): - def __init__(self, mem_manager: "MemoryManager") -> None: - super().__init__() - self.mem_manager = mem_manager - - @abstractmethod - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - pass - - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - raise NotImplementedError() - - # cpu cache 的相关操作接口 - def load_cpu_cache_to_gpu( - self, mem_indexes: torch.Tensor, page_indexes: torch.Tensor, cpu_cache_client, req: "InferReq" - ): - raise NotImplementedError() - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client, - req: "InferReq", - ): - raise NotImplementedError() - - # dp 间共享 kv 的操作接口, 提升dp 模式下的kv 共享效率，降低调度的难度 - def copy_kv_from_other_dp_ranks( - self, - mem_managers: List["MemoryManager"], - move_token_indexes: torch.Tensor, - token_dp_indexes: torch.Tensor, - mem_indexes: torch.Tensor, - dp_size_in_node: int, - rank_in_dp: int, - ): - raise NotImplementedError() diff --git a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py deleted file mode 100644 index 6be23b5cfe..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py +++ /dev/null @@ -1,207 +0,0 @@ -from typing import TYPE_CHECKING - -import torch -from .normal import NormalMemOperator -from .base import BaseMemManagerOperator -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args - -if TYPE_CHECKING: - from ..deepseek4_mem_manager import DeepseekV4MemoryManager - from lightllm.common.req_manager.deepseek4 import DeepseekV4ReqManager - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - -logger = init_logger(__name__) - - -class Deepseek2MemOperator(NormalMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager - - mem_manager: Deepseek2MemoryManager = self.mem_manager - - from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank + rope_dim == mem_manager.kv_buffer.shape[-1] - - destindex_copy_kv( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - mem_manager.kv_buffer[layer_index][:, :, :kv_lora_rank], - mem_manager.kv_buffer[layer_index][:, :, kv_lora_rank:], - ) - return - - -class Deepseek3_2MemOperator(Deepseek2MemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.deepseek3_2mem_manager import Deepseek3_2MemoryManager - - mem_manager: Deepseek3_2MemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank + rope_dim == mem_manager.kv_buffer.shape[-1] - (144 // 2) - - destindex_copy_kv( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - mem_manager.kv_buffer[layer_index][:, :, :kv_lora_rank], - mem_manager.kv_buffer[layer_index][:, :, kv_lora_rank : (kv_lora_rank + rope_dim)], - ) - return - - -class FP8PerTokenGroupQuantDeepseek3_2MemOperator(BaseMemManagerOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( - FP8PerTokenGroupQuantDeepseek3_2MemoryManager, - ) - - mem_manager: FP8PerTokenGroupQuantDeepseek3_2MemoryManager = self.mem_manager - from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_kv_flashmla_fp8 import ( - destindex_copy_kv_flashmla_fp8, - ) - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank == 512, f"Expected kv_lora_rank=512, got {kv_lora_rank}" - - flashmla_bytes_per_token = mem_manager.flashmla_bytes_per_token - - o_nope = mem_manager.kv_buffer[layer_index][:, :, :512].view(torch.float8_e4m3fn) - o_scale = mem_manager.kv_buffer[layer_index][:, :, 512:528].view(torch.float32) - o_rope = mem_manager.kv_buffer[layer_index][:, :, 528:flashmla_bytes_per_token].view(torch.bfloat16) - destindex_copy_kv_flashmla_fp8( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - o_nope, - o_scale, - o_rope, - ) - return - - -class DeepseekV4MemOperator(BaseMemManagerOperator): - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - """仅复制压缩 KV 和 indexer K；SWA、compressor 状态由 ReqManager 恢复。""" - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - src = src_mem_index.to(device="cuda", dtype=torch.int64, non_blocking=True) - dst = dst_mem_index.to(device="cuda", dtype=torch.int64, non_blocking=True) - for buffer in self.mem_manager.kv_buffer: - copy_kv_buffer_to_kv_buffer(src, dst, buffer.view(1, buffer.shape[0], -1)) - - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - # 模型推理层通过融合算子写入 SWA，compressor/indexer 直接写入压缩 KV。 - raise NotImplementedError("DeepSeek-V4 writes packed KV using request-owned SWA slots") - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ) -> None: - from lightllm.utils.dist_utils import get_current_rank_in_dp - from lightllm.server.router.model_infer.infer_batch import g_infer_context - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_kv_buffer_to_cpu_cache, copy_state_cache - - # MLA KV and compression states are replicated across TP ranks. - if get_current_rank_in_dp() != 0: - return - - mem_manager: "DeepseekV4MemoryManager" = self.mem_manager - req_manager: "DeepseekV4ReqManager" = g_infer_context.req_manager - page_token_num = get_env_start_args().cpu_cache_token_page_size - - big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - max_kv_len = (len(mem_indexes) // page_token_num) * page_token_num - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * page_token_num - for seq_len in range(start_kv_len, max_kv_len + 1, page_token_num): - big_page_buffer_ids_cpu.append(req.hybrid_len_to_big_page_id[seq_len]) - - # 碎页：将尾部小页状态复制到临时大页，KV 由拷贝算子按实际 token 数处理。 - has_tail_page = len(mem_indexes) % page_token_num != 0 - if has_tail_page: - assert req.tail_small_page_buffer_id is not None - temp_page_id = mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID - copy_state_cache( - src_state=req_manager.small_page_buffers.buffer[req.tail_small_page_buffer_id], - dst_state=mem_manager.big_page_buffers.buffer[temp_page_id], - ) - big_page_buffer_ids_cpu.append(temp_page_id) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - copy_kv_buffer_to_cpu_cache( - mem_manager=mem_manager, - mem_indexes=mem_indexes, - page_indexes=page_indexes, - page_readies=page_readies, - big_page_buffer_ids=big_page_buffer_ids_gpu, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - ) - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ) -> None: - from lightllm.server.router.model_infer.infer_batch import g_infer_context - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_cpu_cache_to_kv_buffer, copy_state_cache - - mem_manager: "DeepseekV4MemoryManager" = self.mem_manager - req_manager: "DeepseekV4ReqManager" = g_infer_context.req_manager - page_token_num = get_env_start_args().cpu_cache_token_page_size - - big_page_num = len(mem_indexes) // page_token_num - max_kv_len = (req.cur_kv_len // page_token_num) * page_token_num - big_page_buffer_ids_cpu = [] - for _ in range(big_page_num): - page_id = mem_manager.big_page_buffers.alloc_one_state_cache() - assert page_id is not None - req.hybrid_len_to_big_page_id[max_kv_len] = page_id - big_page_buffer_ids_cpu.append(page_id) - max_kv_len -= page_token_num - big_page_buffer_ids_cpu.reverse() - - # 碎页：先将状态读入临时大页，再恢复请求运行态。 - has_tail_page = len(mem_indexes) % page_token_num != 0 - if has_tail_page: - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - copy_cpu_cache_to_kv_buffer( - mem_manager=mem_manager, - mem_indexes=mem_indexes, - big_page_buffer_ids=big_page_buffer_ids_gpu, - page_indexes=page_indexes, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - ) - - if has_tail_page and req.cur_kv_len == req.hybrid_cache_len: - g_infer_context.radix_cache.free_one_small_page_buffer() - small_page_buffer_id = req_manager.small_page_buffers.alloc_one_state_cache() - if small_page_buffer_id is not None: - req.tail_small_page_buffer_id = small_page_buffer_id - # 临时大页会被复用，先将碎页状态保留到小页池。 - copy_state_cache( - src_state=mem_manager.big_page_buffers.buffer[big_page_buffer_ids_cpu[-1]], - dst_state=req_manager.small_page_buffers.buffer[small_page_buffer_id], - ) - - req_manager.restore_big_page_state(big_page_buffer_idx=big_page_buffer_ids_cpu[-1], req=req) diff --git a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py b/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py deleted file mode 100644 index 4b1f0ac207..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -logger = init_logger(__name__) - - -class FP8StaticPerHeadQuantMemOperator(BaseMemManagerOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.fp8_static_per_head_quant_mem_manager import ( - FP8StaticPerHeadQuantMemManager, - ) - - mem_manager: FP8StaticPerHeadQuantMemManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import ( - destindex_copy_kv_fp8, - ) - - scales = mem_manager.scales - destindex_copy_kv_fp8( - kv, - mem_index, - scales[layer_index], - mem_manager.kv_buffer[layer_index].view(torch.float8_e4m3fn), - ) - return - - -class FP8StaticPerTensorQuantMemOperator(BaseMemManagerOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import ( - destindex_copy_kv_fp8, - ) - - destindex_copy_kv_fp8( - kv, - mem_index, - mem_manager.scales[layer_index], - mem_manager.kv_buffer[layer_index].view(torch.float8_e4m3fn), - is_per_tensor_quant=True, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py deleted file mode 100644 index e935539e37..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ /dev/null @@ -1,210 +0,0 @@ -import torch -import triton -from typing import List -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - -logger = init_logger(__name__) - - -class LinearAttMemOperator(BaseMemManagerOperator): - """ - 只用于非量化的linear att 混合 full att的模型，列入 qwen3.5 - """ - - def __init__(self, mem_manager): - super().__init__(mem_manager) - self.linear_config = mem_manager.linear_config - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size - max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu = [] - for i in range(big_page_num): - page_id = mem_manager.big_page_buffers.alloc_one_state_cache() - assert page_id is not None - req.hybrid_len_to_big_page_id[max_kv_len] = page_id - big_page_buffer_ids_cpu.append(page_id) - max_kv_len -= args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu.reverse() - - # 碎页情况的处理 - has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 - if has_tail_page: - padded_token_num = triton.cdiv( - len(mem_indexes), args.cpu_cache_token_page_size - ) * args.cpu_cache_token_page_size - len(mem_indexes) - mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) - - # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - - assert len(big_page_buffer_ids_gpu) == len(page_indexes) - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_cpu_cache_to_kv_buffer, - ) - - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes, - big_page_buffer_ids=big_page_buffer_ids_gpu, - page_indexes=page_indexes, - gpu_full_att_kv_state=mem_manager.kv_buffer, - cpu_kv_conv_state=mem_manager.big_page_buffers.conv_state_cache.buffer, - cpu_kv_ssm_state=mem_manager.big_page_buffers.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - tp_rank=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - big_page_token_num=args.cpu_cache_token_page_size, - linear_config=self.linear_config, - ) - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=big_page_buffer_ids_cpu[-1], - req=req, - ) - - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - args = get_env_start_args() - if not hasattr(self, "big_page_ids_buffer_store"): - self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") - # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， - # 再从大页上拷贝到对应的运行态页面上 - self.mem_indexes_buffer = torch.empty( - (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" - ) - - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size - for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): - page_id = req.hybrid_len_to_big_page_id[seq_len] - big_page_buffer_ids_cpu.append(page_id) - - if len(mem_indexes) % args.cpu_cache_token_page_size != 0: - # 存在不满大页的碎页的页面存在需要复制的情况 - dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert dst_len <= self.mem_indexes_buffer.shape[0] - dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) - dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) - mem_indexes = dst_mem_indexes - assert req.tail_small_page_buffer_id is not None - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_linear_att_state_to_linear_att_state, - ) - - src_conv_state, src_ssm_state = g_infer_context.radix_cache.small_page_buffers.get_state_cache( - buffer_idx=req.tail_small_page_buffer_id - ) - dst_conv_state, dst_ssm_state = mem_manager.big_page_buffers.get_state_cache( - buffer_idx=mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, - ) - copy_linear_att_state_to_linear_att_state( - src_conv_state=src_conv_state, - src_ssm_state=src_ssm_state, - dst_conv_state=dst_conv_state, - dst_ssm_state=dst_ssm_state, - ) - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) - - assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) - - big_page_buffer_ids_cpu = torch.tensor( - big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True - ) - assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] - big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] - big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_kv_buffer_to_cpu_cache, - ) - - copy_kv_buffer_to_cpu_cache( - mem_indexes=mem_indexes, - page_indexes=page_indexes, - page_readies=page_readies, - big_page_buffer_ids=big_page_buffer_ids_gpu, - gpu_kv_full_att_state=mem_manager.kv_buffer, - cpu_kv_conv_state=mem_manager.big_page_buffers.conv_state_cache.buffer, - cpu_kv_ssm_state=mem_manager.big_page_buffers.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - tp_rank=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - big_page_token_num=args.cpu_cache_token_page_size, - linear_config=self.linear_config, - ) - return - - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - # Qwen3Next 需要调整 layer_index - layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import ( - destindex_copy_kv, - ) - - destindex_copy_kv(kv, mem_index, mem_manager.kv_buffer[layer_index]) - return - - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - copy_kv_buffer_to_kv_buffer( - src_mem_index.cuda(non_blocking=True), dst_mem_index.cuda(non_blocking=True), self.mem_manager.kv_buffer - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/normal.py b/lightllm/common/kv_cache_mem_manager/operator/normal.py deleted file mode 100644 index 3c53ace079..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/normal.py +++ /dev/null @@ -1,115 +0,0 @@ -import torch -from typing import TYPE_CHECKING, List -from .base import BaseMemManagerOperator -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - -logger = init_logger(__name__) - - -class NormalMemOperator(BaseMemManagerOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import ( - destindex_copy_kv, - ) - - destindex_copy_kv(kv, mem_index, mem_manager.kv_buffer[layer_index]) - return - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu - - load_cpu_kv_to_gpu( - gpu_mem_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=None, - cpu_kv_cache=cpu_cache_client.cpu_kv_cache_tensor, - cpu_kv_cache_scale=None, - page_indexes=page_indexes, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - assert len(mem_indexes) // args.cpu_cache_token_page_size == len(page_indexes) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu - - offload_gpu_kv_to_cpu( - token_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=None, - cpu_kv_cache=cpu_cache_client.cpu_kv_cache_tensor, - cpu_kv_cache_scale=None, - page_indexes=page_indexes, - page_readies=page_readies, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - def copy_kv_from_other_dp_ranks( - self, - mem_managers: List, - move_token_indexes: torch.Tensor, - token_dp_indexes: torch.Tensor, - mem_indexes: torch.Tensor, - dp_size_in_node: int, - rank_in_dp: int, - ): - if not hasattr(self, "mem_ptrs_tensor"): - # 构建一个2D tensor，shape为(layer_num, mem_num) - mems_ptr_list = [] - for i in range(0, len(mem_managers)): - mems_ptr_list.append(mem_managers[i].kv_buffer.data_ptr()) - self.mem_ptrs_tensor = torch.tensor(mems_ptr_list, dtype=torch.uint64, device="cpu", pin_memory=True) - - from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_for_dp - - # 一次性传输所有层 - kv_trans_for_dp( - input_mems=self.mem_ptrs_tensor.cuda(non_blocking=True), - input_idx=move_token_indexes, - input_dp_idx=token_dp_indexes, - output=self.mem_manager.kv_buffer, - output_idx=mem_indexes, - dp_size_in_node=dp_size_in_node, - rank_in_dp=rank_in_dp, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/quant.py b/lightllm/common/kv_cache_mem_manager/operator/quant.py deleted file mode 100644 index a3a1c1d01c..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/quant.py +++ /dev/null @@ -1,129 +0,0 @@ -import torch -from typing import TYPE_CHECKING, List -from .base import BaseMemManagerOperator -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - -logger = init_logger(__name__) - - -class QuantScaleMemOperator(BaseMemManagerOperator): - """ - 对于kv cache中包含独立的对应每个token的scale变量的memManager使用。 - """ - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - cpu_cache_meta = cpu_cache_client.kv_cache_tensor_meta - cpu_kv_cache = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, 0 : cpu_cache_meta.head_dim] - cpu_kv_cache_scale = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, cpu_cache_meta.head_dim :].view( - mem_manager.scale_buffer.dtype - ) - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu - - load_cpu_kv_to_gpu( - gpu_mem_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=mem_manager.scale_buffer, - cpu_kv_cache=cpu_kv_cache, - cpu_kv_cache_scale=cpu_kv_cache_scale, - page_indexes=page_indexes, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - assert len(mem_indexes) // args.cpu_cache_token_page_size == len(page_indexes) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - cpu_cache_meta = cpu_cache_client.kv_cache_tensor_meta - cpu_kv_cache = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, 0 : cpu_cache_meta.head_dim] - cpu_kv_cache_scale = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, cpu_cache_meta.head_dim :].view( - mem_manager.scale_buffer.dtype - ) - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu - - offload_gpu_kv_to_cpu( - token_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=mem_manager.scale_buffer, - cpu_kv_cache=cpu_kv_cache, - cpu_kv_cache_scale=cpu_kv_cache_scale, - page_indexes=page_indexes, - page_readies=page_readies, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - -class PPLInt4KVMemOperator(QuantScaleMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager - - mem_manager: PPLINT4KVMemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import ( - destindex_copy_int4kv, - ) - - destindex_copy_int4kv( - kv, - mem_index, - mem_manager.kv_buffer[layer_index], - mem_manager.scale_buffer[layer_index], - quant_group_size=mem_manager.group_quant_size, - ) - return - - -class PPLInt8KVMemOperator(QuantScaleMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - mem_manager: PPLINT8KVMemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( - destindex_copy_quantize_kv, - ) - - destindex_copy_quantize_kv( - kv, - mem_index, - mem_manager.kv_buffer[layer_index], - mem_manager.scale_buffer[layer_index], - quant_group_dim=mem_manager.group_quant_size, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py deleted file mode 100755 index 6b02cbed63..0000000000 --- a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py +++ /dev/null @@ -1,50 +0,0 @@ -import torch -from typing import Tuple, Any -from .mem_manager import MemoryManager -from .operator import PPLInt4KVMemOperator - - -class PPLINT4KVMemoryManager(MemoryManager): - - operator_class = PPLInt4KVMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9): - self.kv_dtype = torch.int8 - self.group_quant_size = 8 - super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - k_scale = self.scale_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - v_scale = self.scale_buffer[layer_index][:, self.head_num :, :] - return (k, k_scale), (v, v_scale) - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim // 2 * self.layer_num * torch._utils._element_size( - self.kv_dtype - ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size( - self.dtype - ) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim // 2), dtype=torch.int8, device="cuda" - ) - self.scale_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim // self.group_quant_size), - dtype=dtype, - device="cuda", - ) - - def _free_buffers(self): - self.kv_buffer = None - self.scale_buffer = None - - def get_index_kv_buffer(self, index): - return {"kv_buffer": self.kv_buffer[:, index], "scale_buffer": self.scale_buffer[:, index]} - - def load_index_kv_buffer(self, index, load_tensor_dict): - self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - self.scale_buffer[:, index].copy_(load_tensor_dict["scale_buffer"]) diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py deleted file mode 100755 index 0f5e8e9370..0000000000 --- a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py +++ /dev/null @@ -1,49 +0,0 @@ -import torch -from typing import Tuple, Any -from .mem_manager import MemoryManager -from .operator import PPLInt8KVMemOperator - - -class PPLINT8KVMemoryManager(MemoryManager): - operator_class = PPLInt8KVMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9): - self.kv_dtype = torch.int8 - self.group_quant_size = 8 - super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - k_scale = self.scale_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - v_scale = self.scale_buffer[layer_index][:, self.head_num :, :] - return (k, k_scale), (v, v_scale) - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size( - self.kv_dtype - ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size( - self.dtype - ) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=torch.int8, device="cuda" - ) - self.scale_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim // self.group_quant_size), - dtype=dtype, - device="cuda", - ) - - def _free_buffers(self): - self.kv_buffer = None - self.scale_buffer = None - - def get_index_kv_buffer(self, index): - return {"kv_buffer": self.kv_buffer[:, index], "scale_buffer": self.scale_buffer[:, index]} - - def load_index_kv_buffer(self, index, load_tensor_dict): - self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - self.scale_buffer[:, index].copy_(load_tensor_dict["scale_buffer"]) diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py deleted file mode 100644 index 1962fa71b0..0000000000 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ /dev/null @@ -1,464 +0,0 @@ -import torch -import triton -from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager -from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager -from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager -from .operator import ( - FP8StaticPerHeadQuantMemOperator, - FP8StaticPerTensorQuantMemOperator, - LinearAttMemOperator, -) -from typing import Tuple, Any, List - -logger = init_logger(__name__) - - -class Qwen3NextMemManager(MemoryManager): - operator_class = LinearAttMemOperator - - def __init__( - self, - size, - dtype, - num_kv_heads, - head_dim, - full_att_layer_num, - linear_config: LinearAttCacheConfig, - always_copy=False, - mem_fraction=0.9, - ): - self.linear_config = linear_config - - super().__init__(size, dtype, num_kv_heads, head_dim, full_att_layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - return super().get_att_input_params(layer_index) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - # TODO 初始化线性 att 对应的部分 buffer. - self._init_linear_att_buffers() - return - - def _init_linear_att_buffers(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - # 申请大页可能需要对应的资源, 多申请了两个linear att的状态，理论上这个状态 - # 永远不会被 alloc 申请到，只会在 cpu cache中，用于过渡和存储碎页情况下的 - # cpu cache 的页面拷贝。 - self.big_page_buffers = LinearAttCacheManager( - size=triton.cdiv(self.size, big_page_token_num) + 2, - linear_config=self.linear_config, - keep_num=2, - ) - - self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 - self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 - return - - def _free_buffers(self): - super()._free_buffers() - self._free_linear_att_buffers() - return - - def _free_linear_att_buffers(self): - self.big_page_buffers = None - return - - def write_to_shm(self, req_manager): - self.req_to_conv_state = req_manager.req_to_conv_state - self.req_to_ssm_state = req_manager.req_to_ssm_state - # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 - # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 - # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 - # 启动会报 "Pointer argument cannot be accessed from Triton (cpu tensor?)"。 - # 跨进程消费方并不使用 cpu 侧大页 state cache，序列化期间临时剔除以保住 pinned。 - big_page_buffers = self.big_page_buffers - self.big_page_buffers = None - try: - return super().write_to_shm(req_manager) - finally: - self.big_page_buffers = big_page_buffers - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - self.att_state_page_helper = self._create_att_state_page_helper() - self.att_state_page_helper.assert_page_size() - return kv_move_buffer - - def _create_att_state_page_helper(self): - return Qwen3NextLinearAttPageHelper(self) - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().write_mem_to_page_kv_move_buffer( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().read_page_kv_move_buffer_to_mem( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - -class _FP8StaticPerHeadQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerHeadQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class _FP8StaticPerTensorQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerTensorQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class FP8StaticPerHeadQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerHeadQuantMemManager): - operator_class = _FP8StaticPerHeadQuantLinearAttMemOperator - - -class FP8StaticPerTensorQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerTensorQuantMemManager): - operator_class = _FP8StaticPerTensorQuantLinearAttMemOperator - - -class Qwen3NextLinearAttPageHelper: - def __init__(self, mem_manager: "Qwen3NextMemManager"): - self.mem_manager = mem_manager - self.linear_config = mem_manager.linear_config - self.req_to_conv_state = mem_manager.req_to_conv_state - self.req_to_ssm_state = mem_manager.req_to_ssm_state - self.global_linear_k_heads = self.linear_config.global_linear_k_heads - self.global_linear_v_heads = self.linear_config.global_linear_v_heads - - self.global_q_dim = self.global_linear_k_heads * self.linear_config.head_linear_k_dim - self.global_k_dim = self.global_q_dim - self.global_v_heads = self.global_linear_v_heads - self.global_v_dim = self.global_v_heads * self.linear_config.head_linear_v_dim - # conv state follows mixed_qkv layout: [q, k, v], each as a flat channel block. - self.conv_shape = ( - self.linear_config.linear_layer_num, - self.global_q_dim + self.global_k_dim + self.global_v_dim, - self.linear_config.conv_kernel_size - 1, - ) - self.ssm_shape = ( - self.linear_config.linear_layer_num, - self.global_v_heads, - self.linear_config.head_linear_k_dim, - self.linear_config.head_linear_v_dim, - ) - - self.conv_nbytes = ( - self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.req_to_conv_state.buffer.element_size() - ) - ssm_alignment = self.req_to_ssm_state.buffer.element_size() - # 做一下字节对齐，防止切出来的不对齐，导致一些操作的性能下降。 - self.ssm_offset = ((self.conv_nbytes + ssm_alignment - 1) // ssm_alignment) * ssm_alignment - self.ssm_nbytes = ( - self.ssm_shape[0] - * self.ssm_shape[1] - * self.ssm_shape[2] - * self.ssm_shape[3] - * self.req_to_ssm_state.buffer.element_size() - ) - self.state_nbytes = self.ssm_offset + self.ssm_nbytes - - def assert_page_size(self): - kv_move_buffer = self.mem_manager.kv_move_buffer - page_nbytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() - assert ( - page_nbytes >= self.state_nbytes - ), f"nixl kv move page bytes {page_nbytes} is smaller than global linear att state bytes {self.state_nbytes}" - return - - def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, dp_world_size: int): - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - assert len(dp_mems) == dp_world_size - for mem in dp_mems: - assert hasattr(mem, "req_to_conv_state") and hasattr(mem, "req_to_ssm_state") - assert mem.linear_config.linear_layer_num == self.linear_config.linear_layer_num - assert mem.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size - assert mem.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim - assert mem.linear_config.head_linear_v_dim == self.linear_config.head_linear_v_dim - assert mem.linear_config.num_linear_k_heads == self.linear_config.num_linear_k_heads - assert mem.linear_config.num_linear_v_heads == self.linear_config.num_linear_v_heads - return dp_mems - - def view_page_to_linear_att_state(self, page_index: int): - page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - conv_page = page_bytes[0 : self.conv_nbytes].view(self.req_to_conv_state.buffer.dtype).view(self.conv_shape) - ssm_page = ( - page_bytes[self.ssm_offset : self.ssm_offset + self.ssm_nbytes] - .view(self.req_to_ssm_state.buffer.dtype) - .view(self.ssm_shape) - ) - return conv_page, ssm_page - - def write_req_to_page( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._write_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def read_page_to_req( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._read_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def _get_req_state_indexes(self, req_idx: int): - mtp_size = get_env_start_args().mtp_step + 1 - # Conv is one widened slot per request; SSM keeps the historical S+1 block layout. - return req_idx, req_idx * mtp_size - - def _write_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) - conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] - ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] - self._copy_conv_state_to_page(conv_state, conv_page, mem, tp_index) - self._copy_ssm_state_to_page(ssm_state, ssm_page, mem, tp_index) - return - - def _copy_conv_state_to_page( - self, - conv_state: torch.Tensor, - conv_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=True - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - global_q_page[:, global_dim_start:global_dim_end, :].copy_( - local_q_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - global_k_page[:, global_dim_start:global_dim_end, :].copy_( - local_k_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=True - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - global_v_page[:, global_dim_start:global_dim_end, :].copy_( - local_v_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - return - - def _copy_ssm_state_to_page( - self, - ssm_state: torch.Tensor, - ssm_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=True, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_page[:, global_head_start:global_head_end, :, :].copy_( - ssm_state[:, local_head_start:local_head_end, :, :], - non_blocking=True, - ) - return - - def _get_head_slice( - self, - tp_index: int, - local_heads: int, - global_heads: int, - tp_world_size: int, - is_write: bool, - ): - if local_heads == 0 or global_heads == 0: - return None - total_local_heads = local_heads * tp_world_size - repeat_count = max(1, total_local_heads // global_heads) - if is_write and repeat_count > 1 and tp_index % repeat_count != 0: - return None - unique_tp_index = tp_index // repeat_count - global_head_start = unique_tp_index * local_heads - global_head_end = min(global_head_start + local_heads, global_heads) - local_head_start = 0 - local_head_end = global_head_end - global_head_start - if local_head_end <= local_head_start: - return None - return local_head_start, local_head_end, global_head_start, global_head_end - - def _copy_page_to_conv_state( - self, - conv_page: torch.Tensor, - conv_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=False - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - local_q_state[:, local_dim_start:local_dim_end, :].copy_( - global_q_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - local_k_state[:, local_dim_start:local_dim_end, :].copy_( - global_k_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=False - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - local_v_state[:, local_dim_start:local_dim_end, :].copy_( - global_v_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - return - - def _copy_page_to_ssm_state( - self, - ssm_page: torch.Tensor, - ssm_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=False, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_state[:, local_head_start:local_head_end, :, :].copy_( - ssm_page[:, global_head_start:global_head_end, :, :], - non_blocking=True, - ) - return - - def _read_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) - conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] - ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] - self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index) - self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index) - return diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index 7f2e6ac9b5..711915bd04 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,15 +1,7 @@ from .base import ReqManager -from .deepseek4 import DeepseekV4ReqManager -from .linear_att import ReqManagerForMamba -from .glm5_next import Glm5NextReqManager -from .hybrid_base import HybridAttentionReqManager from .req_sampling_params import ReqSamplingParamsManager __all__ = [ "ReqManager", - "DeepseekV4ReqManager", - "HybridAttentionReqManager", - "ReqManagerForMamba", - "Glm5NextReqManager", "ReqSamplingParamsManager", ] diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index cfe5d716ee..0088c6b82b 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -2,7 +2,6 @@ import torch -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger @@ -52,7 +51,7 @@ def is_all_free(self): class ReqManager: - def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryManager): + def __init__(self, max_request_num, max_sequence_length): from .req_sampling_params import ReqSamplingParamsManager # 这里对最大请求数量的管理在默认上多申请了一个，主要是 index 为 max_request_num 代表 @@ -60,7 +59,7 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana # 的那个batch size 进行运行，所有 padding 的请求都会使用预留的这个请求管理 id 进行处理 # 这样让 DP 的实现更为简化一些。 self.req_list = _ReqLinkedList(max_request_num) - page_size = get_env_start_args().page_size + self.page_size = page_size = get_env_start_args().page_size max_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size self.req_to_token_indexs = torch.zeros( (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cuda" @@ -68,43 +67,31 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num - self.mem_manager = None - if mem_manager is not None: - self.bind_mem_manager(mem_manager) - def bind_mem_manager(self, mem_manager: MemoryManager): - self.mem_manager = mem_manager - - self.init_hold_request_indexs() - return - - def init_hold_request_indexs(self): + def bind_hold_page(self, token_indices): + """Map padding requests to the resource owner's reserved physical page.""" assert ( self.req_list.is_all_free() ), "hold request indexes can only be initialized when all requests are released" - - # HOLD_REQUEST_ID 对应的请求行供 DP padding、overlap microbatch 等占位请求使用。将该行 - # 按 page_size 划分后，每一页都映射到 mem_manager 额外保留的同一个物理页；这样占位请求 - # 无论访问哪一个逻辑位置，都会落到合法且不会参与正常分配的 KV cache 地址上。 hold_row = self.req_to_token_indexs[self.HOLD_REQUEST_ID] - hold_page = torch.tensor( - self.mem_manager.HOLD_TOKEN_MEMINDEXES, + hold_page = torch.as_tensor( + token_indices, dtype=hold_row.dtype, device=hold_row.device, ) - hold_row.view(-1, self.mem_manager.page_size).copy_(hold_page) - return + if hold_page.ndim != 1 or hold_page.numel() != self.page_size: + raise ValueError("HOLD indices must describe exactly one physical page") + hold_row.view(-1, self.page_size).copy_(hold_page) def alloc(self): return self.req_list.alloc() - def free(self, free_req_indexes: List[int], free_token_index): + def free(self, free_req_indexes: List[int]): for req_index in free_req_indexes: self.req_list.free(req_index) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") - self.mem_manager.free(free_token_index) def free_req(self, free_req_index: int): self.req_list.free(free_req_index) @@ -112,10 +99,6 @@ def free_req(self, free_req_index: int): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") return - def free_token(self, free_token_index): - self.mem_manager.free(free_token_index) - return - def free_all(self): self.req_list = _ReqLinkedList(self.max_request_num) return diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py deleted file mode 100644 index 492dfeb73c..0000000000 --- a/lightllm/common/req_manager/glm5_next.py +++ /dev/null @@ -1,35 +0,0 @@ -from lightllm.common.state_cache_manager import LayerCache -from lightllm.utils.envs_utils import get_env_start_args - -from .linear_att import ReqManagerForMamba - - -class Glm5NextReqManager(ReqManagerForMamba): - """KDA runtime state and the NSA indexer's incomplete four-token pool.""" - - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config): - # Both checkpoint sizes are multiples of this hash page. A restored - # prefix therefore has no incomplete K-pool to serialize or replay. - assert ( - get_env_start_args().linear_att_hash_page_size % linear_config.index_kpool == 0 - ), "GLM K-pool requires cache pages aligned to index_kpool" - super().__init__(max_request_num, max_sequence_length, mem_manager, linear_config) - self.req_to_indexer_tail = LayerCache( - size=max_request_num + 1, - # Raw keys/gates are computation state, even when packed KV is FP8. - dtype=linear_config.conv_state_dtype, - shape=(linear_config.index_kpool + self.mtp_step, 2 * linear_config.index_head_dim), - layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), - device="cuda", - ) - - def get_indexer_tail_buffer(self, layer_index): - return self.req_to_indexer_tail.buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] - - def init_hybrid_attention_state(self, req): - super().init_hybrid_attention_state(req) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() - - def restore_state(self, req, state_cache_manager, buffer_idx): - super().restore_state(req, state_cache_manager, buffer_idx) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py deleted file mode 100644 index 1f78003317..0000000000 --- a/lightllm/common/req_manager/hybrid_base.py +++ /dev/null @@ -1,69 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List - -import torch - -from .base import ReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class HybridAttentionReqManager(ReqManager, ABC): - """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 - - 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 - linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， - 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 - - 请求的 GPU 计算状态由具体实现管理；small_page_buffers 持有 CPU 小页快照； - big_page_buffers 引用 mem_manager 的 CPU 大页快照，与 full KV 容量一起创建、调整。 - 大小页不包含额外的 GPU 运行态，命中后恢复到请求状态，不重新分配 checkpoint 池。 - 公共缓存流程负责 checkpoint 槽位分配、边界、匹配与淘汰；本接口负责运行态与 checkpoint 保存/恢复。 - """ - - def __init__(self, max_request_num, max_sequence_length, mem_manager): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.small_page_buffers = None - - @property - def big_page_buffers(self): - return self.mem_manager.big_page_buffers - - @abstractmethod - def create_small_page_cache_manager(self, size: int): - """创建并持有 CPU 小页池，返回同一池供 radix 使用；size 是槽位数，不是 token 数。""" - - def init_hybrid_attention_state(self, req: "InferReq"): - """无前缀缓存命中时初始化 GPU 运行态；需要初始化的模型覆写。""" - - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - """将指定大页槽位的 CPU checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.big_page_buffers, big_page_buffer_idx) - - def restore_small_page_state(self, req: "InferReq"): - """将 req.shared_kv_node 对应的小页 checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.small_page_buffers, req.shared_kv_node.small_page_buffer_idx) - - @abstractmethod - def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): - """CPU checkpoint → 请求 GPU 运行态；大小页共用，不负责前缀匹配或 full KV 索引恢复。""" - - def save_big_page_states(self, b_req_idx: torch.Tensor, reqs: List["InferReq"], buffer_indexes: List[int]): - """批量保存请求 GPU 运行态到已分配的大页槽位，buffer_indexes 中的 -1 表示跳过。 - - b_req_idx 与 reqs 分别为同一批请求的 GPU 索引张量和 CPU 请求对象列表。 - 默认逐请求保存，模型可覆盖为批量拷贝算子。 - """ - for req, buffer_idx in zip(reqs, buffer_indexes): - if buffer_idx != -1: - self.save_state(req, buffer_idx, self.big_page_buffers) - - @abstractmethod - def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager): - """保存本轮 prefill 终点的运行态；在更新 req.cur_kv_len 前调用，大小页共用。""" - - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): - """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" - return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py deleted file mode 100644 index 5ae9f4966c..0000000000 --- a/lightllm/common/req_manager/linear_att.py +++ /dev/null @@ -1,119 +0,0 @@ -from typing import TYPE_CHECKING, List - -import torch - -from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig, LinearAttCacheManager -from lightllm.utils.envs_utils import get_env_start_args - -from .hybrid_base import HybridAttentionReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class ReqManagerForMamba(HybridAttentionReqManager): - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.mtp_step = get_env_start_args().mtp_step - # 因为在mtp的推理中，需要标记每个请求对应的mtp index状态(conv state 和 ssm state)，在mtp对应序列中 - # 的真实位置，所以需要需要一个标记来记录，不然算子无法找到真实的处理起点。 - self.req_to_mtp_state_index = ( - torch.zeros((max_request_num + 1,), dtype=torch.int32, device="cuda") if self.mtp_step > 0 else None - ) - # 突然想到， 在linear att 开启mtp的模式中，现在的prefill linear att 算子默认是从0的位置读取信息进行操作 - # 所以不能支持 prefill decode mixed 操作了，因为一个decode过的请求，重新用prefill 算子跑，会出现读错linear - # 状态位置的问题。导致bug, 在这里加个断言，以后可以支持上 TODO - if self.mtp_step > 0: - assert get_env_start_args().enable_prefill_decode_mixed is False - - self.big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - self.linear_config = linear_config - - self.req_to_conv_state = LayerCache( - size=(max_request_num + 1), - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step), - layer_num=self.linear_config.linear_layer_num, - device="cuda", - ) - self.req_to_ssm_state = LayerCache( - size=(max_request_num + 1) * (self.mtp_step + 1), - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cuda", - ) - return - - def init_hybrid_attention_state(self, req: "InferReq"): - conv_index = req.req_idx - ssm_start = req.req_idx * (self.mtp_step + 1) - self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) - # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future - # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN). - self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return - - def create_small_page_cache_manager(self, size: int): - self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) - return self.small_page_buffers - - def save_big_page_states(self, b_req_idx: torch.Tensor, reqs: List["InferReq"], buffer_indexes: List[int]): - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer - - buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) - state_cache_manager = self.big_page_buffers - copy_linear_att_state_to_kv_buffer( - b_req_idx=b_req_idx, - big_page_buffer_ids=buffer_indexes, - gpu_conv_state=self.req_to_conv_state.buffer, - gpu_ssm_state=self.req_to_ssm_state.buffer, - cpu_kv_conv_state=state_cache_manager.conv_state_cache.buffer, - cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, - mtp_step=self.mtp_step, - ) - return - - def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager: LinearAttCacheManager): - # checkpoint 只保存标准 conv 窗口和请求的基准 SSM 状态，不包含 MTP 扩展运行态。 - req_idx = req.req_idx - conv_cache_width = self.linear_config.get_conv_state_shape()[-1] - gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] - gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] - dst_conv_state, dst_ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) - dst_conv_state.copy_(gpu_conv_state, non_blocking=True) - dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) - - def get_mamba_cache(self, layer_idx_in_all: int): - layer_idx_in_linear = self.linear_config.get_linear_state_layer_index(layer_idx_in_all) - conv_states = self.req_to_conv_state.buffer[layer_idx_in_linear] - ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] - return conv_states, ssm_states - - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): - from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update - - linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.req_to_mtp_state_index, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=verify_width, - ) - - def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): - conv_state, ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) - conv_dest = req.req_idx - ssm_dest = req.req_idx * (self.mtp_step + 1) - conv_cache_width = conv_state.shape[-1] - self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state - self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py deleted file mode 100644 index 758c6a88d4..0000000000 --- a/lightllm/common/state_cache_manager/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -from .base import StateCacheManager -from .deepseek4 import DeepseekV4CacheConfig, DeepseekV4StateCacheManager -from .layer_cache import LayerCache -from .linear_att import LinearAttCacheConfig, LinearAttCacheManager -from .glm5_next import Glm5NextCacheConfig - - -def get_hybrid_cache_config(): - """Return the model-specific layout used by hybrid CPU/disk cache pages.""" - from transformers.configuration_utils import PretrainedConfig - from lightllm.utils.envs_utils import get_env_start_args - - args = get_env_start_args() - model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) - model_type = model_cfg["model_type"] - if model_type == "deepseek_v4": - return DeepseekV4CacheConfig.from_model_config(model_cfg, args) - if model_type in ("glm5_next", "glm5_next_text"): - return Glm5NextCacheConfig.from_model_config(model_cfg, args) - if model_type in ("qwen3_5", "sensenova_loop", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): - return LinearAttCacheConfig.load_from_args() - raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/base.py b/lightllm/common/state_cache_manager/base.py deleted file mode 100644 index aa131fe784..0000000000 --- a/lightllm/common/state_cache_manager/base.py +++ /dev/null @@ -1,57 +0,0 @@ -import collections -from abc import ABC, abstractmethod -from typing import List, Optional - -from lightllm.utils.log_utils import init_logger - - -logger = init_logger(__name__) - - -class StateCacheManager(ABC): - """CPU checkpoint 槽位池；大小页分别实例化，具体状态布局由子类负责。 - - 尾部 keep_num 个槽位保留给 CPU cache 碎页传输，不参与普通分配。 - 本类不管理 GPU 运行态，也不判断页面边界、前缀匹配或淘汰策略。 - """ - - def __init__(self, size: int, keep_num: int = 0): - self.size = size - self.keep_num = keep_num - assert 0 <= keep_num <= size, f"invalid keep_num {keep_num} for size {size}" - self.free_list = collections.deque(range(size - keep_num)) - - @abstractmethod - def get_state_cache(self, buffer_idx: int): - """返回指定槽位的状态视图；可以是单个 Tensor 或多个 Tensor。""" - - def alloc_one_state_cache(self) -> Optional[int]: - return None if not self.free_list else self.free_list.popleft() - - def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: - if need_size > len(self.free_list): - logger.error(f"warn no enough cache need_size {need_size} free_size {len(self.free_list)}") - return None - return [self.free_list.popleft() for _ in range(need_size)] - - def free_state_cache(self, free_indexes: List[int]): - alloc_upper_bound = self.size - self.keep_num - for idx in free_indexes: - assert 0 <= idx < alloc_upper_bound, ( - f"free index {idx} out of alloc range [0, {alloc_upper_bound}), " f"reserved tail num {self.keep_num}" - ) - self.free_list.extend(free_indexes) - assert ( - len(self.free_list) <= alloc_upper_bound - ), f"free cache num {len(self.free_list)} should not be larger than alloc size {alloc_upper_bound}" - - def get_free_cache_num(self): - return len(self.free_list) - - def get_used_cache_num(self): - # Preserve the existing accounting: reserved slots count as used. - return self.size - len(self.free_list) - - def clear_to_init_state(self): - """重置空闲槽位；子类同时清零自身的 checkpoint buffer。""" - self.free_list = collections.deque(range(self.size - self.keep_num)) diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py deleted file mode 100644 index f0928c2c18..0000000000 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ /dev/null @@ -1,64 +0,0 @@ -import dataclasses -import torch - -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.torch_dtype_utils import get_torch_dtype - -from .linear_att import LinearAttCacheConfig - - -@dataclasses.dataclass -class Glm5NextCacheConfig(LinearAttCacheConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" - - index_kpool: int = 4 - index_head_dim: int = 128 - - INDEXER_K_BYTES = 132 - # FlashMLA V4 layout: 448 FP8 values, 64 BF16 values, seven E8M0 scales and one padding byte. - FP8_MLA_KV_BYTES = 584 - - @classmethod - def from_model_config(cls, config, args): - config = config.get("text_config", config) - linear = config["linear_attn_config"] - tp = args.tp // args.dp - assert linear["num_heads"] % tp == 0 - dtype = get_torch_dtype(args.data_type) - use_fp8_kv = args.llm_kv_type == "fp8kv_dsa" - full_att_dtype = torch.uint8 if use_fp8_kv else dtype - mla_kv_bytes = cls.FP8_MLA_KV_BYTES if use_fp8_kv else config["kv_lora_rank"] * dtype.itemsize - # FlashMLA loads 16-byte vectors; keep every packed token row aligned. - packed_kv_bytes = ((mla_kv_bytes + cls.INDEXER_K_BYTES + 15) // 16) * 16 - # Raw index keys and compression scores only live in request tails. - return cls( - tp_world_size=tp, - full_att_all_num_kv_heads=1, - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=1, - full_att_head_dim=packed_kv_bytes // full_att_dtype.itemsize, - global_linear_k_heads=linear["num_heads"], - global_linear_v_heads=linear["num_heads"], - num_linear_k_heads=linear["num_heads"] // tp, - num_linear_v_heads=linear["num_heads"] // tp, - head_linear_k_dim=linear["head_dim"], - head_linear_v_dim=linear["head_dim"], - conv_kernel_size=linear["short_conv_kernel_size"], - conv_state_dtype=dtype, - ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - layer_types=config["layer_types"], - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, - index_kpool=config["index_kpool"], - index_head_dim=config["index_head_dim"], - ) - - def get_cpu_cache_full_att_bytes(self): - args = get_env_start_args() - page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - assert page_tokens == args.cpu_cache_token_page_size - return ( - self.full_att_head_dim - * self.full_att_dtype.itemsize - * self.get_full_att_kv_layer_num_with_draft_model() - * page_tokens - ) diff --git a/lightllm/common/state_cache_manager/layer_cache.py b/lightllm/common/state_cache_manager/layer_cache.py deleted file mode 100644 index 9ac72acca3..0000000000 --- a/lightllm/common/state_cache_manager/layer_cache.py +++ /dev/null @@ -1,37 +0,0 @@ -from typing import Tuple -import torch -import numpy as np -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -class LayerCache: - def __init__( - self, - size: int, - dtype: torch.dtype, - shape: Tuple[int, ...], - layer_num: int, - device: torch.device, - size_first: bool = False, - ): - self.size_first = size_first - self.size = size - self.dtype = dtype - self.shape = shape - self.layer_num = layer_num - self.device = device - if not self.size_first: - if device == "cpu": - self.buffer = torch.zeros((self.layer_num, size, *shape), dtype=dtype, device="cpu", pin_memory=True) - else: - self.buffer = torch.zeros((self.layer_num, size, *shape), dtype=dtype, device=device) - else: - if device == "cpu": - self.buffer = torch.zeros((size, self.layer_num, *shape), dtype=dtype, device="cpu", pin_memory=True) - else: - self.buffer = torch.zeros((size, self.layer_num, *shape), dtype=dtype, device=device) - - def get_cell_size(self): - return np.prod(self.shape) * self.layer_num * torch._utils._element_size(self.dtype) diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py deleted file mode 100644 index 496791d489..0000000000 --- a/lightllm/common/state_cache_manager/linear_att.py +++ /dev/null @@ -1,198 +0,0 @@ -import torch -import dataclasses -import triton - -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.log_utils import init_logger -from lightllm.utils.torch_dtype_utils import get_torch_dtype - -from .base import StateCacheManager -from .layer_cache import LayerCache - - -logger = init_logger(__name__) - - -@dataclasses.dataclass -class LinearAttCacheConfig: - tp_world_size: int - # full att 的参数 - full_att_all_num_kv_heads: int - full_att_dtype: torch.dtype - full_att_num_kv_heads: int # 这个是 tp 后的head头数量 - full_att_head_dim: int - - # linear att 的参数 - global_linear_k_heads: int - global_linear_v_heads: int - num_linear_k_heads: int - num_linear_v_heads: int - head_linear_k_dim: int - head_linear_v_dim: int - conv_kernel_size: int - conv_state_dtype: torch.dtype - ssm_state_dtype: torch.dtype - layer_types: tuple[str, ...] - draft_full_att_kv_layer_num: int = 0 - linear_layer_num: int = dataclasses.field(init=False) - all_layer_num: int = dataclasses.field(init=False) - - def __post_init__(self): - self.layer_types = tuple(self.layer_types) - indices = [] - full_count = linear_count = 0 - for kind in self.layer_types: - if kind == "linear_attention": - indices.append(linear_count) - linear_count += 1 - else: - indices.append(full_count) - full_count += 1 - # MTP draft layers only use full-attention KV slots after the main model. - indices.extend(range(full_count, full_count + self.draft_full_att_kv_layer_num)) - self._layer_indices = tuple(indices) - self.linear_layer_num = linear_count - self.all_layer_num = len(self.layer_types) - - def get_linear_state_layer_index(self, layer_index: int) -> int: - return self._layer_indices[layer_index] - - def get_conv_dim(self): - # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 - # 由于 k_dim = q_dim, k_heads = q_heads, 所以第一项和第二项的计算 - # 形式相同，但是实际内在含义是不同的。 - return ( - self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_v_dim * self.num_linear_v_heads - ) - - def get_main_model_full_att_layer_num(self): - full_att_layer_num = self.all_layer_num - self.linear_layer_num - return full_att_layer_num - - def get_full_att_kv_layer_num_with_draft_model(self): - return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num - - def get_full_att_kv_layer_index(self, layer_index: int) -> int: - """Map a global target/draft layer index to the packed full-attention cache.""" - - return self._layer_indices[layer_index] - - def get_conv_state_shape(self): - # Base committed sliding-window state, without speculative MTP tail. - return (self.get_conv_dim(), self.conv_kernel_size - 1) - - def get_mtp_conv_state_shape(self, mtp_step: int): - # Working state with room for S speculative tokens before acceptance. - return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step) - - def get_ssm_state_shape(self): - return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim) - - def get_conv_state_bytes_per_layer(self): - return self.get_conv_dim() * (self.conv_kernel_size - 1) * self.conv_state_dtype.itemsize - - def get_ssm_state_bytes_per_layer(self): - return self.num_linear_v_heads * self.head_linear_k_dim * self.head_linear_v_dim * self.ssm_state_dtype.itemsize - - def get_cpu_cache_big_page_bytes(self): - a = self.get_cpu_cache_full_att_bytes() - b = self.get_cpu_cache_conv_bytes() - c = self.get_cpu_cache_ssm_bytes() - - return triton.cdiv(a + b + c, 16) * 16 - - def get_cpu_cache_full_att_bytes(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - assert big_page_token_num == get_env_start_args().cpu_cache_token_page_size - full_att_bytes = 2 * self.full_att_all_num_kv_heads * self.full_att_head_dim * self.full_att_dtype.itemsize - a = full_att_bytes * self.get_full_att_kv_layer_num_with_draft_model() * big_page_token_num - return a - - def get_cpu_cache_conv_bytes(self): - b = self.get_conv_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return b - - def get_cpu_cache_ssm_bytes(self): - c = self.get_ssm_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return c - - @staticmethod - def load_from_args() -> "LinearAttCacheConfig": - args = get_env_start_args() - model_path = args.model_dir - from transformers.configuration_utils import PretrainedConfig - - model_cfg, _ = PretrainedConfig.get_config_dict(model_path) - llm_config = model_cfg.get("text_config", model_cfg) - layer_types = llm_config["layer_types"] - if model_cfg["model_type"] == "sensenova_loop": - from lightllm.models.sensenova_loop.config import get_loop_layer_types - - layer_types = get_loop_layer_types(model_cfg["loop_settings"], llm_config) - tp_world_size = args.tp // args.dp - full_att_dtype = ( - torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) - ) - return LinearAttCacheConfig( - tp_world_size=tp_world_size, - full_att_all_num_kv_heads=llm_config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=max(1, llm_config["num_key_value_heads"] // tp_world_size), - full_att_head_dim=llm_config["head_dim"], - global_linear_k_heads=llm_config["linear_num_key_heads"], - global_linear_v_heads=llm_config["linear_num_value_heads"], - num_linear_k_heads=max(1, llm_config["linear_num_key_heads"] // tp_world_size), - num_linear_v_heads=max(1, llm_config["linear_num_value_heads"] // tp_world_size), - head_linear_k_dim=llm_config["linear_key_head_dim"], - head_linear_v_dim=llm_config["linear_value_head_dim"], - conv_kernel_size=llm_config["linear_conv_kernel_dim"], - conv_state_dtype=get_torch_dtype(args.data_type), - ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - layer_types=layer_types, - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), - ) - - -class LinearAttCacheManager(StateCacheManager): - """CPU pinned conv/SSM checkpoint，两个 buffer 均保持 size-first 布局。""" - - def __init__( - self, - size: int, - linear_config: LinearAttCacheConfig, - keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 - ): - super().__init__(size, keep_num) - self.linear_config = linear_config - # init the layer cache - self.conv_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_conv_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.ssm_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.clear_to_init_state() - return - - def get_state_cache(self, buffer_idx: int): - return self.conv_state_cache.buffer[buffer_idx, ...], self.ssm_state_cache.buffer[buffer_idx, ...] - - def clear_to_init_state(self): - self.conv_state_cache.buffer.zero_() - self.ssm_state_cache.buffer.zero_() - super().clear_to_init_state() - return diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index c8a78e6b62..1b1bada67c 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -28,7 +28,8 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -43,7 +44,8 @@ def _context_attention_kernel( def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 7368ed1dc3..d44e6f63da 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -24,6 +24,14 @@ def __init__(self, kvargs): super().__init__(kvargs) return + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _init_config(self): super()._init_config() # rename key diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 3254031056..58fd1ba930 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -103,7 +103,8 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - kv = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + kv = state.get_layer_view(layer_id=self.layer_num_) out = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), @@ -119,7 +120,8 @@ def _decompress_kv( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - compressed_kv = infer_state.mem_manager.kv_buffer[self.layer_num_] + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + compressed_kv = state.get_layer_view(self.layer_num_) total_token_num = infer_state.total_token_num sampled_compressed_kv_nope = self.alloc_tensor( diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 26a7ecd81f..a70f21b6b8 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -4,9 +4,8 @@ from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args, get_added_mtp_kv_layer_num +from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args from lightllm.distributed.communication_op import dist_group_manager from lightllm.common.basemodel.attention import get_mla_decode_att_backend_class, get_mla_prefill_att_backend_class @@ -57,18 +56,13 @@ def _init_custom(self): def _verify_params(self): return super()._verify_params() - def _init_mem_manager(self): - manager_class = select_mem_manager_class() + @classmethod + def _full_attention_options(cls, config, args): + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options - self.mem_manager = manager_class( - self.max_total_token_num, - dtype=self.data_type, - head_num=1, - head_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + options = get_full_attention_options(config, args) + options.update(heads=1, dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], mla=True) + return options def _init_to_get_yarn_rotary(self): from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask, get_deepseek_mscale diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py index b9be73e278..a086349983 100644 --- a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py +++ b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py @@ -194,7 +194,7 @@ def _fwd_kernel_calcu_index_and_block_seq( infer_state.batch_size = Z infer_state.max_kv_seq_len = N_CTX infer_state.total_token_num = Z * N_CTX - infer_state.req_manager = ReqManager(Z, N_CTX, None) + infer_state.req_manager = ReqManager(Z, N_CTX) infer_state.req_manager.req_to_token_indexs = req_to_token_indexs infer_state.b_req_idx = b_req_idx infer_state.b_seq_len = b_seq_len diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index d6eaebe2fd..a812f8750f 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -99,9 +99,10 @@ def _context_attention_kernel( }, ) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) mla_out = infer_state.prefill_att_state.prefill_att( q=q_all, - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -141,9 +142,10 @@ def _token_attention_kernel( }, ) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) o_tensor = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -196,11 +198,12 @@ def _get_indices( q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) destindex_copy_indexer_ks( K_fp8=k_fp8, K_scale=k_scale, DestLoc=infer_state.mem_index, - O_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + O_buffer=state.get_layer_view(self.layer_idx_).indexer_k, ) weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale @@ -216,7 +219,7 @@ def _get_indices( mtp_step = get_env_start_args().mtp_step # Use efficient Triton kernel to extract FP8 keys and scales from buffer k_fp8_, k_scale_ = extract_indexer_ks( - I_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + I_buffer=state.get_layer_view(self.layer_idx_).indexer_k, b_seq_len=infer_state.b_seq_len, b_req_idx=infer_state.b_req_idx, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 57b6739e6f..197b25fafe 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -12,6 +12,12 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): # infer class transformer_layer_infer_class = Deepseek3_2TransformerLayerInfer + @classmethod + def _full_attention_options(cls, config, args): + options = super()._full_attention_options(config, args) + options["packed"] = True + return options + def _init_att_backend(self): self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0)(model=self) self.decode_att_backend = get_nsa_decode_att_backend_class(index=0)(model=self) diff --git a/lightllm/models/deepseek_mtp/model.py b/lightllm/models/deepseek_mtp/model.py index e2b2a56137..c5c4a9aa61 100644 --- a/lightllm/models/deepseek_mtp/model.py +++ b/lightllm/models/deepseek_mtp/model.py @@ -32,10 +32,6 @@ def _init_req_manager(self): self.req_manager = self.main_model.req_manager return - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.pre_post_weight = self.pre_and_post_weight_class( diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index eabbfe95e0..5c769842b9 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -2,8 +2,6 @@ import json import torch from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.models.gemma3.infer_struct import Gemma3InferStateInfo from lightllm.models.gemma3.layer_infer.post_layer_infer import Gemma3PostLayerInfer from lightllm.models.gemma3.layer_infer.pre_layer_infer import Gemma3PreLayerInfer @@ -138,16 +136,13 @@ def _init_custom(self): self.head_dim_ = 256 self._init_to_get_rotary() - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=torch.bfloat16, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=256, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + @classmethod + def _full_attention_options(cls, config, args): + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options + + options = get_full_attention_options(config, args) + options.update(dim=256, dtype=torch.bfloat16) + return options def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2f0c01dbf6..742b2fc996 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -21,7 +21,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): time. The KV cache layout is uniform (sliding shape: num_kv_heads=16, head_dim=256); full-attention layers pack their (4, 512) tensor into the first 8 heads of the 16-head slot at cache-write time, then reshape on - read. See Gemma4TpPartModel._init_mem_manager for context. + read. See Gemma4TpPartModel._full_attention_options for context. """ def __init__(self, layer_num, network_config): @@ -195,7 +195,8 @@ def _att_control(self): def _get_layer_kv(self, infer_state: InferStateInfo): # KV-shared layers read from the target layer's cache slot. layer_idx = self.kv_share_target_layer_ if self.is_kv_shared_ else self.layer_num_ - _k_raw, _v_raw = infer_state.mem_manager.get_att_input_params(layer_index=layer_idx) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k_raw, _v_raw = state.get_layer_view(layer_id=layer_idx) # _k_raw / _v_raw shape (S, cache_slot_num, cache_slot_dim). Use .view # (not .reshape) so any non-contiguous layout from a future mem_manager # backend fails loudly instead of silently copying — slice + view is diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index b68a1a4399..61e933ef5f 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -2,7 +2,6 @@ import json import torch from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.build_utils import repair_config from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo @@ -11,7 +10,7 @@ from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer from lightllm.models.gemma4.layer_weights.pre_and_post_layer_weight import Gemma4PreAndPostLayerWeight from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.distributed.communication_op import dist_group_manager @@ -31,7 +30,7 @@ class Gemma4TpPartModel(LlamaTpPartModel): def __init__(self, kvargs): # head_dim_ is used by the default _init_to_get_rotary which we # override; still set it to the sliding-layer head_dim for consistency - # with the mem manager and any generic helpers. + # with the full-attention storage and any generic helpers. self.head_dim_ = 256 super().__init__(kvargs) return @@ -80,7 +79,10 @@ def _verify_params(self): ) return - def _init_mem_manager(self): + @classmethod + def _full_attention_options(cls, config, args): + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options + # Uniform per-layer KV cache layout. The per-layer cache slot must fit # whichever layer type has the largest per-token K/V width: sliding # (num_key_value_heads * head_dim) or full @@ -90,21 +92,16 @@ def _init_mem_manager(self): # (2*512 > 2*256), so it uses 4 storage slots of 256 dims. # Gemma4TransformerLayerInfer.__init__ computes the same value and # uses it to pack/unpack K/V at write/read time. - head_dim = self.config["head_dim"] - num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] - sliding_total = self.config["num_key_value_heads"] * self.config["head_dim"] - full_total = num_global_kv * self.config["global_head_dim"] + head_dim = config["head_dim"] + num_global_kv = config.get("num_global_key_value_heads") or config["num_key_value_heads"] + sliding_total = config["num_key_value_heads"] * head_dim + full_total = num_global_kv * config["global_head_dim"] per_token_k_width = max(sliding_total, full_total) - head_num_per_rank = (per_token_k_width // head_dim) // self.tp_world_size_ - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=head_num_per_rank, - head_dim=head_dim, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + if per_token_k_width % head_dim: + raise ValueError("Gemma4 KV width must contain complete storage heads") + options = get_full_attention_options(config, args) + options.update(heads=per_token_k_width // head_dim, dim=head_dim) + return options def _init_att_backend(self): # Gemma-4 has per-layer heterogeneous attention: sliding layers use diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index b31f79e502..2e65a13d79 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -5,8 +5,6 @@ from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num class Gemma_2bTpPartModel(LlamaTpPartModel): @@ -36,13 +34,10 @@ def _verify_params(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 return - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"], - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + @classmethod + def _full_attention_options(cls, config, args): + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options + + options = get_full_attention_options(config, args) + options.update(dim=config["hidden_size"] // config["num_attention_heads"], replicated=True) + return options diff --git a/lightllm/models/glm4_moe_lite_mtp/model.py b/lightllm/models/glm4_moe_lite_mtp/model.py index 2e4ba5c86b..245af80505 100644 --- a/lightllm/models/glm4_moe_lite_mtp/model.py +++ b/lightllm/models/glm4_moe_lite_mtp/model.py @@ -24,6 +24,15 @@ def _pre_init(self, kvargs: dict): self.main_model: TpPartBaseModel = kvargs.pop("main_model") self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models") + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + count = config.get("num_nextn_predict_layers", 1) + config.update(num_hidden_layers=count, n_layer=count) + return config + def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached @@ -31,9 +40,6 @@ def _init_custom(self): def _init_req_manager(self): self.req_manager = self.main_model.req_manager - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _init_weights(self, start_layer_index=None): assert start_layer_index is None diff --git a/lightllm/models/glm5_next/att_resource_config.py b/lightllm/models/glm5_next/att_resource_config.py new file mode 100644 index 0000000000..571ccdd084 --- /dev/null +++ b/lightllm/models/glm5_next/att_resource_config.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import dataclasses + +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +@dataclasses.dataclass +class Glm5NextResourceConfig(Qwen3NextResourceConfig): + """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + + # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + index_kpool: int = 4 + # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 + index_head_dim: int = 128 + + # 当前 packed MLA 行除 latent 外预留的字节数，不是 dtype 元素数。 + # 包含 128 字节量化 indexer key、4 字节 scale，并向上补齐到 16 字节边界。 + INDEX_PADDING_BYTES = 144 + + @classmethod + def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict, *, include_draft: bool = True): + config = config.get("text_config", config) + linear = config["linear_attn_config"] + tp = args.tp // args.dp + assert linear["num_heads"] % tp == 0 + layers = config["num_hidden_layers"] + assert config["layer_types"] == [ + "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) + ] + dtype = get_torch_dtype(args.data_type) + # Raw index keys and compression scores only live in request tails. + packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize + return cls( + tp_world_size=tp, + full_att_all_num_kv_heads=1, + full_att_dtype=dtype, + full_att_num_kv_heads=1, + full_att_head_dim=packed_dim, + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // tp, + num_linear_v_heads=linear["num_heads"] // tp, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=dtype, + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=4, + all_layer_num=layers, + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() + if include_draft and args.mtp_mode is not None + else 0, + index_kpool=config["index_kpool"], + index_head_dim=config["index_head_dim"], + ) diff --git a/lightllm/models/glm5_next/att_resource_manager.py b/lightllm/models/glm5_next/att_resource_manager.py new file mode 100644 index 0000000000..e9b57b6dd0 --- /dev/null +++ b/lightllm/models/glm5_next/att_resource_manager.py @@ -0,0 +1,160 @@ +"""Glm5Next attention resources with explicit model-owned State objects.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig +from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + +if TYPE_CHECKING: + import torch + from collections.abc import Sequence + from lightllm.common.att_resources.allocators.full_att_kv_allocator import FullAttKVAllocator + from lightllm.common.att_resources.full_att.base_state import FullAttentionState + from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + + +def init_model_att_tail(model: Glm5NextTpPartModel) -> None: + """按当前模型自己的 full-attention 层创建 tail，直接加入共享管理器。""" + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options, get_resource_config + + config = get_resource_config(model.config, model_class=type(model)) + options = get_full_attention_options( + config, model.args, model_class=type(model), + model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + ) + if not options["layer_ids"]: + return + cfg = model.linear_config + state = create_kpool_state( + config=KPoolStateConfig( + layer_ids=options["layer_ids"], pool_size=cfg.index_kpool, + head_dim=cfg.index_head_dim, dtype=cfg.full_att_dtype, + tp_rank=model.tp_rank_, tp_world_size=model.tp_world_size_, + ), + request_capacity=model.max_req_num + 1, mtp_step=model.args.mtp_step, device="cuda", + ) + model.att_resource_manager.set_att_state(model.att_resource_role, state) + + +class Glm5NextAttResourceManager(AttResourceManager): + # ======================================================================== + # 一、模型资源：创建 State 与绑定存储 + # ======================================================================== + + def __init__( + self, + cache_config: Glm5NextResourceConfig, + request_capacity: int, + mtp_step: int, + *, + tp_rank: int = 0, + device: str | torch.device | int = "cuda", + role: str = "main", + linear_layer_ids: Sequence[int] | None = None, + full_layer_ids: Sequence[int] | None = None, + ): + super().__init__(request_capacity=request_capacity) + self.cache_config = cache_config + self.cpu_tail_restore = True + layers = ( + tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) + if linear_layer_ids is None + else tuple(linear_layer_ids) + ) + config = ConvSSMStateConfig( + layer_ids=layers, + global_k_heads=cache_config.global_linear_k_heads, + global_v_heads=cache_config.global_linear_v_heads, + key_dim=cache_config.head_linear_k_dim, + value_dim=cache_config.head_linear_v_dim, + conv_width=cache_config.conv_kernel_size - 1, + conv_dtype=cache_config.conv_state_dtype, + ssm_dtype=cache_config.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=cache_config.tp_world_size, + ) + if layers: + state = create_kda_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) + self.set_att_state(role, state) + self.att_tail_type = f"indexer_kpool_c{cache_config.index_kpool}_tail" + own_full_layers = ( + tuple( + range( + cache_config.full_attention_interval - 1, + cache_config.all_layer_num, + cache_config.full_attention_interval, + ) + ) + if full_layer_ids is None + else tuple(full_layer_ids) + ) + tail_layers = [(role, own_full_layers)] + if role == "main" and cache_config.draft_full_att_kv_layer_num: + tail_layers.append( + ( + "draft0", + tuple( + range( + cache_config.all_layer_num, + cache_config.all_layer_num + cache_config.draft_full_att_kv_layer_num, + ) + ), + ) + ) + for role, layers in tail_layers: + if layers: + state = create_kpool_state( + config=KPoolStateConfig( + layer_ids=layers, + pool_size=cache_config.index_kpool, + head_dim=cache_config.index_head_dim, + dtype=cache_config.full_att_dtype, + tp_rank=tp_rank, + tp_world_size=cache_config.tp_world_size, + ), + request_capacity=request_capacity, mtp_step=mtp_step, device=device, + ) + self.set_att_state(role, state) + + def bind_storage( + self, + full_state: FullAttentionState, + full_allocator: FullAttKVAllocator, + *, + big_page_capacity: int, + small_page_capacity: int, + draft_full_state: FullAttentionState | None = None, + ): + """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" + assert self.full_att_allocator is None, "attention storage is already bound" + assert all( + layer < self.cache_config.all_layer_num for layer in full_state.config.layer_ids + ), "main State must only contain main layers" + assert (draft_full_state is not None) == bool( + self.cache_config.draft_full_att_kv_layer_num + ), "draft State must match the model configuration" + if draft_full_state is not None: + assert all( + layer >= self.cache_config.all_layer_num for layer in draft_full_state.config.layer_ids + ), "draft State must only contain draft layers" + + state = self.main_kda_state + state.initialize_page_pools( + big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, + pin_memory=state.conv_state.is_cuda, + ) + big_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) if big_page_capacity else None + small_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None + self.set_att_state("main", full_state) + if draft_full_state is not None: + self.set_att_state("draft0", draft_full_state) + self.full_att_allocator = full_allocator + self.bounded_att_big_page_allocator = big_allocator + self.bounded_att_small_page_allocator = small_allocator + self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 if big_allocator else None + self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 if big_allocator else None + self.finish_state_creation() diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..8f4cddb2d4 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -16,10 +16,11 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.heads = network_config["index_n_heads"] self.dim = network_config["index_head_dim"] self.eps = network_config["rms_norm_eps"] + self.pool_size = network_config["index_kpool"] def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" - if HAS_VLLM: + if HAS_VLLM and hasattr(vllm_ops, "top_k_per_row_decode"): # next_n=1 treats each query as an independent row, including prefill. # The decode entry splits long rows before merging their candidates. # Tested with vLLM 0.22.1; persistent_topk can drop candidates (#51782). @@ -36,8 +37,12 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.eps) gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) raw = torch.cat((k, gate), -1) - tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) - packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + role = infer_state.att_resource_role + tail = getattr( + infer_state.att_resource_manager, f"{role}_indexer_kpool_c{self.pool_size}_tail_state" + ).get_layer_view(self.layer_idx) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + packed_buffer = state.get_layer_view(self.layer_idx).indexer_k compress_pools( raw=raw, tail=tail, diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 620554bc06..b741d0ffc2 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -159,9 +159,10 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): layer_weight=layer_weight, ) del infer_state.get_topk_indices_params + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) return infer_state.prefill_att_state.prefill_att( q=q, - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_prefill=True, @@ -240,9 +241,10 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) return infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_decode=True, diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 6c5ffa6280..2fc9559e65 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -8,10 +8,7 @@ from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer @@ -21,12 +18,40 @@ class Glm5NextTpPartModel(TpPartBaseModel): + is_hybrid_att_model = True pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight transformer_weight_class = Glm5NextTransformerLayerWeight pre_layer_infer_class = Glm5NextPreLayerInfer post_layer_infer_class = LlamaPostLayerInfer transformer_layer_infer_class = Glm5NextTransformerLayerInfer + @classmethod + def create_bounded_attention_definition(cls, model_config, args, *, model_layer_ids, tp_rank=0, tp_world_size=1): + from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config, model_class=cls) + cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) + layers = tuple(layer for i, layer in enumerate(model_layer_ids) if (i + 1) % cfg.full_attention_interval) + if not layers: + return None + state = create_kda_state( + config=ConvSSMStateConfig( + layer_ids=layers, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=1, device="meta", + ) + return state + def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: outer_config = json.load(f) @@ -42,12 +67,6 @@ def _verify_params(self): super()._verify_params() assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" - if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): - raise ValueError("GLM-5.3 Flash supports only --llm_kv_type None or fp8kv_dsa") - if self.args.llm_kv_type == "fp8kv_dsa" and ( - self.config["kv_lora_rank"] != 512 or self.data_type != torch.bfloat16 - ): - raise ValueError("GLM-5.3 Flash fp8kv_dsa requires kv_lora_rank=512 and --data_type bfloat16") def autotune_layers(self): return 4 @@ -65,32 +84,33 @@ def _init_some_value(self): self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - def _init_req_manager(self): - self.linear_config = Glm5NextCacheConfig.from_model_config(self.config, self.args) - self.req_manager = Glm5NextReqManager( - self.max_req_num, - max(self.batch_max_tokens or 0, self.max_seq_length or 0), - None, - linear_config=self.linear_config, - ) + def _init_att_resource_manager(self): + from .att_resource_manager import Glm5NextAttResourceManager, init_model_att_tail + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention, init_model_bounded_attention - def _init_mem_manager(self): - mem_manager_class = FP8Glm5NextMemManager if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextMemManager - self.mem_manager = mem_manager_class( - size=self.max_total_token_num, - dtype=self.linear_config.full_att_dtype, - num_kv_heads=1, - head_dim=self.linear_config.full_att_head_dim, - full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - linear_config=self.linear_config, - mem_fraction=self.mem_fraction, - ) + self.linear_config = Glm5NextResourceConfig.from_model_config(self.config, self.args, include_draft=False) + if self.args.linear_att_hash_page_size % self.linear_config.index_kpool: + raise ValueError("GLM cache pages must align to index_kpool") + if self.args.mtp_step and self.args.enable_prefill_decode_mixed: + raise ValueError("MTP linear attention does not support mixed prefill/decode") + if not self.is_mtp_draft_model: + self.att_resource_manager = Glm5NextAttResourceManager( + self.linear_config, self.max_req_num + 1, self.args.mtp_step, + tp_rank=self.tp_rank_, linear_layer_ids=(), full_layer_ids=(), + ) + init_model_bounded_attention(self) + init_model_att_tail(self) + init_model_full_attention(self) + + @classmethod + def _full_attention_options(cls, config, args): + if args.llm_kv_type != "None": + raise ValueError(f"GLM K-pool does not support llm_kv_type={args.llm_kv_type}") + cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) + return dict(heads=1, dim=config["kv_lora_rank"], layer_ids=cfg.get_full_att_layer_ids(), packed=True) def _init_att_backend(self): - backend_class = ( - Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend - ) - self.prefill_att_backend = backend_class(model=self) + self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index c82962fda7..43b24ab354 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -61,10 +61,26 @@ def _init_some_value(self): def _init_req_manager(self): self.req_manager = self.main_model.req_manager + + def _init_att_resource_manager(self): + from lightllm.models.glm5_next.att_resource_manager import init_model_att_tail + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + self.linear_config = self.main_model.linear_config + init_model_att_tail(self) + init_model_full_attention(self) + + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + config.update(num_hidden_layers=1, n_layer=1) + return config - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager + @classmethod + def _full_attention_options(cls, config, args): + return dict(heads=1, dim=config["kv_lora_rank"], layer_ids=(0,), packed=True) def _init_att_backend(self): backend_class = ( diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index 490d2dc4c5..1cd117f336 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -72,7 +72,8 @@ def _context_attention_kernel( window_size = (-1, -1) use_sliding_window = False - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -99,7 +100,8 @@ def _token_attention_kernel( window_size = (-1, -1) use_sliding_window = False - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 69acffaa4d..12610ede89 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -44,7 +44,8 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -61,7 +62,8 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + _k, _v = state.get_layer_view(layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor) return o_tensor.view(q.shape) diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 2eab67f807..1aa3037304 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -9,8 +9,6 @@ from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args @@ -34,6 +32,14 @@ def __init__(self, kvargs): super().__init__(kvargs) return + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _init_config(self): super()._init_config() # rename key @@ -52,19 +58,6 @@ def _verify_params(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 return - def _init_mem_manager(self): - head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim_ = self.config.get("head_dim", head_dim_) - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=head_dim_, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_custom(self): """ 模型特殊的一些初始化 diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 32b8c538d3..d4db886ae2 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -6,8 +6,6 @@ from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.mistral.layer_infer.transformer_layer_infer import MistralTransformerLayerInfer -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.envs_utils import get_env_start_args @@ -28,6 +26,14 @@ def __init__(self, kvargs): super().__init__(kvargs) return + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _verify_params(self): assert self.load_way in ["HF"], "mistral only supports HF format to load Now!" assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0 @@ -38,20 +44,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B - head_dim = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim = self.config.get("head_dim", head_dim) - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=head_dim, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_to_get_rotary(self, default_base=10000): if self.config.get("rope_scaling", {}) is None: rope_scaling_factor = 1.0 diff --git a/lightllm/models/mistral_mtp/model.py b/lightllm/models/mistral_mtp/model.py index f17bc0a383..d85b60754e 100644 --- a/lightllm/models/mistral_mtp/model.py +++ b/lightllm/models/mistral_mtp/model.py @@ -45,10 +45,6 @@ def _init_req_manager(self): self.req_manager = self.main_model.req_manager return - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.config["n_layer"] = 1 diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 8eaeec4018..05cd662a03 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -1,7 +1,6 @@ import torch import numpy as np from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight @@ -35,6 +34,14 @@ def _init_config(self): super()._init_config() return + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _verify_params(self): assert self.load_way in ["HF"], "mixtral only supports HF format to load Now!" assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0 @@ -45,18 +52,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - self.mem_manager = MemoryManager( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"], - always_copy=False, - mem_fraction=self.mem_fraction, - ) - return - def _init_to_get_rotary(self, default_base=10000): if self.config.get("rope_scaling", {}) is None: rope_scaling_factor = 1.0 diff --git a/lightllm/models/qwen2/model.py b/lightllm/models/qwen2/model.py index 82f16bc2f1..250cbbe461 100644 --- a/lightllm/models/qwen2/model.py +++ b/lightllm/models/qwen2/model.py @@ -1,8 +1,6 @@ from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight from lightllm.models.qwen2.layer_weights.transformer_layer_weight import Qwen2TransformerLayerWeight from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num class Qwen2TpPartModel(LlamaTpPartModel): @@ -35,18 +33,3 @@ def _init_some_value(self): self.layers_num = self.config["n_layer"] self.vocab_size = self.config["vocab_size"] return - - def _init_mem_manager(self): - head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim_ = self.config.get("head_dim", head_dim_) - tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=tp_k_head_num_, - head_dim=head_dim_, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 422ccf61cf..7a0d917f49 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -10,6 +10,15 @@ class Qwen3_5DFlashModel(Qwen3DFlashModel): pre_and_post_weight_class = Qwen35DFlashPreAndPostLayerWeight + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + config.update(config.get("dflash_config", {})) + config["n_layer"] = config["num_hidden_layers"] + return config + def _init_config(self): super()._init_config() self.config.update(self.config.get("dflash_config", {})) @@ -28,19 +37,6 @@ def _init_custom(self): self.block_size = self.config["block_size"] self.mask_token_id = self.config["mask_token_id"] - def _init_mem_manager(self): - target_mem_manager = self.main_model.mem_manager - draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) - target_kv_shape = ( - target_mem_manager.linear_config.full_att_all_num_kv_heads, - target_mem_manager.head_dim, - ) - assert draft_kv_shape == target_kv_shape, ( - "Qwen3.5 parallel block drafter requires matching draft and target KV shapes, " - f"got draft={draft_kv_shape}, target={target_kv_shape}." - ) - super()._init_mem_manager() - def _init_weights(self, start_layer_index=None): super()._init_weights(start_layer_index=start_layer_index) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 896949d1cf..841653e30b 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -5,6 +5,15 @@ class Qwen3_5DSparkModel(Qwen3DSparkModel): """Adapter for the current Qwen3 DSpark checkpoint with a Qwen3.5 target.""" + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + config.update(config.get("dflash_config", {})) + config["n_layer"] = config["num_hidden_layers"] + return config + def _init_config(self): super()._init_config() self.config.update(self.config.get("dflash_config", {})) @@ -24,16 +33,3 @@ def _init_custom(self): LlamaTpPartModel._init_custom(self) self.block_size = self.config["block_size"] self.mask_token_id = self.config["mask_token_id"] - - def _init_mem_manager(self): - target_mem_manager = self.main_model.mem_manager - draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) - target_kv_shape = ( - target_mem_manager.linear_config.full_att_all_num_kv_heads, - target_mem_manager.head_dim, - ) - assert draft_kv_shape == target_kv_shape, ( - "Qwen3.5 parallel block drafter requires matching draft and target KV shapes, " - f"got draft={draft_kv_shape}, target={target_kv_shape}." - ) - super()._init_mem_manager() diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py index b0f55af7f3..61c920f23b 100644 --- a/lightllm/models/qwen3_5_mtp/model.py +++ b/lightllm/models/qwen3_5_mtp/model.py @@ -14,7 +14,7 @@ class Qwen3_5MTPModel(Qwen3_5TpPartModel): transformer_weight_class = Qwen3_5MTPTransformerLayerWeight transformer_layer_infer_class = Qwen35TransformerLayerInfer - # MTP draft model: reuses the main model's req/mem managers and rope caches, and is + # MTP draft model: reuses the main model's request/attention resources and rope caches, and is # marked so the decode CUDA-graph / padding paths detect it (is_mtp_draft_model). is_mtp_draft_model = True @@ -29,14 +29,18 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1) + return config + def _init_req_manager(self): self.req_manager = self.main_model.req_manager return - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_config(self): super()._init_config() # MTP draft model: reuses the main model's config, but overrides the following: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 9c4a42fffc..61ad02d9c4 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -49,10 +49,6 @@ def _init_custom(self): def _init_req_manager(self): self.req_manager = self.main_model.req_manager - def _init_mem_manager(self): - # Draft KV uses target-owned cache slots. - self.mem_manager = self.main_model.mem_manager - def _init_att_backend(self): super()._init_att_backend() # FA3 is currently the only backend that supports non-causal block decode. @@ -98,7 +94,8 @@ def _decode(self, model_input: ModelInput) -> ModelOutput: infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens infer_state.position_cos = torch.index_select(self._cos_cached, 0, position_ids) infer_state.position_sin = torch.index_select(self._sin_cached, 0, position_ids) - infer_state.mem_manager = self.mem_manager + infer_state.att_resource_manager = self.att_resource_manager + infer_state.is_draft_model = True infer_state.mem_index = self._select_mem_indexes(model_input) # 这里调用的是 DFlash 覆写后的 context_forward：pre layer 只投影 target hidden， diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b03fd77cc2..b80ebb564f 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -13,7 +13,8 @@ def __init__(self, layer_num, network_config): self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) def _token_attention_kernel(self, q, infer_state, layer_weight): - k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + k, v = state.get_layer_view(layer_id=self.layer_num_) k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) q = q.view(-1, self.tp_q_head_num_, self.head_dim_) @@ -34,8 +35,9 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): - storage_head_num = infer_state.mem_manager.head_num - storage_head_dim = infer_state.mem_manager.head_dim + state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + storage_head_num = state.config.head_num + storage_head_dim = state.config.head_dim storage_kv = cache_kv.view(cache_kv.shape[0], 2 * storage_head_num, storage_head_dim) return super()._post_cache_kv(storage_kv, infer_state, layer_weight) diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py index 3542b37188..e4485995b1 100644 --- a/lightllm/models/qwen3_dflash2/model.py +++ b/lightllm/models/qwen3_dflash2/model.py @@ -18,6 +18,15 @@ class Qwen3DFlash2Model(Qwen3DFlashModel): post_layer_infer_class = Qwen3DFlash2PostLayerInfer transformer_layer_infer_class = Qwen3DFlash2TransformerLayerInfer + @classmethod + def _resource_config(cls, model_config): + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config) + config.update(config.get("dflash_config", {})) + config["n_layer"] = config["num_hidden_layers"] + return config + def _init_config(self): super()._init_config() dflash_config = self.config.get("dflash_config", {}) @@ -61,18 +70,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_mem_manager(self): - main_mem_manager = self.main_model.mem_manager - draft_head_num = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - draft_width = draft_head_num * self.config["head_dim"] - target_width = main_mem_manager.head_num * main_mem_manager.head_dim - assert draft_width == target_width, ( - "DFlash2 draft and target KV must have equal flat widths: " - f"draft=({draft_head_num}, {self.config['head_dim']}), " - f"target=({main_mem_manager.head_num}, {main_mem_manager.head_dim})" - ) - self.mem_manager = main_mem_manager - def _init_weights(self, start_layer_index=None): super()._init_weights(start_layer_index=start_layer_index) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ diff --git a/lightllm/models/qwen3_eagle/model.py b/lightllm/models/qwen3_eagle/model.py index d6257d3ebf..490d06fba3 100644 --- a/lightllm/models/qwen3_eagle/model.py +++ b/lightllm/models/qwen3_eagle/model.py @@ -49,9 +49,6 @@ def _init_custom(self): def _init_req_manager(self): self.req_manager = self.main_model.req_manager - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _init_weights(self, start_layer_index=None): self.pre_post_weight = self.pre_and_post_weight_class( self.data_type, network_config=self.config, quant_cfg=self.quant_cfg diff --git a/lightllm/models/qwen3_moe_mtp/model.py b/lightllm/models/qwen3_moe_mtp/model.py index d9854250e2..e5d5f06781 100644 --- a/lightllm/models/qwen3_moe_mtp/model.py +++ b/lightllm/models/qwen3_moe_mtp/model.py @@ -37,10 +37,6 @@ def _init_req_manager(self): self.req_manager = self.main_model.req_manager return - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None mtp_index = len(self.mtp_previous_draft_models) diff --git a/lightllm/models/qwen3next/__init__.py b/lightllm/models/qwen3next/__init__.py index a9d22c6643..482161743d 100644 --- a/lightllm/models/qwen3next/__init__.py +++ b/lightllm/models/qwen3next/__init__.py @@ -1,3 +1,9 @@ -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel - __all__ = ["Qwen3NextTpPartModel"] + + +def __getattr__(name): + if name == "Qwen3NextTpPartModel": + from .model import Qwen3NextTpPartModel + + return Qwen3NextTpPartModel + raise AttributeError(name) diff --git a/lightllm/models/qwen3next/att_resource_config.py b/lightllm/models/qwen3next/att_resource_config.py new file mode 100644 index 0000000000..3c6f45cc40 --- /dev/null +++ b/lightllm/models/qwen3next/att_resource_config.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import torch +import dataclasses + +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +@dataclasses.dataclass +class Qwen3NextResourceConfig: + # 单个模型的 TP 组大小，通常为 args.tp // args.dp，不包含其他 DP 组。 + tp_world_size: int + + # TP 切分前的全局 KV head 数，K/V 各有这么多个，不相加；GLM 的 MLA latent 记为 1。 + full_att_all_num_kv_heads: int + # full-attention 布局使用的 dtype；普通 KV 为浮点精度，静态 FP8 为编码容器 uint8。 + # GLM 还用此浮点精度保存原始 indexer tail，具体量化字段由各 State 定义。 + full_att_dtype: torch.dtype + # 当前 TP rank 实际保存的 KV head 数；少于 TP 数时按组复制，MLA 为每个 rank 保存 1。 + full_att_num_kv_heads: int + # 普通 GQA 中为单个 K/V head 的维度；GLM packed MLA 中为包含 indexer/padding 的行元素数。 + full_att_head_dim: int + + # TP 切分前的线性 attention K head 数；Q 使用相同的 head 数和维度。 + global_linear_k_heads: int + # TP 切分前的 V head 数；GDN 可与 K 不同，GLM KDA 与 K 相同。 + global_linear_v_heads: int + # 当前 TP rank 的线性 K/Q head 数；Conv 的 Q/K 通道宽度由此与 head_linear_k_dim 相乘。 + num_linear_k_heads: int + # 当前 TP rank 的线性 V head 数，也是每层 SSM 矩阵的 head 轴长度。 + num_linear_v_heads: int + # 每个 Q/K head 的特征维度，对应 SSM 的 key 轴，不随 TP 切分。 + head_linear_k_dim: int + # 每个 V head 的特征维度，对应 SSM 的 value 轴；KDA 中与 key 维度相等。 + head_linear_v_dim: int + # 原始卷积核宽度，单位为 token；基准 Conv 历史保存 kernel_size - 1 个位置。 + # MTP 的额外候选位置由 State.mtp_step 决定，不计入本字段。 + conv_kernel_size: int + # 主模型中线性 attention 层的数量，不含 full-attention 层或追加的 draft 层。 + linear_layer_num: int + # Conv 运行态和大小页快照的元素 dtype，不代表算子的内部累加精度。 + conv_state_dtype: torch.dtype + # SSM 运行态和大小页快照的元素 dtype，可独立于 Conv 精度配置。 + ssm_state_dtype: torch.dtype + # 混合层排列周期：每隔这么多层放一个 full-attention 层，零起始层号为 interval - 1。 + full_attention_interval: int + # 主模型总层数，包含线性与 full-attention 两类层，不包含追加的 draft 层。 + all_layer_num: int + # 需要追加到 full-attention 层布局中的 draft KV 层数；不包含 draft 的布局取 0。 + # 用于合并布局的层号/物理行换算，不代表 draft 请求数或 MTP 候选 token 数。 + draft_full_att_kv_layer_num: int = 0 + + def get_conv_dim(self): + # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 + # 由于 k_dim = q_dim, k_heads = q_heads, 所以第一项和第二项的计算 + # 形式相同，但是实际内在含义是不同的。 + return ( + self.head_linear_k_dim * self.num_linear_k_heads + + self.head_linear_k_dim * self.num_linear_k_heads + + self.head_linear_v_dim * self.num_linear_v_heads + ) + + def get_main_model_full_att_layer_num(self): + full_att_layer_num = self.all_layer_num - self.linear_layer_num + assert full_att_layer_num == self.all_layer_num // self.full_attention_interval + return full_att_layer_num + + def get_full_att_kv_layer_num_with_draft_model(self): + return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num + + def get_full_att_layer_ids(self): + """Logical model layers stored in each physical full-attention row.""" + return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) + tuple( + range(self.all_layer_num, self.all_layer_num + self.draft_full_att_kv_layer_num) + ) + + def get_full_att_kv_layer_index(self, layer_index: int) -> int: + """Map a global target/draft layer index to the packed full-attention cache.""" + + layer_index = int(layer_index) + if layer_index >= self.all_layer_num: + kv_layer_index = layer_index - self.linear_layer_num + else: + kv_layer_index = layer_index // self.full_attention_interval + assert 0 <= kv_layer_index < self.get_full_att_kv_layer_num_with_draft_model(), ( + f"layer {layer_index} maps outside the packed full-attention KV cache: " + f"slot={kv_layer_index}, slots={self.get_full_att_kv_layer_num_with_draft_model()}" + ) + return kv_layer_index + + def get_conv_state_shape(self): + # Base committed sliding-window state, without speculative MTP tail. + return (self.get_conv_dim(), self.conv_kernel_size - 1) + + def get_mtp_conv_state_shape(self, mtp_step: int): + # Working state with room for S speculative tokens before acceptance. + return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step) + + def get_ssm_state_shape(self): + return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim) + + def get_conv_state_bytes_per_layer(self): + return self.get_conv_dim() * (self.conv_kernel_size - 1) * self.conv_state_dtype.itemsize + + def get_ssm_state_bytes_per_layer(self): + return self.num_linear_v_heads * self.head_linear_k_dim * self.head_linear_v_dim * self.ssm_state_dtype.itemsize + + @classmethod + def load_from_args(cls) -> "Qwen3NextResourceConfig": + args = get_env_start_args() + model_path = args.model_dir + from transformers.configuration_utils import PretrainedConfig + + model_cfg, _ = PretrainedConfig.get_config_dict(model_path) + return cls.from_model_config(model_cfg, args) + + @classmethod + def from_model_config( + cls, + model_cfg: dict[str, object], + args: StartArgs | EasyDict, + *, + include_draft: bool = True, + ): + llm_config = model_cfg.get("text_config", model_cfg) + n_layer = llm_config["num_hidden_layers"] + + tp_world_size = args.tp // args.dp + full_att_dtype = ( + torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) + ) + return cls( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=llm_config["num_key_value_heads"], + full_att_dtype=full_att_dtype, + full_att_num_kv_heads=max(1, llm_config["num_key_value_heads"] // tp_world_size), + full_att_head_dim=llm_config["head_dim"], + global_linear_k_heads=llm_config["linear_num_key_heads"], + global_linear_v_heads=llm_config["linear_num_value_heads"], + num_linear_k_heads=max(1, llm_config["linear_num_key_heads"] // tp_world_size), + num_linear_v_heads=max(1, llm_config["linear_num_value_heads"] // tp_world_size), + head_linear_k_dim=llm_config["linear_key_head_dim"], + head_linear_v_dim=llm_config["linear_value_head_dim"], + conv_kernel_size=llm_config["linear_conv_kernel_dim"], + linear_layer_num=n_layer - (n_layer // llm_config["full_attention_interval"]), + conv_state_dtype=get_torch_dtype(args.data_type), + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=llm_config["full_attention_interval"], + all_layer_num=n_layer, + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if include_draft else 0, + ) diff --git a/lightllm/models/qwen3next/att_resource_manager.py b/lightllm/models/qwen3next/att_resource_manager.py new file mode 100644 index 0000000000..d9fdefe593 --- /dev/null +++ b/lightllm/models/qwen3next/att_resource_manager.py @@ -0,0 +1,94 @@ +"""Qwen3Next attention resources with explicit model-owned State objects.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig + +if TYPE_CHECKING: + import torch + from collections.abc import Sequence + from lightllm.common.att_resources.allocators.full_att_kv_allocator import FullAttKVAllocator + from lightllm.common.att_resources.full_att.base_state import FullAttentionState + from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig + + +class Qwen3NextAttResourceManager(AttResourceManager): + # ======================================================================== + # 一、模型资源：创建 State 与绑定存储 + # ======================================================================== + + def __init__( + self, + cache_config: Qwen3NextResourceConfig, + request_capacity: int, + mtp_step: int, + *, + tp_rank: int = 0, + device: str | torch.device | int = "cuda", + role: str = "main", + layer_ids: Sequence[int] | None = None, + ): + super().__init__(request_capacity=request_capacity) + self.cache_config = cache_config + layers = ( + tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) + if layer_ids is None + else tuple(layer_ids) + ) + config = ConvSSMStateConfig( + layer_ids=layers, + global_k_heads=cache_config.global_linear_k_heads, + global_v_heads=cache_config.global_linear_v_heads, + key_dim=cache_config.head_linear_k_dim, + value_dim=cache_config.head_linear_v_dim, + conv_width=cache_config.conv_kernel_size - 1, + conv_dtype=cache_config.conv_state_dtype, + ssm_dtype=cache_config.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=cache_config.tp_world_size, + ) + if layers: + state = create_gdn_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) + self.set_att_state(role, state) + + def bind_storage( + self, + full_state: FullAttentionState, + full_allocator: FullAttKVAllocator, + *, + big_page_capacity: int, + small_page_capacity: int, + draft_full_state: FullAttentionState | None = None, + ): + """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" + assert self.full_att_allocator is None, "attention storage is already bound" + assert all( + layer < self.cache_config.all_layer_num for layer in full_state.config.layer_ids + ), "main State must only contain main layers" + assert (draft_full_state is not None) == bool( + self.cache_config.draft_full_att_kv_layer_num + ), "draft State must match the model configuration" + if draft_full_state is not None: + assert all( + layer >= self.cache_config.all_layer_num for layer in draft_full_state.config.layer_ids + ), "draft State must only contain draft layers" + + state = self.main_gdn_state + state.initialize_page_pools( + big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, + pin_memory=state.conv_state.is_cuda, + ) + big_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) if big_page_capacity else None + small_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None + self.set_att_state("main", full_state) + if draft_full_state is not None: + self.set_att_state("draft0", draft_full_state) + self.full_att_allocator = full_allocator + self.bounded_att_big_page_allocator = big_allocator + self.bounded_att_small_page_allocator = small_allocator + self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 if big_allocator else None + self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 if big_allocator else None + self.finish_state_creation() diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..39f9c42919 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -8,7 +8,6 @@ from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.common.kv_cache_mem_manager import Qwen3NextMemManager from lightllm.common.basemodel.attention.base_att import AttControl from typing import Tuple from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ @@ -245,7 +244,6 @@ def context_attention_forward( if not self.is_linear_attention_layer: return super().context_attention_forward(input_embdings, infer_state, layer_weight) - assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) if torch.cuda.is_current_stream_capturing(): @@ -281,7 +279,6 @@ def token_attention_forward( if not self.is_linear_attention_layer: return super().token_attention_forward(input_embdings, infer_state, layer_weight) - assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) core_attn_out, z = infer_state.decode_att_state1.decode_att( q=None, diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 444ed1f393..5bbc5ffc6f 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -12,19 +12,15 @@ from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( - FP8StaticPerHeadQuantQwen3NextMemManager, - FP8StaticPerTensorQuantQwen3NextMemManager, - Qwen3NextMemManager, -) from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.common.req_manager import ReqManagerForMamba -from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig logger = init_logger(__name__) class Qwen3NextTpPartModel(Qwen3MOEModel): + is_hybrid_att_model = True + # weight class pre_and_post_weight_class = Qwen3NextPreAndPostLayerWeight transformer_weight_class = Qwen3NextTransformerLayerWeight @@ -39,6 +35,33 @@ def __init__(self, kvargs) -> None: self._init_triton() super().__init__(kvargs) + @classmethod + def create_bounded_attention_definition(cls, model_config, args, *, model_layer_ids, tp_rank=0, tp_world_size=1): + from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config, model_class=cls) + cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) + layers = tuple(layer for i, layer in enumerate(model_layer_ids) if (i + 1) % cfg.full_attention_interval) + if not layers: + return None + state = create_gdn_state( + config=ConvSSMStateConfig( + layer_ids=layers, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=1, device="meta", + ) + return state + def _init_triton(self): def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: return torch.empty(size, device="cuda", dtype=torch.int8) @@ -56,45 +79,35 @@ def _init_config(self): super()._init_config() self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - def _init_mem_manager(self): + def _init_att_resource_manager(self): + from .att_resource_manager import Qwen3NextAttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention, init_model_bounded_attention + assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 start_args: StartArgs = get_env_start_args() - if start_args.llm_kv_type == "None": - mem_manager_class = Qwen3NextMemManager - elif start_args.llm_kv_type == "fp8kv_sph": - mem_manager_class = FP8StaticPerHeadQuantQwen3NextMemManager - elif start_args.llm_kv_type == "fp8kv_spt": - mem_manager_class = FP8StaticPerTensorQuantQwen3NextMemManager - else: - raise ValueError(f"Qwen3Next does not support llm_kv_type={start_args.llm_kv_type}") - self.mem_manager = mem_manager_class( - size=self.max_total_token_num, - dtype=self.data_type, - num_kv_heads=self.num_kv_heads, - head_dim=self.config["head_dim"], - full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - linear_config=self.linear_config, - mem_fraction=self.mem_fraction, + self.linear_config = Qwen3NextResourceConfig.from_model_config(self.config, start_args, include_draft=False) + + if start_args.mtp_step and start_args.enable_prefill_decode_mixed: + raise ValueError("MTP linear attention does not support mixed prefill/decode") + if not self.is_mtp_draft_model: + self.att_resource_manager = Qwen3NextAttResourceManager( + self.linear_config, self.max_req_num + 1, start_args.mtp_step, tp_rank=self.tp_rank_, layer_ids=() + ) + init_model_bounded_attention(self) + init_model_full_attention(self) + + @classmethod + def _full_attention_options(cls, config, args): + if args.llm_kv_type not in ("None", "fp8kv_sph", "fp8kv_spt"): + raise ValueError(f"Qwen3Next does not support llm_kv_type={args.llm_kv_type}") + cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) + return dict( + heads=cfg.full_att_all_num_kv_heads, dim=cfg.full_att_head_dim, layer_ids=cfg.get_full_att_layer_ids() ) - def _init_req_manager(self): - self.linear_config = LinearAttCacheConfig.load_from_args() - create_max_seq_len = 0 - - if self.batch_max_tokens is not None: - create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens) - if self.max_seq_length is not None: - create_max_seq_len = max(create_max_seq_len, self.max_seq_length) - - self.req_manager = ReqManagerForMamba( - self.max_req_num, create_max_seq_len, None, linear_config=self.linear_config - ) - return - def _init_att_backend1(self): - # MTP draft models only keep full-attention layers and reuse the main - # model's managers, so they do not need LinearAttBackend. - if getattr(self, "is_mtp_draft_model", False): + role = self.att_resource_role + if getattr(self.att_resource_manager, f"{role}_gdn_state", None) is None: self.prefill_att_backend1 = None self.decode_att_backend1 = None return diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index de7eeac78b..d874a5abad 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -4,8 +4,6 @@ from lightllm.models.starcoder.layer_weights.pre_and_post_layer_weight import StarcoderPreAndPostLayerWeight from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.common.build_utils import repair_config -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel import InferStateInfo @@ -24,6 +22,14 @@ class StarcoderTpPartModel(TpPartBaseModel): def __init__(self, kvargs): super().__init__(kvargs) + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _init_config(self): super()._init_config() # rename key @@ -39,16 +45,13 @@ def _reset_num_key_value_heads(self): def _verify_params(self): assert self.load_way == "HF", "StarCoder only support HF format to load Now!" - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"], - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + @classmethod + def _full_attention_options(cls, config, args): + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options + + options = get_full_attention_options(config, args) + options.update(heads=1, dim=config["hidden_size"] // config["num_attention_heads"], replicated=True) + return options def _init_some_value(self): super()._init_some_value() diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 09107103a4..801897d63d 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -7,8 +7,6 @@ from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.common.build_utils import repair_config -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.common.basemodel import TpPartBaseModel @@ -27,6 +25,14 @@ def __init__(self, kvargs): super().__init__(kvargs) return + def _init_att_resource_manager(self): + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + init_model_full_attention(self) + def _init_config(self): super()._init_config() repair_config(self.config, same_names=["norm_epsilon", "rms_norm_eps", "layer_norm_epsilon"]) @@ -45,17 +51,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_some_value(self): super()._init_some_value() return diff --git a/lightllm/server/multi_level_kv_cache/cache_placement_controller.py b/lightllm/server/multi_level_kv_cache/cache_placement_controller.py index 2bd67e4f93..43b3eea996 100644 --- a/lightllm/server/multi_level_kv_cache/cache_placement_controller.py +++ b/lightllm/server/multi_level_kv_cache/cache_placement_controller.py @@ -229,7 +229,7 @@ def create_cache_placement_controller( if args.enable_disk_cache: disk_capacity_bytes = int(args.disk_cache_storage_size * (1024 ** 3)) - disk_capacity_tokens = disk_capacity_bytes // cpu_cache_meta.calcu_one_page_size() * page_size_tokens + disk_capacity_tokens = disk_capacity_bytes // cpu_cache_meta.page_bytes * page_size_tokens else: disk_capacity_tokens = 0 diff --git a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py b/lightllm/server/multi_level_kv_cache/cpu_cache_client.py index fb9e74c47b..8138b798f6 100644 --- a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py +++ b/lightllm/server/multi_level_kv_cache/cpu_cache_client.py @@ -1,4 +1,5 @@ import ctypes +import torch from lightllm.utils.envs_utils import get_env_start_args, get_disk_cache_prompt_limit_length from typing import List, Optional, Tuple from lightllm.utils.log_utils import init_logger @@ -27,14 +28,11 @@ def __init__(self, only_create_meta_data: bool, init_shm_data: bool): tensor_spec = CpuCacheTensorSpec( shm_key=self.args.cpu_kv_cache_shm_id, shape=( - self.kv_cache_tensor_meta.page_num, - self.kv_cache_tensor_meta.layer_num, - self.kv_cache_tensor_meta.token_page_size, - self.kv_cache_tensor_meta.num_heads, - self.kv_cache_tensor_meta.get_merged_head_dim(), + self.page_num, + self.kv_cache_tensor_meta.page_bytes, ), - dtype=self.kv_cache_tensor_meta.data_type, - size_bytes=self.kv_cache_tensor_meta.calcu_size(), + dtype=torch.uint8, + size_bytes=self.kv_cache_tensor_meta.size_bytes, ) tensor_creator = CpuCacheCreator(tensor_spec=tensor_spec) self.cpu_kv_cache_tensor = tensor_creator.create_or_attach( diff --git a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py index 542ddbd877..fd9534dccb 100644 --- a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py +++ b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py @@ -54,7 +54,8 @@ def __init__( if not cache_dir: cache_dir = os.path.join(tempfile.gettempdir(), f"lightllm_disk_cache_{get_unique_server_name()}") os.makedirs(cache_dir, exist_ok=True) - cache_file = os.path.join(cache_dir, "cache_file") + namespace = cpu_cache_client.kv_cache_tensor_meta.cache_namespace + cache_file = os.path.join(cache_dir, f"cache_file_{namespace}" if namespace else "cache_file") self.max_concurrent_write_tasks = max_concurrent_write_tasks self._page_major_tensor = self._prepare_tensor(cpu_cache_client.cpu_kv_cache_tensor) diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 95f0e4daec..bf563ae791 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -192,8 +192,9 @@ class PDChunckedTransTask: transfer_time_out_secs: int = 66 # kv: 通过 mem_indexes 传输 [start_kv_index, end_kv_index) 的 token KV。 # att_state: 混合注意力模型的请求运行态 buffer（如 linear attention 的 conv/SSM 状态）。 - # start_kv_index == end_kv_index 标记状态对应的 token 位置，mem_indexes 为空。 - # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 mem_manager 负责。 + # start_kv_index == end_kv_index 标记 prefill 状态对应的 token 位置，mem_indexes 为空。 + # tail PD 读写用该位置在固定逻辑窗口与本地 MTP ring 之间转换，不另存长度字段。 + # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 AttResourceManager 负责。 page_kind: str = "kv" # Only valid for the local task owner; remote notify copies may carry the sender-local req_idx. req_idx: Optional[int] = None diff --git a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py index f8516bce2c..8f0a8c6a87 100644 --- a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py @@ -2,7 +2,6 @@ import numpy as np from typing import Tuple, Dict, Set, List, Optional, Sequence from sortedcontainers import SortedSet, SortedDict -from lightllm.common.state_cache_manager import StateCacheManager from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len from .shared_arr import SharedArray from .radix_cache import time_gen @@ -117,11 +116,8 @@ def __init__( rank_in_node: int, hash_page_size: int, big_page_num: int, - kv_cache_mem_manager=None, - small_page_buffers=None, + att_resource_manager, ): - from lightllm.common.kv_cache_mem_manager import MemoryManager - assert hash_page_size >= 1, "hash_page_size must be >= 1" assert big_page_num >= 1, "big_page_num must be >= 1" @@ -130,9 +126,7 @@ def __init__( self.big_page_tokens = hash_page_size * big_page_num self.total_token_num = total_token_num - self.mem_manager: MemoryManager = kv_cache_mem_manager - - self.big_page_buffers: StateCacheManager = self.mem_manager.big_page_buffers + self.att_resource_manager = att_resource_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 @@ -151,7 +145,15 @@ def __init__( self.refed_tokens_num.arr[0] = 0 self.tree_total_tokens_num = SharedArray(f"tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64) self.tree_total_tokens_num.arr[0] = 0 - self.small_page_buffers: StateCacheManager = small_page_buffers + + def _release_pages(self, page_kind, slot_ids): + resources = self.att_resource_manager + allocator = ( + resources.bounded_att_big_page_allocator + if page_kind == "big" + else resources.bounded_att_small_page_allocator + ) + allocator.free(slot_ids) def _discard_node(self, node: HybridAttPagedTreeNode): if node.is_leaf(): @@ -248,7 +250,7 @@ def _insert_helper( new_big_page_buffer_id = len_to_big_page_id.pop(child.node_prefix_total_len, None) if new_big_page_buffer_id is not None: # 因为节点已经存在，所以无法插入，但是要释放对应的buffer_id 节点 - self.big_page_buffers.free_state_cache([new_big_page_buffer_id]) + self._release_pages("big", [new_big_page_buffer_id]) # 已经存在了 sub_prefix_len, ans_node = self._insert_helper( @@ -303,7 +305,7 @@ def _insert_helper( self._add_node(child) else: # 节点已有 checkpoint，释放本次重复申请的槽位。 - self.small_page_buffers.free_state_cache(free_indexes=[block_state_idxs[0]]) + self._release_pages("small", slot_ids=[block_state_idxs[0]]) sub_prefix_len, ans_node = self._insert_helper( child, @@ -585,9 +587,8 @@ def _print_helper(self, node: HybridAttPagedTreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - assert self.mem_manager is not None - if need_token_num > self.mem_manager.allocator.can_use_mem_size: - need_evict_token_num = need_token_num - self.mem_manager.allocator.can_use_mem_size + if need_token_num > self.att_resource_manager.full_att_allocator.can_use_mem_size: + need_evict_token_num = need_token_num - self.att_resource_manager.full_att_allocator.can_use_mem_size release_mems = [] small_page_buffer_ids = [] @@ -598,16 +599,17 @@ def release_mem(mem_index, small_page_buffer_id): self._evict(need_evict_token_num, release_mem) mem_index = torch.concat(release_mems) - self.mem_manager.free(mem_index) + self.att_resource_manager.full_att_allocator.free(mem_index) small_page_buffer_ids = [idx for idx in small_page_buffer_ids if idx is not None] if len(small_page_buffer_ids) > 0: - self.small_page_buffers.free_state_cache(small_page_buffer_ids) + self._release_pages("small", small_page_buffer_ids) return def free_one_small_page_buffer(self): - if self.small_page_buffers is None: + allocator = self.att_resource_manager.bounded_att_small_page_allocator + if allocator is None: return - if self.small_page_buffers.get_free_cache_num() > 0: + if allocator.free_count > 0: return if len(self._evict_tree_set_for_state_cache) == 0: return @@ -616,7 +618,7 @@ def free_one_small_page_buffer(self): self._discard_node(node) assert node.small_page_buffer_idx is not None - self.small_page_buffers.free_state_cache(free_indexes=[node.small_page_buffer_idx]) + self._release_pages("small", slot_ids=[node.small_page_buffer_idx]) node.small_page_buffer_idx = None self._add_node(node) @@ -639,7 +641,7 @@ def _evict(self, need_remove_tokens, evict_callback): if node.is_big_page_node(): assert node.big_page_buffer_idx is not None - self.big_page_buffers.free_state_cache([node.big_page_buffer_idx]) + self._release_pages("big", [node.big_page_buffer_idx]) evict_callback(node.token_mem_index_value, node.small_page_buffer_idx) self.tree_total_tokens_num.arr[0] -= len(node.token_mem_index_value) diff --git a/lightllm/server/router/dynamic_prompt/radix_cache.py b/lightllm/server/router/dynamic_prompt/radix_cache.py index 25df6d5ef7..f0d4b774b3 100644 --- a/lightllm/server/router/dynamic_prompt/radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/radix_cache.py @@ -110,18 +110,17 @@ def match(t1: torch.Tensor, t2: torch.Tensor) -> int: class RadixCache: - def __init__(self, total_token_num, rank_in_node, mem_manager=None, page_size: int = 1): - from lightllm.common.kv_cache_mem_manager import MemoryManager - + def __init__(self, total_token_num, rank_in_node, att_resource_manager=None, page_size: int = 1): self.total_token_num = total_token_num - self.mem_manager: MemoryManager = mem_manager + self.att_resource_manager = att_resource_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 if page_size < 1: raise ValueError(f"page_size must be >= 1, got {page_size}") - if mem_manager is not None and page_size != mem_manager.page_size: + allocator = None if att_resource_manager is None else att_resource_manager.full_att_allocator + if allocator is not None and page_size != allocator.page_size: raise ValueError( - f"RadixCache page_size {page_size} must match mem_manager page_size {mem_manager.page_size}" + f"RadixCache page_size {page_size} must match allocator page_size {allocator.page_size}" ) self.page_size = page_size @@ -546,9 +545,9 @@ def _print_helper(self, node: TreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - assert self.mem_manager is not None - if need_token_num > self.mem_manager.allocator.can_use_mem_size: - need_evict_token_num = need_token_num - self.mem_manager.allocator.can_use_mem_size + allocator = self.att_resource_manager.full_att_allocator + if need_token_num > allocator.can_use_mem_size: + need_evict_token_num = need_token_num - allocator.can_use_mem_size release_mems = [] def release_mem(mem_index): @@ -557,7 +556,7 @@ def release_mem(mem_index): self.evict(need_evict_token_num, release_mem) mem_index = torch.concat(release_mems) - self.mem_manager.free(mem_index) + allocator.free(mem_index) return diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index 189d3b6b78..f2c089cc0e 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -27,7 +27,7 @@ from lightllm.utils.profiler import ProfilerCmd from lightllm.server.router.token_load import TokenLoad from lightllm.server.metrics.manager import MetricClient -from lightllm.common.kv_cache_mem_manager import ReadOnlyStaticsMemoryManager +from lightllm.common.att_resources.allocators import KVCapacityReader from lightllm.utils.graceful_utils import graceful_registry from lightllm.utils.process_check import start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name @@ -67,7 +67,7 @@ def __init__(self, args: StartArgs): self.shm_max_total_token_num = SharedInt("shm_max_total_token_num") self.shm_req_manager = ShmReqManager() # 用共享内存进行共享，router 模块读取进行精确的调度估计 - self.read_only_statics_mem_manager = ReadOnlyStaticsMemoryManager() + self.kv_capacity_reader = KVCapacityReader() # 初始化 radix_cache_client 用于读取 prompt cache 的管理信息 self.radix_cache_client = None @@ -236,7 +236,7 @@ async def loop_for_fwd( token_ratio1 = self.get_used_tokens(dp_index) / self.max_total_token_num token_ratio2 = ( self.max_total_token_num - - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + - self.kv_capacity_reader.get_available_tokens(dp_index) ) / self.max_total_token_num d_i = dp_index estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(d_i) @@ -387,11 +387,11 @@ def get_used_tokens(self, dp_index): if not self.args.disable_dynamic_prompt_cache: return ( self.max_total_token_num - - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + - self.kv_capacity_reader.get_available_tokens(dp_index) - self.radix_cache_client.get_unrefed_tokens_num(dp_index) ) else: - return self.max_total_token_num - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + return self.max_total_token_num - self.kv_capacity_reader.get_available_tokens(dp_index) def _add_req(self, group_req_indexes: GroupReqIndexes): req_group = [] diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index e135be36b6..eacf08242a 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1,6 +1,5 @@ import enum import torch -import torch.distributed as dist import numpy as np import collections import pickle @@ -8,7 +7,8 @@ from sortedcontainers import SortedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union -from lightllm.common.req_manager import ReqManager, HybridAttentionReqManager +from lightllm.common.req_manager import ReqManager +from lightllm.common.att_resources import AttResourceManager from lightllm.utils.infer_utils import mark_start, mark_end from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode @@ -39,6 +39,7 @@ class InferenceContext: backend: Optional["ModeBackend"] = None req_manager: ReqManager = None # gpu 请求管理 + att_resource_manager: AttResourceManager = None radix_cache: Union[HybridAttPagedRadixCache, RadixCache] = None shm_req_manager: ShmReqManager = None # 共享内存请求对象管理 requests_mapping: Dict[int, "InferReq"] = None @@ -64,6 +65,7 @@ def register( self.args = get_env_start_args() self.backend: ModeBackend = backend self.req_manager = req_manager + self.att_resource_manager = backend.model.att_resource_manager self.req_sampling_manager = self.req_manager.req_sampling_params_manager self.radix_cache = radix_cache self.shm_req_manager = shm_req_manager @@ -74,7 +76,7 @@ def register( self.vocab_size = vocab_size - self.is_hybrid_att_model = isinstance(self.req_manager, HybridAttentionReqManager) + self.is_hybrid_att_model = backend.is_hybrid_att_model self.output_constraint_state_factory = self._init_output_constraint_state_factory() return @@ -174,11 +176,11 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 if req.tail_small_page_buffer_id is not None: - self.radix_cache.small_page_buffers.free_state_cache([req.tail_small_page_buffer_id]) + self.radix_cache._release_pages("small", [req.tail_small_page_buffer_id]) req.tail_small_page_buffer_id = None # 释放请求执行期间申请、但不再插入 radix cache 的大页状态。 if req.hybrid_len_to_big_page_id: - self.radix_cache.big_page_buffers.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + self.radix_cache._release_pages("big", list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() # 解除请求对已命中 GPU radix cache 前缀节点的引用。 @@ -292,7 +294,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): # big page 模式下请求在 prefill 跨过 big page 边界后、到达末尾前被 pause / abort。 # 若不释放，会泄漏 big page state slot，并触发 free_a_req_mem 中 dict 为空的断言。 if req.hybrid_len_to_big_page_id: - self.radix_cache.big_page_buffers.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + self.radix_cache._release_pages("big", list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() req.cur_kv_len = shared_kv_len @@ -305,20 +307,6 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): assert False, f"error state: cur_kv_len: {req.cur_kv_len}" - def _save_promptcache_kvbuffer(self): - """ - save prompt cache kv buffer - 这个接口是用于保存非量化的缓存prompt cache资源，是定制场景使用的接口，当前代码中不会有调用。 - 其保存的 kv 会配合量化推理模式, 加载到量化推理的prompt cache中, 提升量化推理的精度。 - like paper: - https://arxiv.org/abs/2403.01241 - """ - prompt_cache_token_id = list(self.radix_cache.root_node.children.values())[0].token_id_key - print(f"prompt_cache_token_id : {prompt_cache_token_id}") - index = range(len(prompt_cache_token_id)) - prompt_cache_kv_buffer = self.radix_cache.mem_manager.get_index_kv_buffer(index) - torch.save(prompt_cache_kv_buffer, f"prompt_cache_rank_{dist.get_rank()}.pt") - @torch.no_grad() def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool = True): if len(finished_request_ids) == 0: @@ -348,7 +336,8 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool if free_token_index: free_token_index = custom_cat(free_token_index) - self.req_manager.free(free_req_index, free_token_index) + self.att_resource_manager.full_att_allocator.free(free_token_index) + self.req_manager.free(free_req_index) finished_req_ids_set = set(finished_request_ids) self.infer_req_ids = [_id for _id in self.infer_req_ids if _id not in finished_req_ids_set] @@ -358,8 +347,8 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool f"free a batch state:\n" f"radix refed token num {self.radix_cache.get_refed_tokens_num()}\n" f"radix hold token num {self.radix_cache.get_tree_total_tokens_num()}\n" - f"mem manager can alloc token num {self.req_manager.mem_manager.allocator.can_use_mem_size}\n" - f"mem manager total size {self.req_manager.mem_manager.allocator.size}\n" + f"KV allocator available tokens {self.att_resource_manager.full_att_allocator.can_use_mem_size}\n" + f"KV allocator total tokens {self.att_resource_manager.full_att_allocator.size}\n" ) return @@ -388,7 +377,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if len(free_token_index) != 0: free_token_index = custom_cat(free_token_index) - self.req_manager.free_token(free_token_index) + self.att_resource_manager.full_att_allocator.free(free_token_index) return self def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool): @@ -422,7 +411,7 @@ def get_can_alloc_token_num(self): radix_cache_unref_token_num = ( self.radix_cache.get_tree_total_tokens_num() - self.radix_cache.get_refed_tokens_num() ) - return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num + return self.att_resource_manager.full_att_allocator.can_use_mem_size + radix_cache_unref_token_num def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" @@ -441,8 +430,9 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR # 完整图片可能跨过多个大页结束点，例如页长 256、chunk [128, 768) 覆盖 256/512/768。 # 这些页的 KV 都有效，仍需为 radix 插入和 CPU offload 分配对应槽位。 for page_end in range(first_page_end, last_page_end + 1, big_page_token_num): - big_page_id = self.radix_cache.big_page_buffers.alloc_one_state_cache() - assert big_page_id is not None + slots = self.att_resource_manager.bounded_att_big_page_allocator.alloc(1) + assert slots is not None + big_page_id = slots[0] assert page_end not in req.hybrid_len_to_big_page_id req.hybrid_len_to_big_page_id[page_end] = big_page_id # 当前运行态只对应 chunk 终点；图片内部页尾不保存 state，匹配时也不能在此恢复。 @@ -452,10 +442,11 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR assert len(b_req_idx) == len(big_page_buffer_ids) if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): - self.req_manager.save_big_page_states( - b_req_idx=b_req_idx, - reqs=reqs, - buffer_indexes=big_page_buffer_ids, + selected = [(req, slot) for req, slot in zip(reqs, big_page_buffer_ids) if slot != -1] + self.att_resource_manager.save_bounded_att_pages( + [req.req_idx for req, _ in selected], + page_kind="big", + slot_ids=[slot for _, slot in selected], ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -467,13 +458,15 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR assert req.tail_small_page_buffer_id is None if req.hybrid_cache_len % big_page_token_num != 0: self.radix_cache.free_one_small_page_buffer() - req.tail_small_page_buffer_id = self.radix_cache.small_page_buffers.alloc_one_state_cache() + allocator = self.att_resource_manager.bounded_att_small_page_allocator + slots = allocator.alloc(1) if allocator is not None else None + req.tail_small_page_buffer_id = slots[0] if slots else None if req.tail_small_page_buffer_id is not None: dst_buffer_idx = req.tail_small_page_buffer_id - self.req_manager.save_state( - req=req, - buffer_idx=dst_buffer_idx, - state_cache_manager=self.radix_cache.small_page_buffers, + self.att_resource_manager.save_bounded_att_pages( + [req.req_idx], + page_kind="small", + slot_ids=[dst_buffer_idx], ) return @@ -637,11 +630,7 @@ def _init_all_state(self): # 仅载入调用方显式指定的 token ID 序列；字符串由 detokenization / HTTP 匹配。 self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_token_ids_list() self.multimodal_params = self.multimodal_params.to_dict() - self.mm_input_spans = ( - get_mm_input_spans(self.get_input_token_ids(), self.multimodal_params) - if self.args.disable_chunked_mm_input - else [] - ) + self.mm_input_spans = get_mm_input_spans(self.multimodal_params) if self.args.disable_chunked_mm_input else [] self.shared_kv_node: Union[TreeNode, HybridAttPagedTreeNode] = None self.finish_status = FinishStatus() @@ -679,6 +668,12 @@ def _match_radix_cache(self): self.shm_req.shm_cur_kv_len = self.cur_kv_len return + def _restore_attention_page(self, node, page_kind): + resources = g_infer_context.att_resource_manager + slot = node.big_page_buffer_idx if page_kind == "big" else node.small_page_buffer_idx + resources.restore_bounded_att_pages([self.req_idx], page_kind=page_kind, slot_ids=[slot]) + resources.reset_attention_tail_state([self.req_idx]) + def _hybrid_match_radix_cache(self): assert ( g_infer_context.is_hybrid_att_model is True @@ -713,9 +708,7 @@ def _hybrid_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=share_node.big_page_buffer_idx, req=self - ) + self._restore_attention_page(share_node, "big") else: # 小页匹配 if big_page_is_disable: @@ -728,9 +721,7 @@ def _hybrid_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_small_page_state( - req=self, - ) + self._restore_attention_page(share_node, "small") else: # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 # 对应的小页对应的节点，递归找到对应最近的大叶节点进行返回,然后赋值到req.shared_node 对象上 @@ -742,7 +733,7 @@ def _hybrid_match_radix_cache(self): # 有充足的token 容量时 assert need_tokens % self.args.page_size == 0 radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) - tail_mems = radix_cache.mem_manager.alloc(need_size=need_tokens) + tail_mems = g_infer_context.att_resource_manager.full_att_allocator.alloc(need_tokens) g_infer_context.req_manager.req_to_token_indexs[ self.req_idx, 0:cur_big_page_tokens ] = value_tensor[0:cur_big_page_tokens] @@ -751,8 +742,9 @@ def _hybrid_match_radix_cache(self): ] = tail_mems # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 - radix_cache.mem_manager.operator.copy_mem_to_mem( - value_tensor[cur_big_page_tokens:shared_kv_len], tail_mems + g_infer_context.att_resource_manager.copy_full_att_tokens( + value_tensor[cur_big_page_tokens:shared_kv_len].to(device="cuda", non_blocking=True), + tail_mems.to(device="cuda", non_blocking=True), ) # 尾部 KV 换到新 mem 后，同步拷贝已捕获的 top-k prompt logprobs。 self.prompt_selected_logprobs.copy_capture_slots_if_needed( @@ -760,11 +752,7 @@ def _hybrid_match_radix_cache(self): destination_indexes=tail_mems, ) - self.shared_kv_node = share_node # 只是为了保证 restore_small_page_state 正确调用 - g_infer_context.req_manager.restore_small_page_state( - req=self, - ) - self.shared_kv_node = None + self._restore_attention_page(share_node, "small") big_page_shared_node = radix_cache.deref_to_first_big_page_node(node=share_node) self.shared_kv_node = big_page_shared_node @@ -788,9 +776,7 @@ def _hybrid_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=share_node.big_page_buffer_idx, req=self - ) + self._restore_attention_page(share_node, "big") self.shm_req.shm_cur_kv_len = self.cur_kv_len @@ -802,7 +788,8 @@ def _hybrid_match_radix_cache(self): if self.cur_kv_len == 0: # 说明没有任何命中 - g_infer_context.req_manager.init_hybrid_attention_state(req=self) + g_infer_context.att_resource_manager.reset_bounded_att_state([self.req_idx]) + g_infer_context.att_resource_manager.reset_attention_tail_state([self.req_idx]) return def is_master_req(self): diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6c621b6adb..0e68642414 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,5 +1,3 @@ -import os - import numpy as np import torch import time @@ -15,7 +13,6 @@ from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager -from lightllm.common.req_manager import HybridAttentionReqManager from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput @@ -147,40 +144,50 @@ def init_model(self, kvargs): } self.model, self.is_multimodal = get_model(model_cfg, model_kvargs) self.model: TpPartBaseModel = self.model # for easy typing - set_random_seed(2147483647) - self.is_hybrid_att_model = isinstance(self.model.req_manager, HybridAttentionReqManager) - - if self.is_hybrid_att_model: - self.small_page_buffers = self.model.req_manager.create_small_page_cache_manager( - size=self.args.linear_att_cache_size, + self.draft_models = [] + if self.args.mtp_mode is not None: + self.init_mtp_draft_model(model_kvargs) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.att_resources import AttResourceManager + + models = (self.model, *self.draft_models) + assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + for item in models: + item.prepare_att_resources() + manager: AttResourceManager = self.model.att_resource_manager + assert all(item.max_req_num == self.model.max_req_num for item in models), "models must share the request domain" + assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.initialize_resources( + self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) - else: - self.small_page_buffers = None + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + set_random_seed(2147483647) + self.is_hybrid_att_model = bool(self.model.att_resource_manager.get_bounded_attention_states()) if not self.use_dynamic_prompt_cache: self.radix_cache = None else: if self.is_hybrid_att_model: self.radix_cache = HybridAttPagedRadixCache( - total_token_num=self.model.mem_manager.size, + total_token_num=self.model.att_resource_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, hash_page_size=self.args.linear_att_hash_page_size, big_page_num=self.args.linear_att_page_block_num, - kv_cache_mem_manager=self.model.mem_manager, - small_page_buffers=self.small_page_buffers, + att_resource_manager=self.model.att_resource_manager, ) else: self.radix_cache = RadixCache( - total_token_num=self.model.mem_manager.size, + total_token_num=self.model.att_resource_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, - mem_manager=self.model.mem_manager, + att_resource_manager=self.model.att_resource_manager, page_size=self.args.page_size, ) - if "prompt_cache_kv_buffer" in model_cfg: - assert self.use_dynamic_prompt_cache - self.preload_prompt_cache_kv_buffer(model_cfg) - self.logger.info(f"loaded model class {self.model.__class__}") cache_placement_controller = create_cache_placement_controller( @@ -219,16 +226,26 @@ def init_model(self, kvargs): ) if self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch: - # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便 - # 读取 - self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) + from lightllm.server.router.model_infer.shared_att_resources_utils import publish_att_resources + + # PD 和 DP 使用独立导出，避免 PD 读者收到大小页池或 DP 读者丢失页池。 + if self.args.run_mode in ("prefill", "decode"): + publish_att_resources( + self.model.att_resource_manager, self.model.req_manager.req_to_token_indexs, + kind="pd", reader_count=self.node_world_size, + ) + if self.args.enable_dp_prompt_cache_fetch: + publish_att_resources( + self.model.att_resource_manager, self.model.req_manager.req_to_token_indexs, + kind="dp", reader_count=self.dp_size_in_node - 1, + ) dist.barrier(group=self.node_nccl_group) # 同一 DP 组内只需主 rank 初始化真实的 capture buffer 并执行后续相关操作； # 非主 rank 不需要分配 buffer，避免重复占用内存。 if self.is_master_in_dp: # Capture 只保存 allocator 管理的真实 KV 槽位，HOLD 页对应的 padding 写入由 kernel 过滤。 - kv_cache_size = self.model.mem_manager.size + kv_cache_size = self.model.att_resource_manager.full_att_allocator.size if self.args.enable_prompt_logprobs: mgr = PromptLogprobsCaptureManager.get_instance() if mgr is not None: @@ -248,7 +265,6 @@ def init_model(self, kvargs): self.shm_pd_trans_io_buffer = ShmObjsIOBuffer(tail_str="pd") if self.args.mtp_mode is not None: - self.init_mtp_draft_model(model_kvargs) self.init_spec_engine() if self.args.enable_cpu_cache: @@ -274,7 +290,7 @@ def init_custom(self): def init_dp_kv_shared(self): from lightllm.server.router.model_infer.mode_backend.dp_backend.dp_shared_kv_trans import DPKVSharedMoudle - from lightllm.common.kv_cache_mem_manager import MemoryManager + from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources torch.cuda.set_device(get_current_device_id()) @@ -284,17 +300,25 @@ def init_dp_kv_shared(self): backend=self, ) - # Collect mem_managers from all ranks - self.mem_managers = [] - for rank_idx in range(self.node_world_size): + # 每个 DP 只保留对应 TP 分片，统一以节点内 DP rank 索引。 + self.dp_att_resource_managers = {} + self.dp_req_to_token_indices = {} + # 取节点内的 TP 起点，兼容单个 TP 组跨节点、rank_in_dp 超过本节点 rank 范围的情况。 + for dp_rank, rank_idx in enumerate( + range(self.rank_in_node % self.dp_world_size, self.node_world_size, self.dp_world_size) + ): if rank_idx != self.rank_in_node: - self.mem_managers.append(MemoryManager.loads_from_shm(rank_idx)) + att_resource_manager, req_to_token_indices = load_att_resources(rank_idx, kind="dp") + att_resource_manager.register_page_pools_pin_memory() else: - self.mem_managers.append(self.model.mem_manager) + att_resource_manager = self.model.att_resource_manager + req_to_token_indices = self.model.req_manager.req_to_token_indexs + self.dp_att_resource_managers[dp_rank] = att_resource_manager + self.dp_req_to_token_indices[dp_rank] = req_to_token_indices return def get_max_total_token_num(self): - return self.model.mem_manager.size + return self.model.att_resource_manager.full_att_allocator.size def infer_loop(self): raise NotImplementedError() @@ -311,9 +335,10 @@ def init_mtp_draft_model(self, main_kvargs: dict): spec_mode = self.args.mtp_mode is_chained_draft = spec_mode in ("vanilla_with_att", "vanilla_no_att") - os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" - draft_model_count = self.max_draft_step if is_chained_draft else 1 + from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + + assert 0 < draft_model_count <= MAX_DRAFT_MODELS, "at most 36 draft models are supported" draft_model_dirs = self.args.mtp_draft_model_dir assert draft_model_dirs is not None assert len(draft_model_dirs) >= draft_model_count @@ -322,7 +347,7 @@ def init_mtp_draft_model(self, main_kvargs: dict): draft_model_cfg, _ = PretrainedConfig.get_config_dict(draft_model_dirs[i]) draft_model_kvargs = { "weight_dir": draft_model_dirs[i], - "max_total_token_num": self.model.mem_manager.size, + "max_total_token_num": main_kvargs["max_total_token_num"], "load_way": main_kvargs["load_way"], "max_req_num": main_kvargs.get("max_req_num", 1000), "max_seq_length": main_kvargs.get("max_seq_length", 1024 * 5), @@ -632,7 +657,7 @@ def _alloc_req_kv_mem( old_hold_kv_len = req_obj.hold_kv_len new_hold_kv_len = old_hold_kv_len + alloc_token_num - mem_indexes = g_infer_context.req_manager.mem_manager.alloc(alloc_token_num) + mem_indexes = g_infer_context.att_resource_manager.full_att_allocator.alloc(alloc_token_num) # 高频调度路径允许异步写入请求索引表，其他调用方默认保持原有的同步拷贝语义。 g_infer_context.req_manager.req_to_token_indexs[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( mem_indexes, non_blocking=no_blcoking_copy @@ -1009,27 +1034,6 @@ def _dp_all_reduce_decode_req_num(self, decode_reqs: List[InferReq]) -> int: max_decode_num = self.dp_reduce_tensor.item() return max_decode_num - def preload_prompt_cache_kv_buffer(self, model_cfg): - self.logger.info("Preload prompt cache kv buffer.") - cur_rank = dist.get_rank() - prompt_cache_kv_buffer_path = os.path.join( - self.weight_dir, model_cfg["prompt_cache_kv_buffer"][f"rank_{cur_rank}"] - ) - page_size = self.args.page_size - intact_kv_len = len(model_cfg["prompt_cache_token_ids"]) // page_size * page_size - if intact_kv_len == 0: - return - - prompt_cache_kv_buffer = torch.load(prompt_cache_kv_buffer_path, weights_only=True, map_location="cpu") - prompt_cache_kv_buffer = {name: buffer[:, :intact_kv_len] for name, buffer in prompt_cache_kv_buffer.items()} - intact_kv_index = self.radix_cache.mem_manager.alloc(intact_kv_len) - self.radix_cache.mem_manager.load_index_kv_buffer(intact_kv_index, prompt_cache_kv_buffer) - intact_token_ids = torch.tensor( - model_cfg["prompt_cache_token_ids"][:intact_kv_len], dtype=torch.int64, device="cpu" - ) - self.radix_cache.insert(intact_token_ids, intact_kv_index) - self.radix_cache.match_prefix(intact_token_ids, update_refs=True) - def init_rank_infos(self): self.node_world_size = get_node_world_size() self.rank_in_node = get_current_rank_in_node() diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 46586959f1..af49f208e4 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -182,7 +182,7 @@ def _master_req_to_radix_cache(self, master_req: InferReq): prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value) old_prefix_len = 0 if master_req.shared_kv_node is None else master_req.shared_kv_node.node_prefix_total_len assert old_prefix_len <= master_req.cur_kv_len - self.model.mem_manager.free( + self.model.att_resource_manager.full_att_allocator.free( self.model.req_manager.req_to_token_indexs[master_req.req_idx][old_prefix_len:prefix_len] ) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 413a86b94e..00e0c6b00c 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -1,15 +1,12 @@ # 该文件用于提供在数据dp并行的推理模式下，共享kv cache trans相关的功能函数模块 -import time import numpy as np import dataclasses import torch from typing import List -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_dp_rank_in_node from lightllm.server.core.objs.shm_array import ShmArray from ...infer_batch import InferReq -from lightllm.utils.dist_utils import get_current_device_id from lightllm.server.router.model_infer.infer_batch import g_infer_context import torch.distributed as dist @@ -24,28 +21,49 @@ def __init__(self, max_req_num: int, dp_size_in_node: int, backend): self.backend: DPChunkedPrefillBackend = backend self.max_req_num = max_req_num - # 0 代表 kv_len, 1 代表 radix_cache_len + self.big_page_tokens = ( + backend.args.linear_att_hash_page_size * backend.args.linear_att_page_block_num + if backend.is_hybrid_att_model else 0 + ) + max_big_pages = backend.args.max_req_total_len // self.big_page_tokens if self.big_page_tokens else 0 + # 原共享信息只保存 KV 长度和请求槽位，大页列表使用独立共享数组。 self.shared_req_infos = ShmArray( name="dp_shared_req_infos", shape=(self.max_req_num, dp_size_in_node, 2), dtype=np.int64, ) self.shared_req_infos.create_shm() + self.shared_req_big_page_ids: ShmArray | None = None + if max_big_pages: + # 第 i 项对应前缀长度 (i + 1) * big_page_tokens，未使用位置填 -1。 + self.shared_req_big_page_ids = ShmArray( + name="dp_shared_req_big_page_ids", + shape=(self.max_req_num, dp_size_in_node, max_big_pages), + dtype=np.int64, + ) + self.shared_req_big_page_ids.create_shm() self.dp_rank_in_node = get_dp_rank_in_node() assert get_env_start_args().diverse_mode is False - def fill_reqs_info(self, reqs: List[InferReq]): - """ - 填充请求的 kv 信息到共享内存中 - """ + def fill_reqs_info(self, reqs: List[InferReq]) -> None: + """发布每个请求的 KV 长度、请求槽位及完整有序大页列表。""" dist.barrier(group=self.backend.node_nccl_group) if self.backend.is_master_in_dp: - self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._KV_LEN_INDEX] = [ - req.cur_kv_len for req in reqs - ] - self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._REQ_IDX_INDEX] = [ - req.req_idx for req in reqs - ] + for req_index, req in enumerate(reqs): + row = self.shared_req_infos.arr[req_index, self.dp_rank_in_node] + row[self._KV_LEN_INDEX] = req.cur_kv_len + row[self._REQ_IDX_INDEX] = req.req_idx + if self.shared_req_big_page_ids is None: + continue + page_ids = self.shared_req_big_page_ids.arr[req_index, self.dp_rank_in_node] + page_ids.fill(-1) # 批次复用共享行时，清除上次较长请求留下的页 ID。 + # radix 页按路径从根到叶排列；请求持有的页按其前缀长度放到对应位置。 + if req.shared_kv_node is not None: + shared_pages = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + page_ids[:len(shared_pages)] = shared_pages + for token_len, page_id in req.hybrid_len_to_big_page_id.items(): + if token_len <= req.cur_kv_len: + page_ids[token_len // self.big_page_tokens - 1] = page_id return def build_shared_kv_trans_tasks( @@ -59,10 +77,9 @@ def build_shared_kv_trans_tasks( dist.barrier(group=self.backend.node_nccl_group) trans_tasks: List[TransTask] = [] - rank_max_radix_cache_lens = np.max( - self.shared_req_infos.arr[0 : len(reqs), :, self._KV_LEN_INDEX], axis=1, keepdims=False - ) - # 如果发现自己是dp_rank 最小， radix_cache_len 最长的请求，则将数据写入到共享内存中。 + available_lens = self.shared_req_infos.arr[0 : len(reqs), :, self._KV_LEN_INDEX] + rank_max_radix_cache_lens = np.max(available_lens, axis=1) + # 为本 DP 负责的请求选择可用前缀最长的来源。 for req_index, req, max_req_radix_cache_len, req_dp_rank in zip( list(range(len(reqs))), reqs, rank_max_radix_cache_lens, req_dp_ranks ): @@ -75,67 +92,82 @@ def build_shared_kv_trans_tasks( if is_current_dp_handle and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): assert req.hold_kv_len == req.cur_kv_len - mem_indexes = self.backend._alloc_req_kv_mem(req, alloc_token_num) - assert mem_indexes is not None - # mem_indexes 只描述需要复制的逻辑 KV；页尾预留槽位已经由 - # _alloc_req_kv_mem 写入请求表，但不参与本次跨 DP 传输。 - mem_indexes = mem_indexes[:trans_size] - max_kv_len_dp_rank = self.shared_req_infos.arr[req_index, :, self._KV_LEN_INDEX].argmax() - max_kv_len_req_idx = int(self.shared_req_infos.arr[req_index, max_kv_len_dp_rank, self._REQ_IDX_INDEX]) - max_kv_len_mem_manager_index = max_kv_len_dp_rank * self.backend.dp_world_size + self.backend.rank_in_dp - max_kv_len_mem_manager: MemoryManager = self.backend.mem_managers[max_kv_len_mem_manager_index] - max_kv_len_mem_indexes = max_kv_len_mem_manager.req_to_token_indexs[ - max_kv_len_req_idx, req.cur_kv_len : req.cur_kv_len + trans_size + src_dp_rank = int(available_lens[req_index].argmax()) + src_info = self.shared_req_infos.arr[req_index, src_dp_rank] + src_req_idx = int(src_info[self._REQ_IDX_INDEX]) + src_big_page_ids, dst_big_page_ids = [], [] + if self.backend.is_hybrid_att_model: + # 只复制新增前缀跨过的完整大页，已有前缀页继续使用本地引用。 + first_page = req.cur_kv_len // self.big_page_tokens + end_page = target_kv_len // self.big_page_tokens + if end_page > first_page: + src_big_page_ids = self.shared_req_big_page_ids.arr[ + req_index, src_dp_rank, first_page:end_page + ].tolist() + assert all(page_id >= 0 for page_id in src_big_page_ids), "DP source is missing prefix big pages" + allocator = self.backend.model.att_resource_manager.bounded_att_big_page_allocator + dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) + assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" + dst_token_indexes = self.backend._alloc_req_kv_mem(req, alloc_token_num) + assert dst_token_indexes is not None + # 只搬运有效 KV；页尾预留槽位已写入请求表，不参与本次传输。 + dst_token_indexes = dst_token_indexes[:trans_size] + src_token_indexes = self.backend.dp_req_to_token_indices[src_dp_rank][ + src_req_idx, req.cur_kv_len:target_kv_len ] trans_tasks.append( TransTask( - req=req, - mem_indexes=mem_indexes, - max_kv_len_dp_rank=int(max_kv_len_dp_rank), - max_kv_len_mem_manager_index=int(max_kv_len_mem_manager_index), - max_kv_len_mem_indexes=max_kv_len_mem_indexes, + dst_req=req, + src_dp_rank=src_dp_rank, + src_req_idx=src_req_idx, + src_token_indexes=src_token_indexes, + dst_token_indexes=dst_token_indexes, + src_big_page_ids=src_big_page_ids, + dst_big_page_ids=dst_big_page_ids, ) ) return trans_tasks - def kv_trans(self, trans_tasks: List["TransTask"]): - # kv 传输 - if len(trans_tasks) > 0: - max_kv_len_mem_indexes = [] - max_kv_len_dp_ranks = [] - mem_indexes = [] - - for i, trans_task in enumerate(trans_tasks): - max_kv_len_mem_indexes.append(trans_task.max_kv_len_mem_indexes) - max_kv_len_dp_ranks.extend([trans_task.max_kv_len_dp_rank] * len(trans_task.max_kv_len_mem_indexes)) - mem_indexes.append(trans_task.mem_indexes) - - max_kv_len_mem_indexes_tensor = torch.cat(max_kv_len_mem_indexes).to(dtype=torch.int64, device="cuda") - max_kv_len_dp_ranks_tensor = torch.tensor(max_kv_len_dp_ranks, dtype=torch.int32, device="cuda") - mem_indexes_tensor = torch.cat(mem_indexes).to(dtype=torch.int64, device="cuda") - self.backend.model.mem_manager.operator.copy_kv_from_other_dp_ranks( - mem_managers=self.backend.mem_managers, - move_token_indexes=max_kv_len_mem_indexes_tensor, - token_dp_indexes=max_kv_len_dp_ranks_tensor, - mem_indexes=mem_indexes_tensor, - dp_size_in_node=self.backend.dp_size_in_node, - rank_in_dp=self.backend.rank_in_dp, + def kv_trans(self, trans_tasks: List["TransTask"]) -> None: + target = self.backend.model.att_resource_manager + for task in trans_tasks: + source = self.backend.dp_att_resource_managers[task.src_dp_rank] + # 源索引来自远端 CUDA 请求表；allocator 返回 CPU 索引，在此一次性准备目标 CUDA int64 索引。 + dst_token_indexes = task.dst_token_indexes.to(device="cuda", dtype=torch.long, non_blocking=True) + target.copy_from_other_dp( + remote_att_manager=source, + src_token_indexes=task.src_token_indexes, + dst_token_indexes=dst_token_indexes, + src_big_page_ids=task.src_big_page_ids, + dst_big_page_ids=task.dst_big_page_ids, + src_req_idx=task.src_req_idx, + dst_req_idx=task.dst_req.req_idx, ) - # This copy runs during request initialization, not once per PD transfer page. - self.backend.logger.debug(f"dp_i {self.dp_rank_in_node} transfer kv tokens num: {len(mem_indexes_tensor)}") + # Source requests may be released on other ranks immediately afterwards. + # All ranks participate, including those without a local transfer task. + dist.barrier(group=self.backend.node_nccl_group) - for trans_task in trans_tasks: - trans_task.req.cur_kv_len += len(trans_task.mem_indexes) - assert trans_task.req.cur_kv_len <= trans_task.req.hold_kv_len + for task in trans_tasks: + req = task.dst_req + # 复制完成后，才把目标大页发布为请求持有的快照；后续 radix 插入或释放沿用此表。 + if task.dst_big_page_ids: + first_page_end = (req.cur_kv_len // self.big_page_tokens + 1) * self.big_page_tokens + for page_offset, page_id in enumerate(task.dst_big_page_ids): + token_len = first_page_end + page_offset * self.big_page_tokens + req.hybrid_len_to_big_page_id[token_len] = page_id + req.cur_kv_len += len(task.dst_token_indexes) + assert req.cur_kv_len <= req.hold_kv_len if self.backend.is_master_in_dp: - trans_task.req.shm_req.shm_cur_kv_len = trans_task.req.cur_kv_len + req.shm_req.shm_cur_kv_len = req.cur_kv_len @dataclasses.dataclass class TransTask: - req: InferReq - mem_indexes: torch.Tensor - max_kv_len_dp_rank: int - max_kv_len_mem_manager_index: int - max_kv_len_mem_indexes: torch.Tensor + dst_req: InferReq # 本 DP 接收前缀的请求，持有目标请求槽位和新分配的资源。 + src_dp_rank: int # 来源 DP 在当前节点内的编号。 + src_req_idx: int # 远端请求槽位，用于复制 prefill 基准运行态。 + src_token_indexes: torch.Tensor # 来源 CUDA 请求表中的有效 KV 索引。 + dst_token_indexes: torch.Tensor # allocator 返回的目标 CPU 索引，提交时统一转为 CUDA int64。 + src_big_page_ids: List[int] # 按前缀长度递增排列的源大页槽位。 + dst_big_page_ids: List[int] # 与源列表一一配对的新分配目标大页槽位。 diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 736e52f439..dbd48b37f6 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -1,4 +1,3 @@ -import threading import torch.distributed as dist import torch import dataclasses @@ -8,11 +7,9 @@ from collections import deque from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.config_utils import is_hybrid_att_model from lightllm.utils.envs_utils import get_env_start_args from ..infer_batch import InferReq from lightllm.utils.dist_utils import create_new_group_for_current_dp -from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu, load_cpu_kv_to_gpu from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.log_utils import init_logger from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len @@ -35,11 +32,12 @@ def __init__(self, backend): dist.barrier(group=self.offload_sync_group) self.offload_sync_tensor = torch.empty((1,), dtype=torch.int32, device="cuda") - self.page_index_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.int32, device="cuda") - self.page_ready_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.bool, device="cuda") - self.cpu_cache_handle_queue: Deque[TransTask] = deque() self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False) + resources = self.backend.model.att_resource_manager + meta = self.cpu_cache_client.kv_cache_tensor_meta + if meta.page_bytes != resources.get_cpu_page_bytes(): + raise ValueError("CPU cache service and model resource page sizes differ") @lru_cache() def need_sync_compute_stream(self) -> bool: @@ -77,7 +75,6 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): continue page_len_list = req.shm_req.token_hash_page_len_list.get_all() - page_len_start_list = [0] + page_len_list assert len(page_list) <= len(page_len_list) matched_page_num = len(page_list) @@ -113,43 +110,13 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= matched_page_num need_pages = page_list[ready_page_num:matched_page_num] - mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) - assert mem_indexes is not None - - if self.need_sync_compute_stream(): - # TODO fa3 现在必须使用同步模式, 未来需要移除 - torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) - # g_infer_context.get_overlap_stream().synchronize() - - mem_manager = self.backend.model.mem_manager - req_manager = self.backend.model.req_manager - - mem_indexes_cuda = mem_indexes.cuda(non_blocking=True) - page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda( - non_blocking=True - ) - # hybrid 页面加载必须按完整 page 处理，否则可能缺失恢复运行态所需的 checkpoint，所以 - # 这里需要进行pad操作，使操作的页面是完整的。 - _start = page_len_start_list[ready_page_num] - - _end = req.cur_kv_len - assert 0 <= _start <= _end, f"invalid pad range [{_start}, {_end}]" - mem_indexes_cuda = torch.cat( - [req_manager.req_to_token_indexs[req.req_idx, _start:_end], mem_indexes_cuda] - ) - - assert len(mem_indexes_cuda) == match_tokens - page_len_start_list[ready_page_num] - - # 更新 req 状态。 - idle_token_num -= need_token_num - req.cur_kv_len = req.cur_kv_len + need_token_num - - mem_manager.operator.load_cpu_cache_to_gpu( - mem_indexes=mem_indexes_cuda, - page_indexes=page_indexes_cuda, - cpu_cache_client=self.cpu_cache_client, + self.read_req_from_cpu( req=req, + cpu_page_indices=need_pages, + page_lengths=page_len_list[ready_page_num:matched_page_num], + start_token_len=page_len_list[ready_page_num - 1] if ready_page_num else 0, ) + idle_token_num -= need_token_num torch.cuda.current_stream().synchronize() @@ -166,6 +133,63 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): self.cpu_cache_client.lock.release() return + def read_req_from_cpu( + self, + req: InferReq, + cpu_page_indices: List[int], + page_lengths: List[int], + *, + start_token_len: int, + ) -> None: + """恢复已选定的请求前缀页面，分配资源并登记完整大页引用。 + + start_token_len 是首个 CPU 页的逻辑起点，可能早于已有 GPU 命中长度。 + 调用方负责选择命中范围、确认容量以及搬运完成后的 CPU 页引用释放。 + """ + if not cpu_page_indices: + return + resources = self.backend.model.att_resource_manager + req_manager = self.backend.model.req_manager + need_token_num = page_lengths[-1] - req.cur_kv_len + + # 阶段一：分配新增 KV，并补齐首个 CPU 页在 GPU 上已经命中的 token 索引。 + mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) + assert mem_indexes is not None + if self.need_sync_compute_stream(): + # FA3 的资源搬运需等待已有计算流完成。 + torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) + assert 0 <= start_token_len <= req.cur_kv_len, "invalid CPU page start" + token_indices = torch.cat([ + req_manager.req_to_token_indexs[req.req_idx, start_token_len:req.cur_kv_len], + mem_indexes.cuda(non_blocking=True), + ]) + assert token_indices.numel() == page_lengths[-1] - start_token_len + + # 阶段二：只为完整 bounded 页分配普通大页，碎页 LOAD 槽位由资源层管理。 + big_page_slot_ids = None + if self.backend.is_hybrid_att_model: + page_tokens = self.args.cpu_cache_token_page_size + complete_page_lengths = [length for length in page_lengths if length % page_tokens == 0] + assert len(complete_page_lengths) == len(page_lengths) - int(page_lengths[-1] % page_tokens != 0), ( + "only the final CPU page may be partial" + ) + allocator = resources.bounded_att_big_page_allocator + big_page_slot_ids = allocator.alloc(len(complete_page_lengths)) + assert big_page_slot_ids is not None, "not enough bounded state slots for CPU cache load" + + # 阶段三：统一恢复各类 State；异常直接向上传播，不在此处回收或继续执行。 + resources.read_req_from_cpu( + self.cpu_cache_client.cpu_kv_cache_tensor, + cpu_page_indices, + token_indices, + layouts=resources.get_cpu_transfer_layouts(), + big_page_slot_ids=big_page_slot_ids, + req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, + ) + if big_page_slot_ids is not None: + req.hybrid_len_to_big_page_id.update(zip(complete_page_lengths, big_page_slot_ids)) + req.cur_kv_len = page_lengths[-1] + def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> List[InferReq]: """ 将满足cpu kv cache 卸载条件的请求进行处理, 并返回真的满足退出条件的请求list。 @@ -208,7 +232,7 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L g_infer_context.get_overlap_stream().synchronize() # 发起将请求的 kv cache 卸载到 cpu cache 中的任务 - trans_task = self._start_kv_cache_offload_task(req=req, cpu_kv_cache_stream=cpu_stream) + trans_task = self.write_req_to_cpu(req=req, cpu_kv_cache_stream=cpu_stream) # 根据是否成功创建了卸载任务，决定是否将请求加入到处理队列中 if trans_task is not None: @@ -222,17 +246,24 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L return true_finished_reqs - def _start_kv_cache_offload_task( + def write_req_to_cpu( self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream ) -> Optional["TransTask"]: + """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers + producer = torch.cuda.current_stream() + # 阶段一：等待当前生产流已提交的数据写入，不阻塞主机线程。 + cpu_kv_cache_stream.wait_stream(producer) + # TODO：将 bounded CPU→CPU 快照搬运改为当前流上的 kernel。 + # 现有主机 copy 不受 wait_stream 约束，kernel 落地前源快照仍须预先就绪。 with torch.cuda.stream(cpu_kv_cache_stream): # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() page_len_list = req.shm_req.token_hash_page_len_list.get_all() assert len(token_hash_list) == len(page_len_list) + # 阶段二：确定本次可写出范围，申请 CPU 页。 if self.backend.is_master_in_dp: find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) @@ -257,6 +288,12 @@ def _start_kv_cache_offload_task( finally: self.cpu_cache_client.lock.release() + # 分配不足时返回 -1 并补齐剩余位置，只保留成功分配的连续前缀。 + # 第一个 -1 及后续页面（包括可能的尾部碎页）均不参与本次写出。 + if -1 in page_list: + allocated_page_count = page_list.index(-1) + page_list = page_list[:allocated_page_count] + ready_list = ready_list[:allocated_page_count] item_size = len(page_list) if item_size == 0: dist.broadcast_object_list([0], group=self.gloo_group, group_src=0) @@ -276,26 +313,40 @@ def _start_kv_cache_offload_task( page_list = broadcast_data["page_list"] ready_list = broadcast_data["ready_list"] - page_indexes = torch.tensor(page_list, dtype=torch.int32, device="cpu", pin_memory=True) - page_readies = torch.tensor(ready_list, dtype=torch.bool, device="cpu", pin_memory=True) - assert len(page_indexes) <= self.page_index_buffer.shape[0] - cuda_page_indexes = self.page_index_buffer[: len(page_indexes)] - cuda_page_readies = self.page_ready_buffer[: len(page_readies)] - cuda_page_indexes.copy_(page_indexes, non_blocking=True) - cuda_page_readies.copy_(page_readies, non_blocking=True) + page_indexes = torch.tensor(page_list, dtype=torch.int32) + page_readies = torch.tensor(ready_list, dtype=torch.bool) move_token_num = page_len_list[item_size - 1] assert req.cur_kv_len >= move_token_num token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num] - mem_manager = self.backend.model.mem_manager - - mem_manager.operator.offload_gpu_kv_to_cpu_cache( - mem_indexes=token_indexes, - page_indexes=cuda_page_indexes, - page_readies=cuda_page_readies, - cpu_cache_client=self.cpu_cache_client, - req=req, + # 阶段三：只查找实际分配的 CPU 页所需快照；完整大页必须存在，不做缺失回退。 + resources = self.backend.model.att_resource_manager + big_page_slot_ids, small_page_slot_ids = None, None + if self.backend.is_hybrid_att_model: + page_tokens = self.args.cpu_cache_token_page_size + radix_cache = g_infer_context.radix_cache + shared_slots = radix_cache.get_big_page_ids_by_node(req.shared_kv_node) if radix_cache is not None else [] + big_slots = {(i + 1) * page_tokens: slot for i, slot in enumerate(shared_slots)} + big_slots.update(req.hybrid_len_to_big_page_id) + big_page_slot_ids = [ + big_slots[length] for length in page_len_list[:item_size] if length % page_tokens == 0 + ] + small_page_slot_ids = [] + # 只有实际分配的前缀包含原始尾页，才可能使用小页快照。 + # 若分配遇到 -1，item_size 已截短，尾页不会进入此分支。 + includes_tail_page = item_size == len(page_len_list) + if includes_tail_page and page_len_list[-1] % page_tokens != 0: + small_page_slot_ids.append(req.tail_small_page_buffer_id) + # 整批提交并保留 token/page 对齐，由资源层跳过 ready 页。 + resources.write_req_to_cpu( + token_indices=token_indexes, + cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, + cpu_page_indices=page_list, + layouts=resources.get_cpu_transfer_layouts(), + cpu_page_readies=ready_list, + big_page_slot_ids=big_page_slot_ids, + small_page_slot_ids=small_page_slot_ids, ) # 这个操作只是为了在offload 对应的cuda stream中，同步标记下对应的kv cache offload 操作已经完成， diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py index 786c60b30a..fa2183182f 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py @@ -111,7 +111,7 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: .detach() .cpu() ) - self.model.mem_manager.free(mem_indexes) + self.model.att_resource_manager.full_att_allocator.free(mem_indexes) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len if self.is_master_in_dp: diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index 783a58523a..672cb59ac8 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -8,7 +8,9 @@ import pickle from typing import List, Dict, Union, Optional from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.att_resources import AttResourceManager +from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import ( PDChunckedTransTask, PDChunckedTransTaskGroup, @@ -83,20 +85,20 @@ def _init_env( task_out_queue.put("proc_start") - # 从共享内存读取所有rank的mem_manager + # 从共享内存读取所有 rank 的 attention 资源 node_world_size = args.tp // args.nnodes - mem_managers: List[MemoryManager] = [ - MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size) + resource_managers: List[AttResourceManager] = [ + load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] - task_out_queue.put("get_mem_managers_ok") + task_out_queue.put("get_resources_ok") manager = _DecodeTransModule( args=args, device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + resource_managers=resource_managers, up_status_in_queue=up_status_in_queue, ) assert manager is not None @@ -119,7 +121,7 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + resource_managers: List[AttResourceManager], up_status_in_queue: Optional[mp.SimpleQueue], ): self.args = args @@ -127,12 +129,12 @@ def __init__( self.device_id = device_id self.task_in_queue = task_in_queue self.task_out_queue = task_out_queue - self.mem_managers = mem_managers + self.resource_managers = resource_managers self.up_status_in_queue = up_status_in_queue - cur_mem_manager: MemoryManager = self.mem_managers[device_id] - kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer( - page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size + kv_move_buffer = allocate_transfer_buffer( + self.resource_managers[device_id], self.args.pd_kv_page_num, device=f"cuda:{device_id}" ) + self.kv_move_buffer = kv_move_buffer self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -170,13 +172,23 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] - cur_mem.read_page_kv_move_buffer_to_mem( - mem_indexes=[0], + copy_page( + self.resource_managers, + self.kv_move_buffer, + write=True, + token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), page_index=0, dp_index=dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, + tp_size=self.dp_world_size, + ) + copy_page( + self.resource_managers, + self.kv_move_buffer, + write=False, + token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), + page_index=0, + dp_index=dp_index, + tp_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -396,15 +408,19 @@ def read_page_to_mems_loop(self): copy_end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.stream(stream=self.copy_cuda_stream): copy_start_event.record(self.copy_cuda_stream) - cur_mem = self.mem_managers[self.device_id] - cur_mem.read_page_kv_move_buffer_to_mem( - trans_task.mem_indexes, + copy_page( + self.resource_managers, + self.kv_move_buffer, + write=False, + token_indices=torch.tensor( + trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device + ), page_index=trans_task.dst_page_index, dp_index=trans_task.decode_dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, - page_kind=trans_task.page_kind, + tp_size=self.dp_world_size, + kind=trans_task.page_kind, req_idx=trans_task.req_idx, + state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) copy_end_event.record(self.copy_cuda_stream) self.success_queue.put((copy_end_event, copy_start_event, trans_task)) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index d3f2c0214d..a665dca08d 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -56,7 +56,7 @@ def __init__( assert args.run_mode in ["prefill", "decode"], args.run_mode self.is_prefill_node = args.run_mode == "prefill" self.capture_telemetry = False - self.num_pages, self.page_size, self.num_layers, self.kv_head_num, self.head_dims = kv_move_buffer.shape + self.num_pages = kv_move_buffer.shape[0] self.host_ip = host_ip or get_hostname_ip() assert self.host_ip is not None, "can not get host ip for NcclKVTransporter" @@ -91,10 +91,7 @@ def local_page_mem_desc(self) -> bytes: return pickle.dumps( { "num_pages": self.num_pages, - "page_size": self.page_size, - "num_layers": self.num_layers, - "kv_head_num": self.kv_head_num, - "head_dims": self.head_dims, + "page_shape": tuple(self.kv_move_buffer.shape[1:]), "dtype": str(self.kv_move_buffer.dtype), } ) @@ -114,6 +111,11 @@ def connect_add_remote_agent(self, remote_agent: PDAgentMetadata): metadata.agent_name == remote_agent.agent_name ), f"Peer name {metadata.agent_name} does not match remote name {remote_agent.agent_name}" + page = pickle.loads(remote_agent.page_reg_desc) + if page.get("page_shape") != tuple(self.kv_move_buffer.shape[1:]) or page.get("dtype") != str( + self.kv_move_buffer.dtype + ): + raise ValueError("PD peers must use the same page shape and dtype") self.remote_agents[remote_agent.agent_name] = remote_agent logger.info(f"Added NCCL remote agent {remote_agent.agent_name} at {metadata.host_ip}:{metadata.control_port}") return @@ -434,6 +436,7 @@ def __init__( def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[ThreadedServer, int]: last_error = None for cur_port in range(port_min, port_max + 1): + server = None try: server = ThreadedServer( _NcclControlService(self), @@ -446,10 +449,16 @@ def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[Thr "allow_setattr": True, }, ) + # SO_REUSEADDR permits simultaneous binds before either socket + # listens. Claim the port here so EADDRINUSE retries in this + # thread instead of killing the background server after return. + server.listener.listen(server.backlog) threading.Thread(target=server.start, daemon=True).start() logger.info(f"NCCL RPyC control channel listen on {host_ip}:{cur_port}") return server, cur_port except OSError as e: + if server is not None: + server.close() last_error = e if e.errno == errno.EADDRINUSE: logger.info(f"NCCL RPyC control port {host_ip}:{cur_port} is in use, try next port") diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py index b80d52a747..147ba3ee25 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py @@ -61,14 +61,15 @@ def get_new_notifs(self) -> Dict[str, list[bytes]]: return self.nixl_agent.get_new_notifs() def _register_kv_move_buffer(self, kv_move_buffer: Tensor): - self.num_pages, self.page_size, self.num_layers, self.kv_head_num, self.head_dims = kv_move_buffer.shape - self.dtype_byte_size = kv_move_buffer.element_size() - self.page_len = self.page_size * self.num_layers * self.kv_head_num * self.head_dims * self.dtype_byte_size + self.num_pages = kv_move_buffer.shape[0] + self.page_len = kv_move_buffer[0].numel() * kv_move_buffer.element_size() self.page_reg_desc = self.nixl_agent.register_memory(kv_move_buffer) self.page_local_xfer_handles = self._create_paged_xfer_handles(self.page_reg_desc, self.num_pages) def _create_paged_xfer_handles(self, reg_desc: "nixlBind.nixlRegDList", page_num: int, agent_name: str = ""): - base_addr, _, device_id, _ = reg_desc[0] + base_addr, registered_bytes, device_id, _ = reg_desc[0] + if registered_bytes != page_num * self.page_len: + raise ValueError("PD peers must use the same page byte capacity") pages_data = [] for page_id in range(page_num): pages_data.append((base_addr + page_id * self.page_len, self.page_len, device_id)) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index 6e2084b04d..2fdb114225 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -8,7 +8,9 @@ import pickle from typing import List, Dict, Optional, Union from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.att_resources import AttResourceManager +from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask from lightllm.utils.graceful_utils import graceful_registry from lightllm.server.core.objs import StartArgs @@ -58,20 +60,20 @@ def _init_env( graceful_registry(inspect.currentframe().f_code.co_name) task_out_queue.put("proc_start") - # 从共享内存读取所有rank的mem_manager + # 从共享内存读取所有 rank 的 attention 资源 node_world_size = args.tp // args.nnodes - mem_managers: List[MemoryManager] = [ - MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size) + resource_managers: List[AttResourceManager] = [ + load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] - task_out_queue.put("get_mem_managers_ok") + task_out_queue.put("get_resources_ok") manager = _PrefillTransModule( args=args, device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + resource_managers=resource_managers, ) assert manager is not None task_out_queue.put("module_ready") @@ -93,19 +95,19 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + resource_managers: List[AttResourceManager], ) -> None: self.args = args self.dp_world_size = self.args.tp // self.args.dp self.device_id = device_id self.task_in_queue = task_in_queue self.task_out_queue = task_out_queue - self.mem_managers = mem_managers + self.resource_managers = resource_managers - cur_mem_manager: MemoryManager = self.mem_managers[device_id] - kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer( - page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size + kv_move_buffer = allocate_transfer_buffer( + self.resource_managers[device_id], self.args.pd_kv_page_num, device=f"cuda:{device_id}" ) + self.kv_move_buffer = kv_move_buffer self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -145,13 +147,14 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] - cur_mem.write_mem_to_page_kv_move_buffer( - mem_indexes=[0], + copy_page( + self.resource_managers, + self.kv_move_buffer, + write=True, + token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), page_index=0, dp_index=dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, + tp_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -203,15 +206,19 @@ def local_copy_kv_loop(self): # 将kv 数据拷贝到 page 上，然后传输给 decode node，让其进行读取。 with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] - cur_mem.write_mem_to_page_kv_move_buffer( - trans_task.mem_indexes, + copy_page( + self.resource_managers, + self.kv_move_buffer, + write=True, + token_indices=torch.tensor( + trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device + ), page_index=trans_task.src_page_index, dp_index=trans_task.prefill_dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, - page_kind=trans_task.page_kind, + tp_size=self.dp_world_size, + kind=trans_task.page_kind, req_idx=trans_task.req_idx, + state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) sync_event = torch.cuda.Event() sync_event.record() diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py new file mode 100644 index 0000000000..30c54d633a --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -0,0 +1,57 @@ +"""Select the local TP group and submit its State-owned PD operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import torch + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from typing import Literal + + +def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, *, device: str | torch.device | int): + if page_count <= 0: + raise ValueError("PD buffer must contain pages") + return torch.empty((page_count, resources.get_pd_page_bytes()), dtype=torch.uint8, device=device) + + +def copy_page( + resources: Sequence[AttResourceManager], + buffer: torch.Tensor, + *, + page_index: int, + dp_index: int, + tp_size: int, + token_indices: torch.Tensor, + kind: Literal["kv", "att_state"] = "kv", + req_idx: int | None = None, + state_token_len: int | None = None, + write: bool, +): + ranks = resources[dp_index * tp_size : (dp_index + 1) * tp_size] + if dp_index < 0 or tp_size <= 0 or len(ranks) != tp_size or not 0 <= page_index < len(buffer): + raise ValueError("invalid PD group or page index") + rank_layouts = [] + for resource in ranks: + if kind == "kv": + layouts = resource.get_pd_full_att_transfer_layouts() + elif kind == "att_state": + layouts = resource.get_pd_bounded_att_or_tail_transfer_layouts() + else: + raise ValueError(f"unknown PD task kind {kind}") + if not layouts or resource.get_pd_page_bytes() != buffer.shape[1]: + raise ValueError("PD page layout does not match this group") + rank_layouts.append(layouts) + for resource, layouts in zip(ranks, rank_layouts): + kwargs = {"kind": kind, "layouts": layouts} + if kind == "kv": + kwargs["token_indices"] = token_indices + else: + kwargs["req_idx"] = req_idx + kwargs["state_token_len"] = state_token_len + if write: + resource.write_to_pd_buffer(buffer, pd_page_index=page_index, **kwargs) + else: + resource.read_from_pd_buffer(buffer, pd_page_index=page_index, **kwargs) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py b/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py index 84173a44ba..f180071fe3 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py @@ -39,7 +39,7 @@ def init_all( up_status_in_queue, ) assert self.task_out_queue.get(timeout=30) == "proc_start" - assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok" + assert self.task_out_queue.get(timeout=60) == "get_resources_ok" return True diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index edde24602d..898a4e08fd 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -24,6 +24,7 @@ def verify_mtp_tokens( b_mtp_index: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" + # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, @@ -31,14 +32,9 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - if backend.is_hybrid_att_model: - backend.model.req_manager.update_mtp_state( - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=backend.max_draft_step + 1, - ) + backend.model.att_resource_manager.commit_bounded_att_state( + b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" + ) return accept_lengths, accepted_index diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py new file mode 100644 index 0000000000..e031d2e850 --- /dev/null +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -0,0 +1,75 @@ +"""分别发布 PD 运行态和 DP shared 资源，使用独立共享内存及读取计数。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal +from multiprocessing.reduction import ForkingPickler + +import torch.multiprocessing as mp +from filelock import FileLock + +from lightllm.utils.device_utils import kv_trans_use_p2p +from lightllm.utils.dist_utils import get_current_rank_in_node +from lightllm.utils.shm_utils import create_or_link_shm, get_service_shm_name + +if TYPE_CHECKING: + import torch + from lightllm.common.att_resources.att_resource_manager import AttResourceManager + + +def publish_att_resources( + resources: AttResourceManager, + request_mapping: torch.Tensor, + *, + kind: Literal["pd", "dp"], + reader_count: int, +) -> None: + """按用途选择导出副本；每个读者独立消费对应类别的 IPC 引用。""" + assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" + if reader_count < 0: + raise ValueError("storage reader count cannot be negative") + if reader_count == 0: + return + if kv_trans_use_p2p(): + from .mode_backend.pd.p2p_fix import reduce_tensor + + mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ + count = reader_count + if kind == "pd": + ipc_resources = resources.create_pd_ipc_copy() + else: + ipc_resources = resources.create_dp_ipc_copy() + with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): + # Each importer consumes its own CUDA IPC reference counters. + payloads = [ForkingPickler.dumps((ipc_resources, request_mapping)).tobytes() for _ in range(count)] + size = len(payloads[0]) + if any(len(payload) != size for payload in payloads): + raise RuntimeError("rank storage exports have inconsistent sizes") + shm = create_or_link_shm( + name=f"rank_storage_{kind}_{get_current_rank_in_node()}", expected_size=8 + count * size, force_mode="create" + ) + shm.buf[:4] = count.to_bytes(4, "little") + shm.buf[4:8] = size.to_bytes(4, "little") + for index, payload in enumerate(payloads): + shm.buf[8 + index * size : 8 + (index + 1) * size] = payload + shm.close() + + +def load_att_resources( + rank_in_node: int, *, kind: Literal["pd", "dp"], +) -> tuple[AttResourceManager, torch.Tensor]: + """读取指定用途的资源；DP 接收方随后通过 manager 统一注册共享页池。""" + assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" + with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): + shm = create_or_link_shm(name=f"rank_storage_{kind}_{rank_in_node}", expected_size=-1, force_mode="link") + try: + count = int.from_bytes(shm.buf[:4], "little") + size = int.from_bytes(shm.buf[4:8], "little") + if count <= 0: + raise RuntimeError("rank storage IPC exports are exhausted") + payload = bytes(shm.buf[8 + (count - 1) * size : 8 + count * size]) + resources, request_mapping = ForkingPickler.loads(payload) + shm.buf[:4] = (count - 1).to_bytes(4, "little") + return resources, request_mapping + finally: + shm.close() diff --git a/lightllm/server/router/req_queue/base_queue.py b/lightllm/server/router/req_queue/base_queue.py index 8d4ac38085..b306a5e904 100644 --- a/lightllm/server/router/req_queue/base_queue.py +++ b/lightllm/server/router/req_queue/base_queue.py @@ -3,7 +3,7 @@ from lightllm.utils.infer_utils import calculate_time from ..batch import Batch, Req from lightllm.server.core.objs import FinishStatus -from lightllm.utils.config_utils import get_fixed_kv_len, get_running_max_req_size_per_dp +from lightllm.utils.config_utils import get_running_max_req_size_per_dp from lightllm.server.core.objs import StartArgs from lightllm.utils.log_utils import init_logger @@ -18,11 +18,7 @@ def __init__(self, args: StartArgs, router, dp_index, dp_size_in_node) -> None: from lightllm.server.router.manager import RouterManager self.router: RouterManager = router - # max_total_token_num - get_fixed_kv_len() 是为了减去被特定 - # 推理模式预先占用了部分token kv 资源，这会导致整体可用的kv 资源 - # 在极端情况下减少，在非特定模式下，get_fixed_kv_len() 返回的都是 - # 0， 不会有任何影响。 - self.max_total_tokens = args.max_total_token_num - get_fixed_kv_len() + self.max_total_tokens = args.max_total_token_num self.running_max_req_size = get_running_max_req_size_per_dp(args) self.waiting_req_list: List[Req] = [] # List of queued requests self.router_token_ratio = args.router_token_ratio # ratio to determine whether the router is busy diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/att_resource_config.py new file mode 100644 index 0000000000..be0b66ab44 --- /dev/null +++ b/lightllm/utils/att_resource_config.py @@ -0,0 +1,215 @@ +"""Shared State definitions for model initialization and CPU page metadata.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import hashlib +import json +from pathlib import Path + +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.full_att import ( + FullAttentionConfig, FullAttentionState, create_gqa_state, create_mla_state, +) +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +def load_static_fp8_calibration( + args: StartArgs | EasyDict, + *, + per_head: bool, + layers: int, + heads: int, + tp_rank: int = 0, + tp_world_size: int = 1, + layer_start: int | None = None, +): + """Load and validate the same fixed parameters for model and cache metadata. + + Returned tensors remain on CPU so validation and capacity planning precede + GPU allocation. TP replicas select the same calibrated head values. + """ + path = args.kv_quant_calibration_config_path + if path is None: + if per_head: + raise ValueError("fp8kv_sph requires kv_quant_calibration_config_path with q_calibration") + return {"scale": "unit"}, torch.ones((layers, 2), dtype=torch.float32), None + with open(path) as f: + cfg = json.load(f) + stored_layers = cfg["num_layers"] + draft_layers = get_added_mtp_kv_layer_num() if layer_start is None else 0 + if layer_start is not None: + if layer_start < 0 or layer_start + layers > stored_layers: + raise ValueError("calibration does not cover this model's attention layers") + elif per_head and ("num_target_layers" in cfg or "num_draft_layers" in cfg): + target, draft = cfg["num_target_layers"], cfg["num_draft_layers"] + if ( + draft < 0 + or target + draft != stored_layers + or target != layers - draft_layers + or (draft_layers > 0 and draft != draft_layers) + ): + raise ValueError("calibration target/draft layer layout does not match the runtime") + elif stored_layers != layers: + raise ValueError(f"calibration num_layers {stored_layers} does not match runtime {layers}") + + if not per_head: + from lightllm.utils.config_utils import get_model_architectures + + if cfg["architectures"] != get_model_architectures(args.model_dir): + raise ValueError("calibration architecture does not match the model") + mode = "per_head" if per_head else "per_tensor" + if cfg.get("quant_type", mode) != mode: + raise ValueError("calibration quantization mode does not match the State") + limits = torch.finfo(torch.float8_e4m3fn) + if cfg.get("qmin", limits.min) != limits.min or cfg.get("qmax", limits.max) != limits.max: + raise ValueError("calibration range does not match float8_e4m3fn") + + def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]): + tensor = torch.tensor(values["scales"], dtype=torch.float32) + if list(values["scales_shape"]) != list(shape) or list(tensor.shape) != list(shape): + raise ValueError(f"calibration scales and scales_shape must match {shape}") + if not bool((torch.isfinite(tensor) & (tensor > 0)).all()): + raise ValueError("calibration scales must be finite and positive") + start = layer_start or 0 + return tensor[start : start + layers] + + def shard(tensor: torch.Tensor, calibrated_heads: int): + if calibrated_heads <= 0 or (heads * tp_world_size) % calibrated_heads: + raise ValueError("calibration heads do not cover the model's KV heads") + expanded = tensor.repeat_interleave(heads * tp_world_size // calibrated_heads, dim=-1) + return expanded[:, tp_rank * heads : (tp_rank + 1) * heads].contiguous() + + if per_head: + calibrated_heads = cfg["num_head"] + all_scales = scale_tensor(cfg, (stored_layers, 2 * calibrated_heads)) + k, v = all_scales.split(calibrated_heads, dim=-1) + scales = torch.cat((shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) + q_cfg = cfg.get("q_calibration") + if not q_cfg: + raise ValueError("fp8kv_sph calibration config requires q_calibration") + q_scales = shard(scale_tensor(q_cfg, (stored_layers, q_cfg["num_head"])), q_cfg["num_head"]) + else: + scales, q_scales = scale_tensor(cfg, (stored_layers, 2)), None + return cfg, scales, q_scales + + +def build_full_attention_state( + args: StartArgs | EasyDict, + *, + heads: int, + dim: int, + layer_ids: Sequence[int], + mla: bool = False, + packed: bool = False, + dtype: torch.dtype | None = None, + replicated: bool = False, + tp_rank: int = 0, + tp_world_size: int = 1, + calibration_layer_start: int | None = None, +) -> FullAttentionState: + """将模型和启动配置转换为算法 create 参数，返回容量规划所需的 meta State。""" + dtype = get_torch_dtype(args.data_type) if dtype is None else dtype + config = FullAttentionConfig( + layer_ids=tuple(layer_ids), + global_head_num=heads, + head_dim=dim, + dtype=dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + replicated=replicated or mla or packed, + ) + if mla or packed: + return create_mla_state( + config=config, kv_type=args.llm_kv_type, packed=packed, page_size=args.page_size, device="meta", + ) + scales = q_scales = None + calibration_id = "" + if args.llm_kv_type in ("fp8kv_sph", "fp8kv_spt"): + calibration, scales, q_scales = load_static_fp8_calibration( + args, + per_head=args.llm_kv_type == "fp8kv_sph", + layers=len(layer_ids), + heads=config.head_num, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + layer_start=calibration_layer_start, + ) + calibration_id = hashlib.sha256(json.dumps(calibration, sort_keys=True).encode()).hexdigest() + return create_gqa_state( + config=config, kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + scales=scales, q_scales=q_scales, calibration_id=calibration_id, + ) + + +def build_cpu_cache_layouts(): + """使用启动流程相同的 meta State 定义，由管理器统一生成 CPU 服务页布局。""" + from transformers.configuration_utils import PretrainedConfig + from lightllm.models import get_draft_model_class + from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + from lightllm.common.basemodel.basemodel_utils import ( + create_full_attention_definition, + create_bounded_attention_definition, + get_resource_config, + ) + from lightllm.utils.llm_utils import get_llm_model_class + + args = get_env_start_args() + main_config, _ = PretrainedConfig.get_config_dict(args.model_dir) + models = [("main", get_llm_model_class(), main_config)] + if args.mtp_mode is not None and args.mtp_mode not in ("vanilla_no_att", "eagle_no_att"): + count = args.mtp_step if args.mtp_mode == "vanilla_with_att" else 1 + assert args.mtp_draft_model_dir is not None and len(args.mtp_draft_model_dir) >= count + assert 0 < count <= MAX_DRAFT_MODELS, "at most 36 draft models are supported" + for index, directory in enumerate(args.mtp_draft_model_dir[:count]): + config, _ = PretrainedConfig.get_config_dict(directory) + models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) + + manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 + layer_start = calibration_start = 0 + for role, model_class, config in models: + count = get_resource_config(config, model_class=model_class)["num_hidden_layers"] + layer_ids = tuple(range(layer_start, layer_start + count)) + full = create_full_attention_definition( + model_class, config, args, model_layer_ids=layer_ids, calibration_layer_start=calibration_start + ) + bounded = create_bounded_attention_definition( + model_class, config, args, model_layer_ids=layer_ids + ) + for state in (full, bounded): + if state is None: + continue + manager.set_att_state(role, state) + layer_start += count + calibration_start += len(full.config.layer_ids) if full is not None else 0 + + manager.finish_state_creation() + if manager.get_bounded_attention_states(): + assert ( + args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + ), "hybrid CPU pages must match the large state-page token boundary" + return manager.get_cpu_transfer_layouts() + + +def cpu_cache_namespace(): + """按 main/draft 模型路径及文件版本隔离磁盘缓存。""" + args = get_env_start_args() + directories = [args.model_dir, *(args.mtp_draft_model_dir or [])] + revisions = [] + for directory in directories: + directory = Path(directory).resolve() + artifacts = [] + for pattern in ("*.safetensors", "*.bin", "*.index.json", "config.json"): + for path in sorted(directory.glob(pattern)): + stat = path.stat() + artifacts.append((path.name, stat.st_size, stat.st_mtime_ns)) + revisions.append((str(directory), artifacts)) + return hashlib.sha256(repr(revisions).encode()).hexdigest() diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 0903b9b33c..8927374915 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -99,9 +99,9 @@ def _validate_flashqla(): from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) - from lightllm.common.state_cache_manager import LinearAttCacheConfig + from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig - linear_config = LinearAttCacheConfig.load_from_args() + linear_config = Qwen3NextResourceConfig.load_from_args() num_k_heads = linear_config.num_linear_k_heads num_v_heads = linear_config.num_linear_v_heads head_k_dim = linear_config.head_linear_k_dim diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 7b258959e9..9457e77391 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -255,42 +255,6 @@ def get_hidden_size(model_path: str) -> Optional[int]: return None -@lru_cache(maxsize=None) -def get_num_key_value_heads(model_path: str) -> int: - num_key_value_heads = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_key_value_heads"]) - if isinstance(num_key_value_heads, int): - return num_key_value_heads - return None - - -@lru_cache(maxsize=None) -def get_num_attention_heads(model_path: str) -> int: - num_attention_heads = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_attention_heads"]) - if isinstance(num_attention_heads, int): - return num_attention_heads - return None - - -@lru_cache(maxsize=None) -def get_head_dim(model_path: str) -> int: - head_dim = _get_config_llm_keyvalue(model_path=model_path, key_name=["head_dim"]) - if isinstance(head_dim, int): - return head_dim - - # calcu head_dim - head_dim = get_hidden_size(model_path=model_path) // get_num_attention_heads(model_path=model_path) - - return head_dim - - -@lru_cache(maxsize=None) -def get_layer_num(model_path: str) -> int: - num_hidden_layers = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_hidden_layers"]) - if isinstance(num_hidden_layers, int): - return num_hidden_layers - return None - - def get_eos_token_ids(model_path: str) -> Optional[List[int]]: # gemma4 special eos_token_id try: @@ -408,19 +372,6 @@ def get_dtype(model_path: str): return torch_dtype -@lru_cache(maxsize=None) -def get_fixed_kv_len(): - start_args = get_env_start_args() - model_cfg = get_config_json(start_args.model_dir) - if "prompt_cache_token_ids" in model_cfg: - fixed_kv_len = len(model_cfg["prompt_cache_token_ids"]) - # 固定 KV 最终会插入 radix cache，只加载完整的模型 KV 页面；不足一页 - # 的尾部直接截断，因此 router 也只扣除实际常驻的页面容量。 - return fixed_kv_len // start_args.page_size * start_args.page_size - else: - return 0 - - @lru_cache(maxsize=None) def has_vision_module(model_path: str) -> bool: try: diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 4f93317d61..923fb9aada 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -12,23 +12,12 @@ get_env_start_args, enable_huge_page, enable_cpu_cache_numa_interleave, - get_llm_data_type, - get_added_mtp_kv_layer_num, ) from lightllm.utils.log_utils import init_logger -from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num, is_hybrid_att_model -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.common.kv_cache_mem_manager import ( - MemoryManager, - PPLINT8KVMemoryManager, - PPLINT4KVMemoryManager, - Deepseek2MemoryManager, -) from typing import List, Tuple, Optional from tqdm import tqdm from lightllm.utils.dist_utils import get_current_device_id -from lightllm.common.state_cache_manager import get_hybrid_cache_config logger = init_logger(__name__) @@ -62,67 +51,22 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir) - mem_manager_class = None if is_hybrid_model else select_mem_manager_class() - if is_hybrid_model: - hybrid_config = get_hybrid_cache_config() - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=1, - layer_num=1, - num_heads=1, - head_dim=hybrid_config.get_cpu_cache_big_page_bytes(), - data_type=torch.uint8, - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is Deepseek2MemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=1, - head_dim=512 + 64, - data_type=get_llm_data_type(), - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is MemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=get_num_key_value_heads(args.model_dir) * 2, - head_dim=get_head_dim(args.model_dir), - data_type=get_llm_data_type(), - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is PPLINT8KVMemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=get_num_key_value_heads(args.model_dir) * 2, - head_dim=get_head_dim(args.model_dir), - data_type=torch.int8, - scale_head_dim=get_head_dim(args.model_dir) // 8, - scale_data_type=get_llm_data_type(), - ) - else: - logger.error(f"not support mem manager: {mem_manager_class} for cpu kv cache") - raise Exception(f"not support mem manager: {mem_manager_class} for cpu kv cache") - - if args.mtp_mode is not None: - # TODO 可能会存在不同mtp模式的精度问题 - if not is_hybrid_model: - # 对于非 hybrid 模型，需要额外增加 mtp 的 kv 层数， - # 对于 hybrid 模型，如 qwen 3.5 mtp，已经将 kv 数据 - # 打包成一个块了，所以不需要额外增加，其 layer_num 一直都保持为 1 - cpu_cache_meta.layer_num += get_added_mtp_kv_layer_num() + from lightllm.utils.att_resource_config import build_cpu_cache_layouts, cpu_cache_namespace + + layouts = build_cpu_cache_layouts() + page_bytes = 0 + for layout in layouts.values(): + page_bytes = max(page_bytes, layout.byte_end) + # State 区间按 8 字节对齐，整个 CPU 物理页再按 16 字节对齐。 + page_bytes = (page_bytes + 15) // 16 * 16 + cpu_cache_meta = CpuKVCacheMeta( + page_num=0, + page_bytes=page_bytes, + cache_namespace=cpu_cache_namespace(), + ) cpu_cache_page_num = int( - (args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.calcu_one_page_size()) + (args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.page_bytes) ) cpu_cache_meta.page_num = cpu_cache_page_num @@ -133,36 +77,19 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": @dataclasses.dataclass class CpuKVCacheMeta: + """External byte-page storage, independent of model tensor dimensions.""" + page_num: int - token_page_size: int - layer_num: int - num_heads: int - head_dim: int - data_type: torch.dtype - scale_head_dim: int - scale_data_type: torch.dtype - - def calcu_size(self): - return self.page_num * self.calcu_one_page_size() - - def calcu_one_page_size(self): - return ( - self.token_page_size - * self.layer_num - * self.num_heads - * (self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize) - ) + page_bytes: int + cache_namespace: str = "" + + def __post_init__(self): + if self.page_num < 0 or self.page_bytes <= 0: + raise ValueError("CPU cache requires a nonnegative page count and positive page size") - def get_merged_head_dim(self): - """ - 返回将head_dim 和 scale_head_dim 看成融合成一个head_dim时候, head_dim的长度。 - """ - assert ( - self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize - ) % self.data_type.itemsize == 0 - return ( - self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize - ) // self.data_type.itemsize + @property + def size_bytes(self): + return self.page_num * self.page_bytes @lru_cache(maxsize=None) diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py new file mode 100644 index 0000000000..8a55f86145 --- /dev/null +++ b/test/att_resources/validate_model_cache.py @@ -0,0 +1,322 @@ +"""Validate resource restoration against an uninterrupted, real model decode. + +Run from the repository root with local model weights, for example: + PYTHONPATH=. python test/att_resources/validate_model_cache.py \ + --model_dir /models/Qwen3.5-0.8B --report /tmp/hybrid-cache.json + +This checks GPU pages, CPU pages, DP copies and PD buffers on one GPU. Network +transport is tested separately. Each CPU page captures a real forward boundary. +Decode uses CUDA graphs by default; pass --disable_cudagraph to check eager mode. +Native MTP can be checked with --mtp_mode vanilla_with_att --mtp_step 1 and +--mtp_draft_model_dir pointing to the checkpoint containing its draft weights. +MTP uses identical candidate packets with varying accepted widths, not sampling. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import json +from pathlib import Path +import socket +import sys + +import numpy as np +import torch +import torch.distributed as dist +from transformers import PretrainedConfig + +from lightllm.distributed import dist_group_manager +from lightllm.models import get_model +from lightllm.server.api_cli import make_argument_parser +from lightllm.utils.dist_utils import init_distributed_env +from lightllm.utils.envs_utils import get_unique_server_name, set_env_start_args, set_unique_server_name +from lightllm.utils.service_shm_cleanup import ServiceShmCleanup + +if TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Sequence + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "benchmark" / "static_inference")) +from static_benchmark import ( + StaticBenchmarkExecutor, + TokenSource, + add_static_benchmark_args, + build_cases, + build_model_kvargs, + init_mtp_draft_models, + normalize_args, +) + + +@torch.inference_mode() +def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpPartBaseModel] = ()): + executor = StaticBenchmarkExecutor(args, model, draft_models, TokenSource(args)) + resources = model.att_resource_manager + (full_state,) = resources.get_full_attention_states("main") + for draft in draft_models: + assert draft.att_resource_manager is resources + assert not hasattr(draft, "full_att_state") + (draft_state,) = resources.get_full_attention_states("draft0") + assert draft_state is not full_state + assert draft_state.capacity == full_state.capacity == resources.full_att_allocator.size + assert draft.req_manager is model.req_manager + prefix_len, page_tokens = args.input_len, args.cpu_cache_token_page_size + if prefix_len % page_tokens: + raise ValueError("input_len must contain complete CPU pages") + pages = list(range(prefix_len // page_tokens)) + lengths = [(page + 1) * page_tokens for page in pages] + hybrid = model.is_hybrid_att_model + big = resources.bounded_att_big_page_allocator + source_slots = big.alloc(len(pages)) if hybrid else None + allocated_slots = list(source_slots) if hybrid else [] + source_req = executor._alloc_req_indexes(1) + resources.reset_bounded_att_state(source_req.tolist()) + resources.reset_attention_tail_state(source_req.tolist()) + tokens = np.random.default_rng(42).integers(100, 10000, (1, prefix_len), dtype=np.int64) + for page, length in zip(pages, lengths): + batch = executor._make_prefill_input(tokens[:, page * page_tokens : length], source_req, page * page_tokens) + output = model.forward(batch) + if draft_models: + executor._fill_mtp_prefill_kv(None, batch, output, executor._argmax_ids(output.logits)) + if hybrid: + resources.save_bounded_att_pages(source_req.tolist(), page_kind="big", slot_ids=[source_slots[page]]) + torch.cuda.synchronize() + next_token = output.logits.argmax(dim=-1).cpu() + source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() + + cpu = torch.zeros((len(pages), resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + resources.write_req_to_cpu( + source_indices, + cpu, + pages, + layouts=resources.get_cpu_transfer_layouts(), + page_kind="big" if hybrid else None, + slot_ids=source_slots, + ) + # Host-side checkpoint repacking consumes the CPU page on the next call. + torch.cuda.synchronize() + full_att_layouts = resources.get_pd_full_att_transfer_layouts() + runtime_layouts = resources.get_pd_bounded_att_or_tail_transfer_layouts() + pd = {"kv": torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda")} + if runtime_layouts: + pd["att_state"] = torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + resources.write_to_pd_buffer( + pd["kv"], + pd_page_index=0, + kind="kv", + layouts=full_att_layouts, + token_indices=source_indices, + ) + if "att_state" in pd: + resources.write_to_pd_buffer( + pd["att_state"], + pd_page_index=0, + kind="att_state", + layouts=runtime_layouts, + req_idx=int(source_req[0]), + state_token_len=prefix_len, + ) + + targets = {} + for kind in ("gpu_pages", "cpu_pages", "dp_copy", "pd_buffer"): + request = executor._alloc_req_indexes(1) + resources.reset_bounded_att_state(request.tolist()) + resources.reset_attention_tail_state(request.tolist()) + executor._ensure_req_kv_capacity(request, torch.tensor([prefix_len], dtype=torch.int32)) + indices = model.req_manager.req_to_token_indexs[int(request[0]), :prefix_len].clone() + assert not set(indices.tolist()) & set(source_indices.tolist()) + if kind == "gpu_pages": + resources.copy_full_att_tokens(source_indices, indices) + if hybrid: + resources.restore_bounded_att_pages(request.tolist(), page_kind="big", slot_ids=[source_slots[-1]]) + resources.reset_attention_tail_state(request.tolist()) + elif kind == "cpu_pages": + slots = big.alloc(len(pages)) if hybrid else None + if hybrid: + allocated_slots.extend(slots) + resources.read_req_from_cpu( + cpu, + pages, + indices, + layouts=resources.get_cpu_transfer_layouts(), + big_page_slot_ids=slots, + req_idx=int(request[0]) if getattr(resources, "cpu_tail_restore", False) else None, + ) + if hybrid: + resources.restore_bounded_att_pages(request.tolist(), page_kind="big", slot_ids=[slots[-1]]) + elif kind == "dp_copy": + resources.copy_from_dp( + resources, + source_indices, + indices, + mappings=resources.dp_copy_mappings(resources), + src_req_indices=source_req.tolist() if hybrid else None, + dst_req_indices=request.tolist() if hybrid else None, + ) + else: + resources.read_from_pd_buffer( + pd["kv"], + pd_page_index=0, + kind="kv", + layouts=full_att_layouts, + token_indices=indices, + ) + if "att_state" in pd: + resources.read_from_pd_buffer( + pd["att_state"], + pd_page_index=0, + kind="att_state", + layouts=runtime_layouts, + req_idx=int(request[0]), + state_token_len=prefix_len, + ) + targets[kind] = request + + errors = {kind: [] for kind in targets} + draft_errors = {kind: [] for kind in targets} if draft_models else None + generated = [] + accepted_widths = [] + sequence_length = prefix_len + step_width = args.mtp_step + 1 if draft_models else 1 + for step in range(args.output_len): + generated.append(int(next_token[0])) + accepted_width = step % step_width + 1 + accepted_widths.append(accepted_width) + + def decode(request: torch.Tensor): + request_rows, mtp_index = executor._build_mtp_decode_index_tensors(request, step_width) + lengths = executor._build_mtp_seq_len(torch.tensor([sequence_length], dtype=torch.int32), step_width) + batch = executor._make_decode_input( + step_width, + request_rows, + mtp_index, + lengths, + next_token.repeat(step_width), + int(lengths[-1]), + ) + output = model.forward(batch) + logits = [output.logits.clone()] + for draft in draft_models: + batch.input_ids = output.logits.argmax(dim=-1) + batch.mtp_draft_input_hiddens = output.mtp_collector.spec_hidden + output = draft.forward(batch) + logits.append(output.logits.clone()) + if draft_models and hybrid: + ( + resources.main_kda_state if hasattr(resources, "main_kda_state") else resources.main_gdn_state + ).commit_bounded_att_state(request, [accepted_width]) + return logits + + reference = decode(source_req) + assert all(torch.isfinite(logits).all() for logits in reference) + for kind, request in targets.items(): + actual = decode(request) + for restored_logits, reference_logits in zip(actual, reference): + torch.testing.assert_close(restored_logits, reference_logits, rtol=1e-4, atol=1e-4) + assert torch.equal(restored_logits.argmax(dim=-1), reference_logits.argmax(dim=-1)) + errors[kind].append(float((actual[0] - reference[0]).abs().max())) + if draft_models: + draft_errors[kind].append([float((a - b).abs().max()) for a, b in zip(actual[1:], reference[1:])]) + next_token = reference[0][accepted_width - 1 : accepted_width].argmax(dim=-1).cpu() + sequence_length += accepted_width + + torch.cuda.synchronize() + if hybrid: + big.free(allocated_slots) + executor._reset_model_cache() + assert resources.full_att_allocator.can_use_mem_size == resources.full_att_allocator.size + return { + "model": args.model_dir, + "model_class": type(model).__name__, + "transport": "local_buffers", + "prefix_tokens": prefix_len, + "cpu_pages": len(pages), + "decode_steps": args.output_len, + "generated_token_ids": generated, + "max_logit_errors": errors, + "draft_models": [type(draft).__name__ for draft in draft_models], + "max_draft_logit_errors": draft_errors, + "accepted_widths": accepted_widths, + "allocator_reclaimed": True, + "cuda_graphs": not args.disable_cudagraph, + } + + +def main(): + parser = make_argument_parser() + add_static_benchmark_args(parser) + parser.add_argument("--report", required=True) + parser.set_defaults( + tp=1, + dp=1, + data_type="bfloat16", + batch_size=1, + max_batch_size=1, + input_len=128, + output_len=4, + batch_max_tokens=256, + max_total_token_num=2048, + running_max_req_size=8, + graph_max_batch_size=2, + max_req_total_len=512, + page_size=16, + cpu_cache_token_page_size=64, + enable_cpu_cache=True, + linear_att_hash_page_size=64, + linear_att_page_block_num=1, + linear_att_cache_size=16, + llm_prefill_att_backend=["fa3", "triton"], + llm_decode_att_backend=["fa3", "triton"], + disable_cudagraph=False, + pd_kv_page_size=2048, + ) + parsed = parser.parse_args() + if parsed.tp != 1 or parsed.dp != 1 or parsed.mtp_mode not in (None, "vanilla_with_att"): + parser.error("this validation uses TP=1, DP=1 and optionally native vanilla_with_att MTP") + with socket.socket() as port: + port.bind(("127.0.0.1", 0)) + parsed.nccl_port = port.getsockname()[1] + args = normalize_args(parsed, build_cases(parsed)) + set_unique_server_name(args) + set_env_start_args(args) + torch.set_num_threads(4) + torch.manual_seed(42) + try: + kvargs = build_model_kvargs(args, 0) + init_distributed_env(kvargs) + dist_group_manager.create_groups(group_size=1) + config, _ = PretrainedConfig.get_config_dict(args.model_dir) + model, _ = get_model(config, kvargs) + draft_models = init_mtp_draft_models(args, kvargs, model) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.att_resources import AttResourceManager + + models = (model, *draft_models) + assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + for item in models: + item.prepare_att_resources() + manager: AttResourceManager = model.att_resource_manager + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.initialize_resources( + model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + report = validate(model, args, draft_models) + Path(args.report).write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2), flush=True) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + ServiceShmCleanup(get_unique_server_name()).cleanup_service_resources() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 279840d5af..1fd3a2cc24 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -568,8 +568,8 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T req_id = int(req_id) max_target_len_by_req[req_id] = max(max_target_len_by_req.get(req_id, 0), int(target_len)) - mem_manager = self.model.req_manager.mem_manager - page_size = mem_manager.page_size + allocator = self.model.att_resource_manager.full_att_allocator + page_size = allocator.page_size for req_id, target_len in max_target_len_by_req.items(): new_hold_kv_len = align_up(target_len, page_size) old_hold_kv_len = self._hold_kv_len_by_req.get(req_id, 0) @@ -577,7 +577,7 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T continue alloc_token_num = new_hold_kv_len - old_hold_kv_len - mem_indexes = mem_manager.alloc(alloc_token_num) + mem_indexes = allocator.alloc(alloc_token_num) self.model.req_manager.req_to_token_indexs[req_id, old_hold_kv_len:new_hold_kv_len].copy_(mem_indexes) self._hold_kv_len_by_req[req_id] = new_hold_kv_len @@ -729,7 +729,7 @@ def _alloc_req_indexes(self, batch_size: int) -> torch.Tensor: return torch.tensor(req_indexes, dtype=torch.int32, device="cpu") def _reset_model_cache(self): - self.model.mem_manager.free_all() + self.model.att_resource_manager.full_att_allocator.free_all() self.model.req_manager.free_all() self._hold_kv_len_by_req.clear() torch.cuda.synchronize() @@ -1182,13 +1182,12 @@ def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) if args.mtp_mode not in MTP_MODES: return [] - os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" draft_models = [] for draft_dir in args.mtp_draft_model_dir: mtp_cfg, _ = PretrainedConfig.get_config_dict(draft_dir) mtp_kvargs = { "weight_dir": draft_dir, - "max_total_token_num": main_model.mem_manager.size, + "max_total_token_num": main_kvargs["max_total_token_num"], "load_way": main_kvargs["load_way"], "max_req_num": main_kvargs["max_req_num"], "max_seq_length": main_kvargs["max_seq_length"], @@ -1271,11 +1270,40 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) torch.cuda.empty_cache() model, _ = get_model(model_cfg, model_kvargs) - cases = resolve_batch_max_prefill_cases(args, cases, model.mem_manager.size) - cases = resolve_profile_decode_cases(args, cases, model.mem_manager.size) + draft_models = init_mtp_draft_models(args, model_kvargs, model) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.att_resources import AttResourceManager + + models = (model, *draft_models) + assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + for item in models: + item.prepare_att_resources() + manager: AttResourceManager = model.att_resource_manager + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.initialize_resources( + model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + cases = resolve_batch_max_prefill_cases(args, cases, model.att_resource_manager.full_att_allocator.size) + cases = resolve_profile_decode_cases(args, cases, model.att_resource_manager.full_att_allocator.size) if defer_cudagraph: init_deferred_cudagraph(args, cases, model_kvargs, model) - draft_models = init_mtp_draft_models(args, model_kvargs, model) + for draft in draft_models: + batch_size = args.graph_max_batch_size + if args.enable_decode_microbatch_overlap: + batch_size //= 2 + batch_size *= draft.mtp_manager.get_decode_tokens_per_request(True) + draft.graph_max_batch_size = draft._align_decode_batch_size(batch_size) + draft.disable_cudagraph = False + draft._init_att_backend() + draft._init_att_backend1() + draft._init_cudagraph() token_source = TokenSource(args) executor = StaticBenchmarkExecutor(args, model, draft_models, token_source) diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index d59142e09f..f28f549022 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -1,276 +1,118 @@ -""" -Speed benchmark for copy_cpu_cache_to_kv_buffer in linear_att_cpu_cache_copy.py. +"""Measure the public CPU-page restore path for one TP shard of a hybrid model. -Test configuration (matching the user's LinearAttCacheConfig): - tp_world_size=8, full_att_all_num_kv_heads=2, full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, full_att_head_dim=256, - num_linear_k_heads=2, num_linear_v_heads=8, - head_linear_k_dim=128, head_linear_v_dim=128, - conv_kernel_size=4, 36 linear layers and 12 full-attention layers, - conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=(linear_attention, linear_attention, linear_attention, full_attention) * 12 +Run from the repository root, for example: + PYTHONPATH=. python test/cpu_cache_kernel/test_speed.py --tokens-per-page 16384 --iterations 100 +The reported time includes validation, KV transfers, and CPU state-page repacking. """ -import os -import json +import argparse +import statistics import time -import triton -import torch -from easydict import EasyDict - -# --------------------------------------------------------------------------- -# Step 0 – set up environment args BEFORE any import that calls -# get_env_start_args() / LinearAttCacheConfig.load_from_args(). -# --------------------------------------------------------------------------- -_env_args = { - "cpu_cache_token_page_size": 2048 * 8, # big_page_token_num - "linear_att_hash_page_size": 2048, - "linear_att_page_block_num": 8, # 512 * 1 == 512 - "data_type": "bfloat16", - "linear_att_ssm_data_type": "float32", - "model_dir": "/tmp/fake_model", # dummy – not used when config is built directly - "tp": 8, - "dp": 1, - "running_max_req_size": 2048, - "enable_cpu_cache": True, -} -os.environ["LIGHTLLM_START_ARGS"] = json.dumps(_env_args) - -# --------------------------------------------------------------------------- -# Step 1 – build LinearAttCacheConfig directly (avoids needing a real model dir) -# --------------------------------------------------------------------------- -from lightllm.common.state_cache_manager import LinearAttCacheConfig - -linear_config = LinearAttCacheConfig( - tp_world_size=8, - full_att_all_num_kv_heads=2, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=256, - global_linear_k_heads=16, - global_linear_v_heads=64, - num_linear_k_heads=2, - num_linear_v_heads=8, - head_linear_k_dim=128, - head_linear_v_dim=128, - conv_kernel_size=4, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - layer_types=("linear_attention", "linear_attention", "linear_attention", "full_attention") * 12, -) -print(f"LinearAttCacheConfig:\n{linear_config}\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 2 – derive sizes from the config -# --------------------------------------------------------------------------- -big_page_token_num = _env_args["cpu_cache_token_page_size"] # 512 -full_att_layer_num = linear_config.get_main_model_full_att_layer_num() # 12 - -full_att_bytes = linear_config.get_cpu_cache_full_att_bytes() # per big page -conv_bytes = linear_config.get_cpu_cache_conv_bytes() -ssm_bytes = linear_config.get_cpu_cache_ssm_bytes() -total_bytes = full_att_bytes + conv_bytes + ssm_bytes -print( - f"Per-page bytes full_att={full_att_bytes:,} conv={conv_bytes:,} " f"ssm={ssm_bytes:,} total={total_bytes:,}", - flush=True, -) -total_bytes = linear_config.get_cpu_cache_big_page_bytes() - -# --------------------------------------------------------------------------- -# Step 3 – allocate tensors -# --------------------------------------------------------------------------- -grid_num = 8 -PAGE_NUM = 1 # number of big pages to copy per call -SEQ_LEN = 2048 * 8 # total sequence length in gpu_full_att_kv_state dim-1 -BIG_PAGE_COUNT = PAGE_NUM # big_page_buffer_ids length == page_indexes length - -# --- GPU tensors --- -mem_indexes = torch.arange(0, big_page_token_num * PAGE_NUM, dtype=torch.int64, device="cpu") -big_page_buffer_ids = torch.arange(0, BIG_PAGE_COUNT, dtype=torch.int64, device="cpu") -page_indexes = torch.arange(0, PAGE_NUM, dtype=torch.int32, device="cpu") - -gpu_full_att_kv_state = torch.empty( - ( - full_att_layer_num, - SEQ_LEN, - 2 * max(1, linear_config.full_att_num_kv_heads // linear_config.tp_world_size), - linear_config.full_att_head_dim, - ), - dtype=linear_config.full_att_dtype, - device="cuda", -) - -# --- CPU tensors --- -buffer_count = triton.cdiv(SEQ_LEN, big_page_token_num) + 2 # matches Qwen3NextMemManager - - -conv_shape = linear_config.get_conv_state_shape() -cpu_kv_conv_state = torch.empty( - (buffer_count, linear_config.linear_layer_num, *conv_shape), - dtype=linear_config.conv_state_dtype, - device="cpu", - pin_memory=True, -) - -ssm_shape = linear_config.get_ssm_state_shape() # (num_linear_v_heads, head_linear_k_dim, head_linear_v_dim) -cpu_kv_ssm_state = torch.empty( - (buffer_count, linear_config.linear_layer_num, *ssm_shape), - dtype=linear_config.ssm_state_dtype, - device="cpu", - pin_memory=True, -) - -# conv_shape = linear_config.get_conv_state_shape() -# cpu_kv_conv_state = torch.empty( -# (buffer_count, linear_config.linear_layer_num, *conv_shape), -# dtype=linear_config.conv_state_dtype, device="cuda", -# ) - -# ssm_shape = linear_config.get_ssm_state_shape() # (num_linear_v_heads, head_linear_k_dim, head_linear_v_dim) -# cpu_kv_ssm_state = torch.empty( -# (buffer_count, linear_config.linear_layer_num, *ssm_shape), -# dtype=linear_config.ssm_state_dtype, device="cuda", -# ) +import torch -# cpu_cache_tensor: [page_num, 1, 1, 1, total_bytes] -cpu_cache_tensor = torch.empty( - (PAGE_NUM, 1, 1, 1, total_bytes), - dtype=torch.uint8, - device="cpu", - pin_memory=True, -) - -# Move GPU tensors to CUDA -mem_indexes_cuda = mem_indexes.cuda(non_blocking=True) -big_page_buffer_ids_cuda = big_page_buffer_ids.cuda(non_blocking=True) -page_indexes_cuda = page_indexes.cuda(non_blocking=True) -gpu_full_att_kv_state = gpu_full_att_kv_state.cuda(non_blocking=True) - -torch.cuda.synchronize() -print("All tensors allocated and moved to GPU.\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 4 – import and warm-up the triton kernel -# --------------------------------------------------------------------------- -from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_cpu_cache_to_kv_buffer, -) - -print("Warming up …", flush=True) -copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.att_resources.full_att import ( + GQAFullAttentionState, + FullAttentionConfig, ) -torch.cuda.synchronize() -print("Warm-up done.\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 5 – benchmark -# --------------------------------------------------------------------------- -WARMUP_ITERS = 10 -BENCH_ITERS = 100 - -print(f"Benchmarking ({BENCH_ITERS} iterations, {PAGE_NUM} pages / {big_page_token_num} tokens each) …", flush=True) - -# Warm-up -for _ in range(WARMUP_ITERS): - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens-per-page", type=int, default=16384) + parser.add_argument("--pages", type=int, default=1) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--warmup", type=int, default=10) + args = parser.parse_args() + if min(args.tokens_per_page, args.pages, args.iterations) <= 0 or args.warmup < 0: + parser.error("page sizes and iterations must be positive; warmup must be nonnegative") + set_env_start_args(StartArgs(cpu_cache_token_page_size=args.tokens_per_page)) + cfg = Qwen3NextResourceConfig( + tp_world_size=8, + full_att_all_num_kv_heads=2, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=256, + global_linear_k_heads=16, + global_linear_v_heads=64, + num_linear_k_heads=2, + num_linear_v_heads=8, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + linear_layer_num=36, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=48, ) -torch.cuda.synchronize() - -# Timed runs -times = [] -for _ in range(BENCH_ITERS): + tokens = args.tokens_per_page * args.pages + full = GQAFullAttentionState( + FullAttentionConfig(tuple(range(12)), 1, 256, torch.bfloat16, 0, 8), + ) + full.initialize_storage(tokens) + resources = Qwen3NextAttResourceManager(cfg, request_capacity=1, mtp_step=0) + resources.bind_storage( + full, + FullAttKVAllocator(tokens), + big_page_capacity=args.pages + 2, + small_page_capacity=0, + ) + pages = list(range(args.pages)) + slots = resources.bounded_att_big_page_allocator.alloc(args.pages) + cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + # Populate every TP region through the production API before measuring restores. + full.k_storage.zero_() + full.v_storage.zero_() + resources.save_bounded_att_pages([0] * args.pages, page_kind="big", slot_ids=slots) torch.cuda.synchronize() - t0 = time.perf_counter() - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, + token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") + resources.write_req_to_cpu( + token_indices, + cpu_page_buffers, + pages, + layouts=resources.get_cpu_transfer_layouts(), + page_kind="big", + slot_ids=slots, ) torch.cuda.synchronize() - t1 = time.perf_counter() - times.append(t1 - t0) - -# --------------------------------------------------------------------------- -# Step 6 – report -# --------------------------------------------------------------------------- -import statistics - -times_ms = [t * 1e3 for t in times] -total_tokens = PAGE_NUM * big_page_token_num - -# Calculate head_scale_size (same logic as in copy_cpu_cache_to_kv_buffer) -if linear_config.full_att_all_num_kv_heads % linear_config.tp_world_size == 0: - head_scale_size = 1 -else: - head_scale_size = linear_config.tp_world_size // linear_config.full_att_all_num_kv_heads -# Each TP rank copies: -# - full_att_bytes / head_scale_size (full attention is sharded by head_scale_size) -# - conv_bytes / tp_world_size (conv state is sharded by tp_rank) -# - ssm_bytes / tp_world_size (ssm state is sharded by tp_rank) -full_att_bytes = linear_config.get_cpu_cache_full_att_bytes() -conv_bytes = linear_config.get_cpu_cache_conv_bytes() -ssm_bytes = linear_config.get_cpu_cache_ssm_bytes() - -bytes_per_page_per_tp = ( - full_att_bytes - * max(1, linear_config.full_att_all_num_kv_heads // linear_config.tp_world_size) - / linear_config.full_att_all_num_kv_heads - + conv_bytes // linear_config.tp_world_size - + ssm_bytes // linear_config.tp_world_size -) -total_bytes_copied = PAGE_NUM * bytes_per_page_per_tp - -print() -print("=" * 60) -print(f" copy_cpu_cache_to_kv_buffer speed benchmark") -print("=" * 60) -print(f" Pages / call : {PAGE_NUM}") -print(f" Tokens / page : {big_page_token_num}") -print(f" Total tokens / call : {total_tokens}") -print(f" Bytes / page (total) : {total_bytes:,}") -print(f" Bytes / page (per TP) : {bytes_per_page_per_tp:,}") -print(f" Total bytes / call : {total_bytes_copied:,} ({total_bytes_copied / 1024**3:.3f} GB)") -print(f" Iterations : {BENCH_ITERS}") -print(f" Mean latency : {statistics.mean(times_ms):.3f} ms") -print(f" Median latency : {statistics.median(times_ms):.3f} ms") -print(f" Std latency : {statistics.stdev(times_ms):.3f} ms") -print(f" Min latency : {min(times_ms):.3f} ms") -print(f" Max latency : {max(times_ms):.3f} ms") -print(f" Throughput (tokens/s) : {total_tokens / statistics.mean(times_ms) * 1e3:,.0f}") -print(f" Throughput (GB/s) : {total_bytes_copied / 1024**3 / statistics.mean(times_ms) * 1e3:.3f}") -print("=" * 60) + def restore(): + resources.read_req_from_cpu( + cpu_page_buffers, + pages, + token_indices, + layouts=resources.get_cpu_transfer_layouts(), + big_page_slot_ids=slots, + ) + + for _ in range(args.warmup): + restore() + torch.cuda.synchronize() + times = [] + for _ in range(args.iterations): + start = time.perf_counter() + restore() + torch.cuda.synchronize() + times.append(time.perf_counter() - start) + token_bytes = ( + full.k_storage[:, :tokens].numel() * full.k_storage.element_size() + + full.v_storage[:, :tokens].numel() * full.v_storage.element_size() + ) + pool = ( + resources.main_kda_state if hasattr(resources, "main_kda_state") else resources.main_gdn_state + ).big_pages_pool + state_bytes = args.pages * sum(t[0].numel() * t.element_size() for t in (pool.conv_state, pool.ssm_state)) + median = statistics.median(times) + print(f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; CPU bytes/page: {resources.get_cpu_page_bytes():,}") + print(f"Mean: {statistics.mean(times) * 1000:.3f} ms; median: {median * 1000:.3f} ms") + print(f"KV + state payload throughput: {(token_bytes + state_bytes) / median / 1024**3:.3f} GiB/s") + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/att_resources/conftest.py b/unit_tests/common/att_resources/conftest.py new file mode 100644 index 0000000000..aac6d6de02 --- /dev/null +++ b/unit_tests/common/att_resources/conftest.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import json +import os +from dataclasses import asdict + +import pytest + + +@pytest.fixture(autouse=True) +def allocator_environment(kv_allocator_environment: None, monkeypatch: pytest.MonkeyPatch): + """独立资源测试也需要启动参数，以选择页池的共享分配策略。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + if "LIGHTLLM_START_ARGS" not in os.environ: + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.fixture +def pd_args(monkeypatch: pytest.MonkeyPatch): + """为独立资源测试提供与生产路径一致的全局 PD 配置。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(pd_kv_page_size=8, cpu_cache_token_page_size=8)))) + get_env_start_args.cache_clear() + yield get_env_start_args() + get_env_start_args.cache_clear() + + +@pytest.fixture +def cpu_args(monkeypatch: pytest.MonkeyPatch): + """为独立资源测试提供与生产路径一致的全局 CPU 页配置。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(cpu_cache_token_page_size=8)))) + get_env_start_args.cache_clear() + yield get_env_start_args() + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py new file mode 100644 index 0000000000..f41e1d9924 --- /dev/null +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -0,0 +1,61 @@ +"""Concrete main/draft composition used to exercise manager transfer operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace +from functools import lru_cache +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.att_resources.full_att import MLAFullAttentionState +from lightllm.common.att_resources.bounded_att import GDNState +from lightllm.common.att_resources.att_tail import KPoolAttentionTailState +from unit_tests.common.att_resources.test_full_state import full_state +from unit_tests.common.att_resources.test_bounded_state import make_state +from unit_tests.common.att_resources.test_tail_state import make_tail + +if TYPE_CHECKING: + import torch + + +class MainDraftTestResourceManager(AttResourceManager): + def __init__(self, device: str | torch.device | int = "cpu", *, with_tail: bool = True, tail_pool_size: int = 4): + super().__init__(request_capacity=3) + self.set_att_state("main", make_state(device=device)) + draft_state = GDNState( + replace(self.main_kda_state.config, global_v_heads=2), + request_capacity=3, mtp_step=self.main_kda_state.mtp_step, device=device, + ) + draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + self.set_att_state("draft0", draft_state) + self.set_att_state("main", full_state(device=device)) + prepared_full_state = MLAFullAttentionState( + self.main_full_att_fp32_state.config, + page_size=self.main_full_att_fp32_state.page_size, device=device, + ) + prepared_full_state.initialize_storage(self.main_full_att_fp32_state.capacity) + self.set_att_state("draft0", prepared_full_state) + self.att_tail_type = f"indexer_kpool_c{tail_pool_size}_tail" + if with_tail: + tail = make_tail(device=device) + state = KPoolAttentionTailState( + replace(tail.config, pool_size=tail_pool_size), + request_capacity=tail.request_capacity, mtp_step=tail.mtp_step, device=tail.tail_buffer.device, + ) + self.set_att_state("main", state) + self.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) + self.bounded_att_big_page_allocator = StatePageAllocator(3) + self.bounded_att_small_page_allocator = StatePageAllocator(2) + self.finish_state_creation() + + @lru_cache(maxsize=None) + def get_cpu_transfer_layouts(self): + raise NotImplementedError("This test manager uses explicitly supplied CPU layouts") + + @lru_cache(maxsize=None) + def get_pd_full_att_transfer_layouts(self): + raise NotImplementedError("This test manager uses explicitly supplied PD layouts") + + @lru_cache(maxsize=None) + def get_pd_bounded_att_or_tail_transfer_layouts(self): + raise NotImplementedError("This test manager uses explicitly supplied PD layouts") diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py new file mode 100644 index 0000000000..38ddc83fbe --- /dev/null +++ b/unit_tests/common/att_resources/model_helpers.py @@ -0,0 +1,152 @@ +"""Select explicit model fields in tests parametrized over Qwen and GLM.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace + +import torch +from lightllm.common.att_resources.types import TensorField +from lightllm.common.att_resources.full_att import ( + FullAttentionState, FullAttentionConfig, StaticFP8FullAttentionState, + GQAFullAttentionState, GroupQuantizedFullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, +) + + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +def rebuild_full_state( + template: FullAttentionState, + *, + config: FullAttentionConfig | None = None, + token_capacity: int | None = None, + page_size: int | None = None, + device: str | torch.device | None = None, + scales: torch.Tensor | None = None, + q_scales: torch.Tensor | None = None, +) -> FullAttentionState: + """测试中按模板重建独立存储，显式传递运行时参数与静态标定值。""" + config = template.config if config is None else config + token_capacity = template.capacity if token_capacity is None else token_capacity + page_size = template.page_size if page_size is None else page_size + device = template.device if device is None else device + if isinstance(template, StaticFP8FullAttentionState): + prepared_full_state = type(template)( + config, page_size=page_size, device=device, + scales=template.scales if scales is None else scales, + q_scales=template.q_scales if q_scales is None else q_scales, + ) + prepared_full_state.initialize_storage(token_capacity) + return prepared_full_state + prepared_full_state = type(template)( + config, page_size=page_size, device=device, + ) + prepared_full_state.initialize_storage(token_capacity) + return prepared_full_state + + +def owned_states(attributes: Mapping[str, object], state_class: type | tuple[type, ...]): + """从显式成员选择指定类型的 State，按资源身份排序。""" + return tuple( + sorted( + (state for state in attributes.values() if isinstance(state, state_class)), + key=lambda state: state.resource_id, + ) + ) + + +def main_linear_state(resources: AttResourceManager): + if hasattr(resources, "main_kda_state"): + return resources.main_kda_state + return resources.main_gdn_state + + +def prepare_model_states(model: TpPartBaseModel, *, layer_ids: Sequence[int] | None = None): + config = model.config.get("text_config", model.config) + if layer_ids is None: + layer_ids = range(config.get("num_hidden_layers", config.get("n_layer"))) + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in layer_ids] + if not hasattr(model, "torch_memory_saver"): + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + model.prepare_att_resources() + + +def create_full_attention_roles( + template: FullAttentionState, main_layer_count: int, +) -> tuple[FullAttentionState, FullAttentionState | None]: + """仅用于测试：按配置层号创建各自拥有独立存储的 main 和单个 draft。""" + layers = template.config.layer_ids + boundary = sum(layer < main_layer_count for layer in layers) + assert boundary and all(layer < main_layer_count for layer in layers[:boundary]) + assert all(layer >= main_layer_count for layer in layers[boundary:]), "draft layers must follow main layers" + states = [] + for start, end in ((0, boundary), (boundary, len(layers))): + if start == end: + states.append(None) + continue + config = replace(template.config, layer_ids=layers[start:end]) + if isinstance(template, StaticFP8FullAttentionState): + state = rebuild_full_state( + template, config=config, scales=template.scales[start:end].clone(), + q_scales=None if template.q_scales is None else template.q_scales[start:end].clone(), + ) + else: + state = rebuild_full_state(template, config=config) + states.append(state) + return states[0], states[1] + + +def read_kv_storage(state: FullAttentionState) -> torch.Tensor | None: + """测试只读比较用：按 K、V 顺序拼接；修改数据须直接操作 State 的真实 buffer。""" + if isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)): + return state.storage + if state.k_storage is None: + return None + return torch.cat((state.k_storage, state.v_storage), dim=2) + + +def full_state_test_buffers(state: FullAttentionState) -> dict[str, torch.Tensor | None]: + """仅供多方案参数化测试枚举实际存储，检查容量、独立分配及字节一致性。""" + if isinstance(state, GroupQuantizedFullAttentionState): + return { + "k": state.k_storage, "v": state.v_storage, + "k_scale": state.k_scales, "v_scale": state.v_scales, + } + if isinstance(state, (GQAFullAttentionState, StaticFP8FullAttentionState)): + return {"k": state.k_storage, "v": state.v_storage} + assert isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) + return {"kv": state.storage} + + +def full_state_test_fields( + state: FullAttentionState, buffers: dict[str, torch.Tensor] | None = None, +) -> dict[str, TensorField]: + """测试中独立描述实际数据，用于校验不同存储方案的传输内容。""" + buffers = full_state_test_buffers(state) if buffers is None else buffers + axes = ("layer", "token", "head", "dim") + if isinstance(state, PackedMLAFullAttentionState): + storage = buffers["kv"] + raw = storage.view(torch.uint8) + if state.config.fp8_nope_dim: + fields = { + "kv_nope": TensorField(tensor=raw[..., :512], axes=axes), + "kv_scale": TensorField(tensor=raw[..., 512:528].view(torch.float32), axes=("layer", "token", "head", "group")), + "kv_rope": TensorField(tensor=raw[..., 528:656].view(torch.bfloat16), axes=axes), + } + else: + fields = {"kv": TensorField(tensor=storage[..., :state.config.kv_dim], axes=axes)} + start = state.indexer_offset + end = start + state.config.indexer_dim + fields["indexer_k"] = TensorField(tensor=raw[..., start:end], axes=axes) + fields["indexer_scale"] = TensorField(tensor=raw[..., end:end + 4].view(torch.float32), axes=("layer", "token", "head", "scale")) + return fields + return { + name: TensorField(tensor=tensor, axes=("layer", "token", "head", "group") if name.endswith("_scale") else axes) + for name, tensor in buffers.items() + } diff --git a/unit_tests/common/att_resources/test_allocators.py b/unit_tests/common/att_resources/test_allocators.py new file mode 100644 index 0000000000..da3c705509 --- /dev/null +++ b/unit_tests/common/att_resources/test_allocators.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import pytest +import torch + +from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator + + +def test_slot_allocation_domains_and_reserved_storage(): + big = StatePageAllocator(5, reserved_slots=2) + small = StatePageAllocator(5) + assert big.alloc(3) == [0, 1, 2] + assert small.alloc(1) == [0] + assert big.alloc(1) is None + with pytest.raises(AssertionError): + big.free([0, 3]) + assert big.free_count == 0 + big.free([0]) + assert big.free_count == 1 and small.free_count == 4 + with pytest.raises(AssertionError): + big.free([0]) + with pytest.raises(AssertionError): + big.free([1, 1]) + assert big.alloc(1) == [0] + + +def test_kv_returns_complete_pages_and_rejects_partial_free_atomically(): + allocator = FullAttKVAllocator(16, 4, pin_memory=False) + tokens = allocator.alloc(8) + snapshot = tokens.clone() + assert tokens.tolist() == list(range(8)) + assert allocator.hold_token_indices == (16, 17, 18, 19) + assert ((tokens.view(-1, 4)[:, 0] % 4) == 0).all() + assert (tokens.view(-1, 4).diff(dim=1) == 1).all() + with pytest.raises(AssertionError): + allocator.free(tokens[:2]) + assert allocator.can_use_mem_size == 8 + more = allocator.alloc(8) + assert set(tokens.tolist()).isdisjoint(more.tolist()) + assert torch.equal(tokens, snapshot) + with pytest.raises(AssertionError, match="available"): + allocator.alloc(4) + # Pages may be returned in any order; tokens within each page remain ordered. + returned = tokens.view(-1, 4).flip(0).reshape(-1) + allocator.free(returned) + assert allocator.can_use_mem_size == 8 + assert more.tolist() == list(range(8, 16)) + reused = allocator.alloc(8) + assert torch.equal(reused, returned) + allocator.free(more) + allocator.free(reused) + assert allocator.can_use_mem_size == 16 + assert set(allocator.alloc(16).tolist()) == set(range(16)) + + +def test_kv_validation_and_empty_operations(): + with pytest.raises(AssertionError, match="HOLD"): + FullAttKVAllocator((1 << 31) - 2, 3, pin_memory=False) + allocator = FullAttKVAllocator(8, 4, pin_memory=False) + for n in [-4, 3]: + with pytest.raises(AssertionError): + allocator.alloc(n) + assert allocator.alloc(0).numel() == 0 + allocator.free([]) + assert allocator.can_use_mem_size == 8 + tokens = allocator.alloc(4) + with pytest.raises(AssertionError, match="complete KV pages"): + allocator.free(tokens[:3]) + with pytest.raises(AssertionError, match="exceeds allocated"): + allocator.free(list(range(8))) + assert allocator.can_use_mem_size == 4 + allocator.free_all() + assert allocator.can_use_mem_size == 8 + assert set(allocator.alloc(8).tolist()) == set(range(8)) + + +@pytest.mark.parametrize("page_size", [1, 4, 16]) +def test_kv_repeated_allocation_and_return_buffer_wrap(page_size: int): + allocator = FullAttKVAllocator(64, page_size, pin_memory=False) + # Keep half the domain live while cycling the other half through the return buffer. + retained = allocator.alloc(32).clone() + for _ in range(16): + current = allocator.alloc(32) + snapshot = current.clone() + assert set(current.tolist()) == set(range(32, 64)) + allocator.free(current.tolist()) + assert torch.equal(current, snapshot) + assert allocator.can_use_mem_size == 32 + allocator.free(retained) + assert set(allocator.alloc(64).tolist()) == set(range(64)) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py new file mode 100644 index 0000000000..bf932d6a17 --- /dev/null +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -0,0 +1,305 @@ +"""Exercise shared-manager startup with independently prepared hybrid drafts.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from contextlib import nullcontext +from dataclasses import asdict +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.mode_backend import base_backend +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.utils.att_resource_config import build_cpu_cache_layouts +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def test_backend_rejects_more_than_36_drafts_before_loading(): + backend = ModeBackend.__new__(ModeBackend) + backend.args = StartArgs(mtp_mode="vanilla_with_att", mtp_step=37) + with pytest.raises(AssertionError, match="at most 36"): + backend.init_mtp_draft_model({}) + assert backend.draft_models == [] + + +@pytest.mark.parametrize("main_linear", [False, True]) +@pytest.mark.parametrize("draft_full", [False, True]) +@pytest.mark.parametrize("draft_count", [1, 2]) +@pytest.mark.parametrize("encoding", ["None", "fp8kv_sph", "fp8kv_spt"]) +@pytest.mark.parametrize("via_backend", [False, True]) +def test_backend_allocates_all_roles_before_warmup( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + main_linear: bool, + draft_full: bool, + draft_count: int, + encoding: str, + via_backend: bool, +): + args = StartArgs( + model_dir=str(tmp_path), + data_type="bfloat16", + llm_kv_type=encoding, + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + linear_att_cache_size=2, + cpu_cache_token_page_size=8, + mtp_mode="vanilla_with_att", + mtp_step=draft_count, + ) + for key in ("CURRENT_RANK_IN_DP", "CURRENT_RANK_IN_NODE", "CURRENT_DEVICE_ID"): + monkeypatch.setenv("LIGHTLLM_" + key, "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + config = dict( + model_type="qwen3_5_text", + architectures=["Qwen3NextForCausalLM"], + num_hidden_layers=4, + n_layer=4, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=64, + head_dim=16, + full_attention_interval=2, + linear_num_key_heads=2, + linear_num_value_heads=4, + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_conv_kernel_dim=3, + ) + draft_config = dict( + config, + num_hidden_layers=2 if draft_full else 1, + n_layer=2 if draft_full else 1, + num_key_value_heads=1, + head_dim=8, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + (draft_dir / "config.json").write_text(json.dumps(draft_config)) + args.mtp_draft_model_dir = [str(draft_dir)] * draft_count + main_full_layers = 2 if main_linear else 4 + if encoding != "None": + calibration_layers = main_full_layers + int(draft_full) * draft_count + scales = torch.arange(1, calibration_layers + 1, dtype=torch.float32) * 0.125 + calibration = { + "architectures": "Qwen3NextForCausalLM", + "num_head": 1, + "num_layers": calibration_layers, + "scales_shape": [calibration_layers, 2], + "scales": scales[:, None].expand(-1, 2).tolist(), + "q_calibration": { + "num_head": 1, + "scales_shape": [calibration_layers, 1], + "scales": (scales[:, None] * 2).tolist(), + }, + "quant_type": "per_head" if encoding == "fp8kv_sph" else "per_tensor", + } + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text(json.dumps(calibration)) + args.kv_quant_calibration_config_path = str(calibration_path) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(args))) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + main_cls = Qwen3NextTpPartModel if main_linear else LlamaTpPartModel + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: main_cls) + monkeypatch.setattr("lightllm.models.get_draft_model_class", lambda *_: Qwen3NextTpPartModel) + models, warmed, padded, captured, prefill_captured, ready = [], set(), set(), set(), set(), [] + checked = [] + mapping = SimpleNamespace(bind_hold_page=lambda indices: None) + + def mark_backends(model: TpPartBaseModel): + model.prefill_att_backend = model.decode_att_backend = object() + model.prefill_att_backend1 = model.decode_att_backend1 = None + model.backends_bound = True + + def warmup(model: TpPartBaseModel): + assert model.backends_bound + assert all(item.att_resource_manager is models[0].att_resource_manager for item in models) + assert model.att_resource_manager.full_att_allocator.size == 32 + assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) + assert all(item.max_total_token_num is None for item in models[models.index(model) + 1 :]) + assert captured == prefill_captured == {id(item) for item in models[: models.index(model)]} + assert not ready + warmed.add(id(model)) + + def capture(model: TpPartBaseModel): + assert warmed == padded == {id(item) for item in models[: models.index(model) + 1]} + assert not ready + captured.add(id(model)) + + def check_max_len(model: TpPartBaseModel): + assert model is models[0] + assert captured == prefill_captured == {id(model)} + assert not ready + checked.append(model) + + offset = 0 + for index in range(1 + draft_count): + cls = Qwen3NextTpPartModel if index else main_cls + model = cls.__new__(cls) + model.is_mtp_draft_model = bool(index) + model.args, model.config = args, draft_config.copy() if index else config.copy() + model.max_req_num, model.max_total_token_num, model.mem_fraction = 2, None, 0.9 + model.max_seq_length, model.batch_max_tokens = 32, 16 + model.tp_world_size_, model.tp_rank_ = 1, 0 + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(offset, offset + model.config["n_layer"])] + offset += len(model.layers_infer) + model.req_manager = mapping + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + if index: + model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() + model._init_att_backend = lambda model=model: mark_backends(model) + model._init_att_backend1 = lambda: None + model._init_hidden_collector = lambda: None + model._autotune_warmup = lambda model=model: warmup(model) + model._init_padded_req = lambda model=model: padded.add(id(model)) + model._init_cudagraph = lambda model=model: capture(model) + model._init_prefill_cuda_graph = lambda model=model: prefill_captured.add(id(model)) + model._check_max_len_infer = lambda model=model: check_max_len(model) + models.append(model) + + def memory_info(device: str | torch.device | int): + manager = models[0].att_resource_manager + states = manager.get_full_attention_states() + assert all(read_kv_storage(state) is None for state in states) + assert manager.draft0_gdn_state.conv_state.is_cuda + assert manager.draft0_gdn_state.big_pages_pool is None + assert manager.full_att_allocator is None + assert not warmed + return manager.get_full_att_token_bytes() * (32 + args.page_size), 0 + + monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) + monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", ready.append) + assert all(not hasattr(item, "att_resource_manager") for item in models) + + class RuntimeReady(Exception): + """Stop after model readiness, before unrelated queue and thread setup.""" + + if via_backend: + backend = ModeBackend.__new__(ModeBackend) + backend.init_rank_infos = lambda: None + backend.init_mtp_draft_model = lambda _: backend.draft_models.extend(models[1:]) + monkeypatch.setattr(base_backend, "init_distributed_env", lambda _: None) + monkeypatch.setattr(base_backend.dist_group_manager, "create_groups", lambda **_: None) + monkeypatch.setattr(base_backend, "TokenLoad", lambda *_: None) + monkeypatch.setattr(base_backend, "get_model", lambda *_: (models[0], False)) + + def mark_ready(value: bool): + ready.append(value) + raise RuntimeReady + + monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", mark_ready) + + def start_runtime(): + if via_backend: + with pytest.raises(RuntimeReady): + backend.init_model( + dict(args=args, world_size=1, load_way="HF", weight_dir=str(tmp_path), max_total_token_num=None) + ) + else: + for index, item in enumerate(models): + item.prepare_att_resources() + assert all(other.att_resource_manager is models[0].att_resource_manager for other in models[: index + 1]) + assert all(not hasattr(other, "att_resource_manager") for other in models[index + 1 :]) + main = models[0] + main.att_resource_manager.initialize_resources( + args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction + ) + assert all(item.max_total_token_num is None for item in models) + for index, item in enumerate(models): + item.init_runtime() + assert all(not getattr(other, "backends_bound", False) for other in models[index + 1 :]) + assert warmed == padded == captured == prefill_captured == { + id(other) for other in models[: index + 1] + } + + start_runtime() + assert captured == {id(model) for model in models} + assert checked == [models[0]] + assert ready == ([True] if via_backend else []) + manager = models[0].att_resource_manager + assert len(manager.get_bounded_attention_states()) == draft_count + int(main_linear) + assert all(not hasattr(model, "mtp_state_cursors") for model in models) + assert all(not hasattr(model, "full_att_state") for model in models) + draft_states = [manager.get_bounded_attention_states(f"draft{i}")[0] for i in range(draft_count)] + assert len({state.conv_state.data_ptr() for state in draft_states}) == draft_count + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), role="main", + ) + assert all(not state.mtp_state_index.any() for state in draft_states) + if main_linear: + assert manager.main_gdn_state.mtp_state_index[0] == draft_count + for i, state in enumerate(draft_states): + assert len(state.config.layer_ids) == 1 + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), role=f"draft{i}", + ) + assert state.mtp_state_index[0] == i + 1 + assert all(not other.mtp_state_index.any() for other in draft_states[i + 1 :]) + if draft_full: + main_state = manager.get_full_attention_states("main")[0] + draft_state = manager.get_full_attention_states(f"draft{i}")[0] + assert main_state.config.head_dim == 16 + assert draft_state.config.head_dim == 8 + assert {buffer.data_ptr() for buffer in full_state_test_buffers(main_state).values()}.isdisjoint( + buffer.data_ptr() for buffer in full_state_test_buffers(draft_state).values() + ) + assert len(draft_state.config.layer_ids) == 1 + else: + assert not manager.get_full_attention_states(f"draft{i}") + + if encoding != "None": + ranges = [("main", 0, main_full_layers)] + [ + (f"draft{i}", main_full_layers + i, main_full_layers + i + 1) for i in range(draft_count) + ] + for role, start, end in ranges: + for state in manager.get_full_attention_states(role): + assert torch.equal(state.scales[:, 0].cpu(), scales[start:end]) + if encoding == "fp8kv_sph": + assert torch.equal(state.q_scales[:, 0].cpu(), scales[start:end] * 2) + + layouts = build_cpu_cache_layouts() + assert (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 == manager.get_cpu_page_bytes() + assert layouts.keys() == manager.get_cpu_transfer_layouts().keys() + for resource_id, layout in layouts.items(): + runtime_layout = manager.get_cpu_transfer_layouts()[resource_id] + assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in layout.regions] == [ + (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in runtime_layout.regions + ] + slot = manager.bounded_att_big_page_allocator.alloc(1)[0] + for i, state in enumerate(manager.get_bounded_attention_states()): + state.conv_state.fill_(i + 3) + state.ssm_state.fill_(i + 7) + manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + manager.reset_bounded_att_state([0]) + assert all(not state.conv_state[:, 0].any() for state in manager.get_bounded_attention_states()) + manager.restore_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + for i, state in enumerate(manager.get_bounded_attention_states()): + assert (state.conv_state[:, 0, :, : state.config.conv_width] == i + 3).all() + assert (state.ssm_state[:, 0] == i + 7).all() + with pytest.raises(AssertionError, match="resources already initialized"): + start_runtime() + assert all(item.att_resource_manager is manager for item in models) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py new file mode 100644 index 0000000000..a710d3de6b --- /dev/null +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.bounded_att import ( + GDNPagesPool, + GDNState, + KDAPagesPool, + KDAState, + ConvSSMStateConfig, +) +from lightllm.common.att_resources.types import ( + AxisSlice, StateTransferLayout, TransferRegion, +) + +if TYPE_CHECKING: + from lightllm.common.att_resources.types import PageKind + + +def make_state( + cls: type[KDAState] | type[GDNState] = KDAState, + *, + device: str | torch.device | int = "cpu", + heads: int = 1, + mtp_step: int = 2, +): + config = ConvSSMStateConfig( + layer_ids=(1, 3), global_k_heads=heads, global_v_heads=heads, key_dim=2, value_dim=2, conv_width=2, + conv_dtype=torch.float32, ssm_dtype=torch.float32, + ) + state = cls(config, request_capacity=3, mtp_step=mtp_step, device=device) + manager = AttResourceManager(request_capacity=3) + manager.set_att_state("main", state) + state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + return state + + +def page_layout(state: KDAState | GDNState): + cfg = state.config + conv_shape = (cfg.conv_channels, len(cfg.layer_ids), cfg.conv_width) + conv_bytes = cfg.conv_channels * len(cfg.layer_ids) * cfg.conv_width * 4 + ssm_shape = (len(cfg.layer_ids), cfg.num_v_heads, cfg.key_dim, cfg.value_dim) + ssm_bytes = len(cfg.layer_ids) * cfg.num_v_heads * cfg.key_dim * cfg.value_dim * 4 + end = 16 + conv_bytes + ssm_bytes + return StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=end, + regions=( + TransferRegion( + field="conv", + byte_start=16, + byte_end=16 + conv_bytes, + dtype=torch.float32, + axes=("channel", "layer", "window"), + shape=conv_shape, + axis_slices=( + AxisSlice(axis="channel", src_start=0, dst_start=0, length=cfg.conv_channels), + AxisSlice(axis="window", src_start=0, dst_start=0, length=cfg.conv_width), + ), + ), + TransferRegion( + field="ssm", + byte_start=16 + conv_bytes, + byte_end=end, + dtype=torch.float32, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_slices=( + AxisSlice(axis="head", src_start=0, dst_start=0, length=cfg.num_v_heads), + AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=cfg.key_dim), + AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=cfg.value_dim), + ), + ), + ), + ) + + +@pytest.mark.parametrize("cls,pool_cls", [(KDAState, KDAPagesPool), (GDNState, GDNPagesPool)]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("page_kind", ["big", "small"]) +def test_prefill_state_is_saved_and_restored( + cls: type[KDAState] | type[GDNState], + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], + device: str | torch.device | int, + page_kind: PageKind, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + config = ConvSSMStateConfig( + layer_ids=(1, 3), global_k_heads=1, global_v_heads=1, key_dim=2, value_dim=2, conv_width=2, + conv_dtype=torch.float32, ssm_dtype=torch.float32, + ) + state = cls(config, request_capacity=3, mtp_step=2, device=device) + assert state.resource_id is None and not hasattr(config, "resource_id") + assert state.big_pages_pool is None and state.small_pages_pool is None + conv, ssm = state.get_layer_view(3) + assert conv.data_ptr() == state.conv_state[1].data_ptr() + assert ssm.data_ptr() == state.ssm_state[1].data_ptr() + for j in range(4): + state.conv_state[:, 0, :, j].fill_(j + 10) + for j in range(3): + state.ssm_state[:, j].fill_(j + 20) + state.mtp_state_index[2].fill_(1) + # 延后创建页池不能重建运行态或重置已接受的 MTP 游标。 + state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device == "cuda") + assert conv.data_ptr() == state.conv_state[1].data_ptr() + assert ssm.data_ptr() == state.ssm_state[1].data_ptr() + assert type(state.big_pages_pool) is pool_cls + assert type(state.small_pages_pool) is pool_cls + assert state.big_pages_pool is not state.small_pages_pool + assert state.mtp_state_index.tolist() == [0, 0, 1] + state.save_bounded_att_pages([0], page_kind=page_kind, slot_ids=[1]) + if device == "cuda": + torch.cuda.synchronize() + pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool + assert torch.equal(pool.conv_state[1, ..., 0], torch.full_like(pool.conv_state[1, ..., 0], 10)) + assert torch.equal(pool.conv_state[1, ..., 1], torch.full_like(pool.conv_state[1, ..., 1], 11)) + assert (pool.ssm_state[1] == 20).all() + state.mtp_state_index[1] = 2 + state.restore_bounded_att_pages([1], page_kind=page_kind, slot_ids=[1]) + assert state.mtp_state_index[1].item() == 0 + assert torch.equal(state.conv_state[:, 1, :, :2].cpu(), pool.conv_state[1]) + assert torch.equal(state.ssm_state[:, 3].cpu(), pool.ssm_state[1]) + + +def test_gdn_dimensions_and_page_restore(): + config = replace( + make_state().config, global_k_heads=1, global_v_heads=2, key_dim=2, value_dim=3 + ) + state = GDNState(config, request_capacity=3, mtp_step=2, device="cpu") + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2, pin_memory=False) + conv = torch.arange(40, dtype=config.conv_dtype).reshape(2, 10, 2) + ssm = torch.arange(24, dtype=config.ssm_dtype).reshape(2, 2, 2, 3) + state.conv_state[:, 0, :, :2].copy_(conv) + state.ssm_state[:, 0].copy_(ssm) + for page_kind in ("big", "small"): + state.save_bounded_att_pages([0], page_kind=page_kind, slot_ids=[0]) + state.reset_bounded_att_state([0]) + for req, page_kind in enumerate(("big", "small"), start=1): + state.restore_bounded_att_pages([req], page_kind=page_kind, slot_ids=[0]) + assert torch.equal(state.conv_state[:, req, :, :2], conv) + assert torch.equal(state.ssm_state[:, req * 3], ssm) + with pytest.raises(AssertionError, match="equal Q/K/V"): + KDAPagesPool(2, config, pin_memory=False) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): + state = make_state(cls) + state.conv_state[:, 0].fill_(7) + state.ssm_state[:, 0:3].fill_(11) + state.save_bounded_att_pages([0], page_kind="small", slot_ids=[0]) + # 推进候选区和游标后，重置应清零全部运行态，同时保留已保存的页快照。 + state.conv_state[:, 0, :, 2:4].fill_(19) + state.ssm_state[:, 2].fill_(23) + state.mtp_state_index[0].fill_(2) + state.reset_bounded_att_state([0]) + assert not state.conv_state[:, 0].any() and not state.ssm_state[:, :3].any() + assert state.mtp_state_index[0] == 0 + assert (state.small_pages_pool.conv_state[0] == 7).all() + assert (state.small_pages_pool.ssm_state[0] == 11).all() + with pytest.raises(AssertionError, match="CUDA tensors"): + state.commit_bounded_att_state(torch.tensor([0]), torch.tensor([1], dtype=torch.int32)) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): + source = make_state(cls) + source.conv_state[:, 0].fill_(13) + source.ssm_state[:, 0].fill_(29) + source.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) + layout = page_layout(source) + buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8) + source.write_page_state_to_cpu("big", [0], buffer, [2], layout=layout) + assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() + target = make_state(cls) + target.read_page_state_from_cpu(buffer, [2], big_page_slot_ids=[1], layout=layout) + target.restore_bounded_att_pages([2], page_kind="big", slot_ids=[1]) + assert (target.conv_state[:, 2, :, :2] == 13).all() + assert (target.ssm_state[:, 6] == 29).all() + # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 + target.read_page_state_from_cpu(buffer, [2], big_page_slot_ids=[0], layout=layout) + target.write_page_state_to_cpu("big", [0], buffer, [0], layout=layout) + assert torch.equal(buffer[0], buffer[2]) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( + cls: type[KDAState] | type[GDNState], + device: str | torch.device | int, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source = make_state(cls, device=device) + source.conv_state.fill_(-1) + source.ssm_state.fill_(-1) + source.conv_state[:, 0, :, :2].fill_(23) + source.ssm_state[:, 0].fill_(47) + layout = page_layout(source) + # PD 只传 Conv/SSM，仍保留非零区域起点以覆盖组合页偏移。 + allocation = torch.full((3, layout.byte_end), 254, dtype=torch.uint8, device=device) + page = allocation[1] + source.write_to_pd_buffer(0, allocation, pd_page_index=1, layout=layout) + target = make_state(cls, device=device, mtp_step=1) + target.read_from_pd_buffer(allocation, 2, pd_page_index=1, layout=layout) + assert (target.conv_state[:, 2, :, :2] == 23).all() + assert (target.ssm_state[:, 4] == 47).all() + assert (allocation[0] == 254).all() and (allocation[2] == 254).all() + assert (page[:16] == 254).all() + assert target.mtp_state_index[2] == 0 + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("src_mtp_step,dst_mtp_step", [(0, 2), (2, 1)]) +def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp_step, dst_mtp_step): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source = make_state(cls, device=device, mtp_step=src_mtp_step) + target = make_state(cls, device=device, mtp_step=dst_mtp_step) + source.conv_state.fill_(-7) + source.ssm_state.fill_(-9) + target.conv_state.fill_(-1) + target.ssm_state.fill_(-2) + target.mtp_state_index.fill_(1) + target.big_pages_pool.conv_state.fill_(91) + target.big_pages_pool.ssm_state.fill_(92) + for req in (0, 2): + source.conv_state[:, req, :, :source.config.conv_width].fill_(10 + req) + source.ssm_state[:, req * (src_mtp_step + 1)].fill_(20 + req) + src_conv = source.conv_state.clone() + src_ssm = source.ssm_state.clone() + + target.copy_run_time_state_from_other_dp( + src_req_indices=[2, 0], dst_req_indices=[0, 1], remote_dp_bounded_state=source, + ) + + for src_req, dst_req in ((2, 0), (0, 1)): + assert (target.conv_state[:, dst_req, :, :2] == 10 + src_req).all() + assert (target.conv_state[:, dst_req, :, 2:] == -1).all() + dst_slot = dst_req * (dst_mtp_step + 1) + assert (target.ssm_state[:, dst_slot] == 20 + src_req).all() + assert (target.ssm_state[:, dst_slot + 1:dst_slot + dst_mtp_step + 1] == -2).all() + assert target.mtp_state_index[dst_req] == 0 + assert (target.conv_state[:, 2] == -1).all() + assert (target.ssm_state[:, 2 * (dst_mtp_step + 1):] == -2).all() + assert target.mtp_state_index[2] == 1 + assert (target.big_pages_pool.conv_state == 91).all() + assert (target.big_pages_pool.ssm_state == 92).all() + assert torch.equal(source.conv_state, src_conv) + assert torch.equal(source.ssm_state, src_ssm) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): + source = make_state(cls) + target = make_state(cls) + for page in range(3): + source.big_pages_pool.conv_state[page].fill_(10 + page) + source.big_pages_pool.ssm_state[page].fill_(20 + page) + target.big_pages_pool.conv_state.fill_(-1) + target.big_pages_pool.ssm_state.fill_(-2) + target.small_pages_pool.conv_state.fill_(31) + target.small_pages_pool.ssm_state.fill_(32) + target.conv_state.fill_(41) + target.ssm_state.fill_(42) + target.mtp_state_index.fill_(1) + src_conv = source.big_pages_pool.conv_state.clone() + src_ssm = source.big_pages_pool.ssm_state.clone() + + target.copy_big_page_state_from_other_dp( + src_page_ids=[2, 0], dst_page_ids=[0, 1], remote_dp_bounded_state=source, + ) + + for src_page, dst_page in ((2, 0), (0, 1)): + assert torch.equal(target.big_pages_pool.conv_state[dst_page], src_conv[src_page]) + assert torch.equal(target.big_pages_pool.ssm_state[dst_page], src_ssm[src_page]) + assert (target.big_pages_pool.conv_state[2] == -1).all() + assert (target.big_pages_pool.ssm_state[2] == -2).all() + assert (target.small_pages_pool.conv_state == 31).all() + assert (target.small_pages_pool.ssm_state == 32).all() + assert (target.conv_state == 41).all() and (target.ssm_state == 42).all() + assert (target.mtp_state_index == 1).all() + assert torch.equal(source.big_pages_pool.conv_state, src_conv) + assert torch.equal(source.big_pages_pool.ssm_state, src_ssm) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py new file mode 100644 index 0000000000..e4dc14b509 --- /dev/null +++ b/unit_tests/common/att_resources/test_capacity.py @@ -0,0 +1,692 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +from dataclasses import replace +from functools import partial +from math import prod + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNState, KDAState +from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.common.att_resources.full_att import ( + FullAttentionState, + GQAFullAttentionState, + FullAttentionConfig, + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + MLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLAFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, +) + +if TYPE_CHECKING: + from easydict import EasyDict + + +def definition(kind: str): + config = FullAttentionConfig((0, 2), 2, 16, torch.bfloat16) + if kind in ("int4", "int8"): + return GroupQuantizedFullAttentionState( + GroupQuantizedAttentionConfig(**vars(config), bits=4 if kind == "int4" else 8), + page_size=4, device="meta", + ) + if kind.startswith("static"): + heads = 2 if kind == "static_head" else 1 + return StaticFP8FullAttentionState( + StaticFP8AttentionConfig( + **vars(config), scale_mode="per_head" if heads == 2 else "per_tensor", calibration_id="test" + ), + page_size=4, device="meta", scales=torch.ones((2, heads * 2), dtype=torch.float32), q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, + ) + if kind.startswith("packed"): + fp8 = kind == "packed_fp8" + kv_dim = 512 if kind == "packed_glm" else 576 + config = replace( + config, global_head_num=1, head_dim=800 if fp8 else kv_dim + 72, dtype=torch.uint8 if fp8 else config.dtype + ) + return PackedMLAFullAttentionState( + PackedMLAAttentionConfig(**vars(config), kv_dim=kv_dim, fp8_nope_dim=512 if fp8 else 0), page_size=4, device="meta" + ) + if kind == "mla": + return MLAFullAttentionState(replace(config, global_head_num=1, head_dim=576), page_size=4, device="meta") + return GQAFullAttentionState(config, page_size=4, device="meta") + + +def allocated_bytes(state: FullAttentionState, *, token_only: bool = False): + buffers = list(full_state_test_buffers(state).values()) + if not token_only and isinstance(state, StaticFP8FullAttentionState): + buffers.append(state.scales) + if state.q_scales is not None: + buffers.append(state.q_scales) + stores = {} + for buffer in buffers: + storage = buffer.untyped_storage() + stores[(buffer.device, storage.data_ptr())] = storage.nbytes() + return sum(stores.values()) + + +@pytest.mark.parametrize( + "kind", ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] +) +def test_budget_matches_token_stores_including_hold_and_padding( + monkeypatch: pytest.MonkeyPatch, + kind: str, + pd_args: EasyDict, +): + template = definition(kind) + state = rebuild_full_state(template, config=template.config, token_capacity=20, device="cpu") + actual = allocated_bytes(state, token_only=True) + planning = AttResourceManager(request_capacity=1) + planning.set_att_state("main", rebuild_full_state(template, config=template.config)) + planning.finish_state_creation() + token_bytes = planning.get_full_att_token_bytes() + pd_page_bytes = planning.get_pd_page_bytes() + assert not hasattr(planning, "pd_page_bytes") + assert not hasattr(planning, "pd_transfer_layouts") + assert token_bytes * (20 + 4) == actual + if kind.startswith("static"): + assert allocated_bytes(state) > actual + assert state.get_token_bytes() == token_bytes + minimum = token_bytes * 4 + assert minimum > 0 + + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", state) + manager.full_att_allocator = FullAttKVAllocator(20, 4, pin_memory=False) + manager.finish_state_creation() + (state,) = manager.get_full_attention_states("main") + assert manager.get_pd_page_bytes() == pd_page_bytes + args = StartArgs(page_size=4, run_mode="prefill", pd_kv_page_size=8, pd_kv_page_num=3) + reserve = manager.get_pd_page_bytes() * args.pd_kv_page_num + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.get_current_device_id", lambda: 0) + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + for budget, expected in ((actual, 20), (actual - 1, 16), (minimum, 0), (minimum - 1, None)): + resources = AttResourceManager(request_capacity=1) + resources.set_att_state("main", rebuild_full_state(template, config=template.config)) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) + if expected is None: + with pytest.raises(MemoryError, match="HOLD"): + resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + elif expected == 0: + with pytest.raises(AssertionError, match="at least one complete page"): + resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + else: + resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + assert resources.full_att_allocator.size == expected + assert resources.get_pd_page_bytes() == manager.get_pd_page_bytes() + + +def test_shared_allocator_accounts_for_distinct_state_stores(): + main, draft = definition("int4"), definition("static_head") + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", main) + manager.set_att_state("draft0", draft) + manager.finish_state_creation() + expected = sum( + allocated_bytes( + rebuild_full_state(state, config=state.config, token_capacity=12, page_size=4, device="cpu"), + token_only=True, + ) + for state in (main, draft) + ) + assert manager.get_full_att_token_bytes() * (12 + 4) == expected + + +@pytest.mark.parametrize("state_cls", [KDAState, GDNState]) +def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( + monkeypatch: pytest.MonkeyPatch, + state_cls: type[KDAState] | type[GDNState], + pd_args: EasyDict, +): + config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=2, global_v_heads=2 if state_cls is KDAState else 4, key_dim=3, + value_dim=3 if state_cls is KDAState else 5, conv_width=1, conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + tail = KPoolAttentionTailState( + KPoolStateConfig(layer_ids=(1,), pool_size=4, head_dim=3, dtype=torch.bfloat16), + request_capacity=3, mtp_step=3, device="cpu", + ) + reference = state_cls(config, request_capacity=3, mtp_step=3, device="cpu") + conv = reference.conv_state[:, 0, :, : config.conv_width] + ssm = reference.ssm_state[:, 0] + bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() + tail_bytes = len(tail.config.layer_ids) * tail.config.pool_size * 2 * tail.config.head_dim * tail.config.dtype.itemsize + manager = AttResourceManager(request_capacity=3) + manager.set_att_state("main", reference) + manager.set_att_state("main", tail) + manager.finish_state_creation() + # PD 页只包含运行态数据，使用第二页验证 8 字节页跨度下的读写。 + assert manager.get_pd_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 + reference.conv_state.fill_(2) + reference.ssm_state.fill_(3) + tail.tail_buffer.fill_(4) + pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8) + layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() + assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "tail"} + manager.write_to_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + reference.conv_state.zero_() + reference.ssm_state.zero_() + tail.tail_buffer.zero_() + manager.read_from_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + assert (reference.conv_state[:, 0, :, : config.conv_width] == 2).all() + assert (reference.ssm_state[:, 0] == 3).all() + expected_tail = torch.zeros_like(tail.tail_buffer[:, 0]) + for position in range(4, 8): + expected_tail[:, position % tail.tail_buffer.shape[2]].fill_(4) + assert torch.equal(tail.tail_buffer[:, 0], expected_tail) + assert not pages[0].any() + for state in (reference, tail): + monkeypatch.setattr( + state, "get_pd_transfer_regions", lambda *args, **kwargs: pytest.fail("cached page size must not rebuild regions") + ) + assert manager.get_pd_page_bytes() == pages.shape[1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = 3 + manager = AttResourceManager(request_capacity=1) + for state_cls, role, dim, dtype in ( + (MLAFullAttentionState, "draft0", 1, torch.bfloat16), + (GQAFullAttentionState, "main", 3, torch.bfloat16), + (GQAFullAttentionState, "main", 1, torch.float32), + ): + config = FullAttentionConfig((0,), 1, dim, dtype) + manager.set_att_state(role, state_cls(config, page_size=1, device="meta")) + manager.finish_state_creation() + # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 + assert manager.get_pd_page_bytes() == 72 + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.initialize_storage(4, 1, big_page_tokens=4, device="cuda") + assert manager.get_pd_page_bytes() == 72 + states = manager.get_full_attention_states() + for i, state in enumerate(states, 1): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.fill_(i) + pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + page = pages[1] + layouts = manager.get_pd_full_att_transfer_layouts() + assert [min(region.byte_start for region in layouts[state.resource_id].regions) for state in states] == [0, 8, 48] + manager.write_to_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + assert (page[:6].view(torch.bfloat16) == 1).all() + for state in states: + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + manager.read_from_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + for i, state in enumerate(states, 1): + assert (read_kv_storage(state)[:, :3] == i).all() + assert not pages[0].any() + + +@pytest.mark.parametrize("layout_first", [False, True]) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_capacity_uses_quantized_field_padding_and_preserves_page_boundaries( + monkeypatch: pytest.MonkeyPatch, pd_args: EasyDict, layout_first: bool, transfer: str +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + class AlignedFieldsState(GroupQuantizedFullAttentionState): + def _get_transfer_regions(self, page_tokens: int): + regions, offset = [], 0 + for region in super()._get_transfer_regions(page_tokens): + offset = (offset + 7) // 8 * 8 + regions.append(replace(region, byte_start=offset, byte_end=offset + region.byte_end - region.byte_start)) + offset += prod(region.shape) * region.dtype.itemsize + return tuple(regions) + + pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 + config = GroupQuantizedAttentionConfig( + layer_ids=(0,), global_head_num=1, head_dim=128, + dtype=torch.float32, bits=8, group_size=128, + ) + manager = AttResourceManager(request_capacity=1) + prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") + prepared_full_state.initialize_storage(8) + manager.set_att_state("main", prepared_full_state) + prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") + prepared_full_state.initialize_storage(8) + manager.set_att_state("draft0", prepared_full_state) + manager.full_att_allocator = FullAttKVAllocator(8, 1, pin_memory=False) + manager.finish_state_creation() + # 每个 State 的裸数据为 792 字节；scale 间 padding 使实际末尾为 796，页段需 800 字节。 + for state in manager.get_full_attention_states(): + assert state.get_token_bytes() * 3 == 792 + monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("capacity must use field layouts")) + get_layouts = manager.get_cpu_transfer_layouts if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts + get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes + if layout_first: + with monkeypatch.context() as patch: + patch.setattr(manager, f"get_{transfer}_page_bytes", lambda: pytest.fail("layout must not query page bytes")) + layouts = get_layouts() + assert get_page_bytes() == 1600 + layouts = get_layouts() + assert get_layouts() is layouts + pages = torch.full( + (3, 1600), 91, dtype=torch.uint8, device="cpu" if transfer == "cpu" else "cuda", pin_memory=transfer == "cpu" + ) + for index, state in enumerate(manager.get_full_attention_states()): + layout = layouts[state.resource_id] + assert (layout.byte_start, layout.byte_end) == (index * 800, (index + 1) * 800) + regions = layout.regions + assert [region.byte_start for region in regions] == [index * 800 + n for n in (0, 384, 768, 784)] + for field_index, field in enumerate(full_state_test_fields(state, full_state_test_buffers(state)).values()): + field.tensor.zero_() + field.tensor[:, :3].fill_(index * 4 + field_index + 1) + if transfer == "cpu": + manager.write_req_to_cpu(torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), pages, [1], layouts=layouts) + manager.read_req_from_cpu(pages, [1], torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), layouts=layouts) + else: + manager.write_to_pd_buffer(pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + manager.read_from_pd_buffer(pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda")) + for state in manager.get_full_attention_states(): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + torch.testing.assert_close(field.tensor[:, 4:7], field.tensor[:, :3]) + assert (pages[[0, 2]] == 91).all() + for offset in (0, 800): + assert (pages[1, offset + 780:offset + 784] == 91).all() + assert (pages[1, offset + 796:offset + 800] == 91).all() + + +@pytest.mark.parametrize("full_page_tokens", [None, 8, 10, 16]) +def test_pd_runtime_capacity_includes_field_padding( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + full_page_tokens: int, +): + class PaddedGDNState(GDNState): + def get_pd_transfer_regions(self): + return tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) if region.field == "ssm" else region + for region in super().get_pd_transfer_regions() + ) + + manager = AttResourceManager(request_capacity=1) + if full_page_tokens is not None: + pd_args.pd_kv_page_size = full_page_tokens + config = FullAttentionConfig((0,), 1, 1, torch.float32) + manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + state = PaddedGDNState( + ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + ), + request_capacity=1, device="cpu", + ) + manager.set_att_state("main", state) + manager.finish_state_creation() + # Conv 为 6 字节，SSM 从偏移 72 开始，占 4 字节；完整运行态需要 80 字节。 + # 即使 KV 页不足，布局仍可独立查询；最终块大小查询选择两类占用的最大值。 + with monkeypatch.context() as patch: + patch.setattr(manager, "get_pd_full_att_transfer_layouts", lambda: pytest.fail("runtime layout must be independent")) + patch.setattr(manager, "get_pd_page_bytes", lambda: pytest.fail("layout must not query page bytes")) + layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() + layout = layouts[state.resource_id] + assert layout.byte_end == 80 and layout.tokens_per_page is None + assert layout.regions[-1].byte_start == 72 + warnings = [] + monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.logger.warning", warnings.append) + expected = max(80, (full_page_tokens or 0) * 8) + assert manager.get_pd_page_bytes() == expected + assert manager.get_pd_bounded_att_or_tail_transfer_layouts() is layouts + # 实际传输页可大于运行态布局，选中第二页后仍从该页起点搬运。 + pages = torch.full((3, expected), 91, dtype=torch.uint8) + state.conv_state.fill_(3) + state.ssm_state.fill_(5) + manager.write_to_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + state.conv_state.zero_() + state.ssm_state.zero_() + manager.read_from_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + assert (state.conv_state == 3).all() and (state.ssm_state == 5).all() + assert (pages[[0, 2]] == 91).all() and (pages[1, 80:] == 91).all() + assert len(warnings) == int(full_page_tokens == 8) + if warnings: + assert "64 bytes" in warnings[0] and "using 80 bytes per transfer page" in warnings[0] + # 相同 manager 的容量查询命中缓存，不重复打印 warning。 + assert manager.get_pd_page_bytes() == expected + assert len(warnings) == int(full_page_tokens == 8) + + +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_page_bytes_refresh_with_storage_dependent_layout( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + transfer: str, +): + class PaddedState(GQAFullAttentionState): + def _get_transfer_regions(self, page_tokens: int): + return tuple( + replace( + region, + byte_start=region.byte_start + self.page_size * 8, + byte_end=region.byte_end + self.page_size * 8, + ) + for region in super()._get_transfer_regions(page_tokens) + ) + + pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 + manager = AttResourceManager(request_capacity=1) + config = FullAttentionConfig((0,), 1, 3, torch.bfloat16) + manager.set_att_state("main", PaddedState(config, page_size=1, device="meta")) + manager.finish_state_creation() + get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes + assert get_page_bytes() == 48 + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.initialize_storage(8, 4, big_page_tokens=8, device="cpu") + assert get_page_bytes() == (80 if transfer == "cpu" else 72) + layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() + assert layouts["main_full_att_bf16_state"].byte_end == 72 + + +@pytest.mark.parametrize("with_tail", [False, True]) +@pytest.mark.parametrize("extra_tokens", [-2, 0, 2]) +def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( + with_tail: bool, + extra_tokens: int, + pd_args: EasyDict, +): + manager = AttResourceManager(request_capacity=1) + config = FullAttentionConfig((0,), 1, 1, torch.float32) + full_state = GQAFullAttentionState(config, page_size=1, device="meta") + manager.set_att_state("main", full_state) + for cls, role in ((KDAState, "main"), (GDNState, "draft0")): + config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + conv_dtype=torch.float32, ssm_dtype=torch.float32, + ) + state = cls(config, request_capacity=1, device="meta") + manager.set_att_state(role, state) + if with_tail: + state = KPoolAttentionTailState( + KPoolStateConfig(layer_ids=(0,), pool_size=4, head_dim=1, dtype=torch.float32), + request_capacity=1, mtp_step=0, device="meta", + ) + manager.set_att_state("main", state) + # 两份运行态各占 16 字节数据，tail 另占 32 字节，无页头及长度、版本字段。 + runtime_page_bytes = 32 + 32 * with_tail + page_tokens = runtime_page_bytes // 8 + extra_tokens + pd_args.pd_kv_page_size = page_tokens + manager.finish_state_creation() + assert manager.get_pd_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) + assert manager.full_att_allocator is None + assert read_kv_storage(manager.main_full_att_fp32_state) is None + assert not hasattr(manager, "pd_transfer_layouts") + + +@pytest.mark.parametrize( + "kind", ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] +) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +@pytest.mark.parametrize("page_tokens", [1, 3]) +def test_cpu_page_bytes_derive_from_cached_transfer_layouts( + monkeypatch: pytest.MonkeyPatch, cpu_args: EasyDict, kind: str, tp_world_size: int, page_tokens: int +): + cpu_args.cpu_cache_token_page_size = page_tokens + template = definition(kind) + config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) + scales = torch.ones((len(config.layer_ids), 2 * config.head_num)) if kind == "static_head" else None + q_scales = torch.ones((len(config.layer_ids), config.head_num)) if kind == "static_head" else None + manager = AttResourceManager(request_capacity=3) + for role in ("main", "draft0"): + manager.set_att_state(role, rebuild_full_state(template, config=config, scales=scales, q_scales=q_scales)) + for state_cls, v_heads, value_dim in ((KDAState, 1, 1), (GDNState, 2, 3)): + # 奇数长度的 BF16 Conv 与 FP32 SSM 之间需要补齐，MTP 候选区不应计入页大小。 + bounded_config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=v_heads, key_dim=1, value_dim=value_dim, conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, tp_rank=tp_world_size - 1, + tp_world_size=tp_world_size, + ) + bounded = state_cls(bounded_config, request_capacity=3, mtp_step=2, device="meta") + manager.set_att_state("main", bounded) + manager.finish_state_creation() + # CPU 容量和布局都只依赖 State 的 CPU 字段描述，不读取 GPU 存储开销或 PD 布局。 + for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): + if isinstance(state, FullAttentionState): + monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("CPU capacity must use field layouts")) + monkeypatch.setattr(state, "get_pd_transfer_regions", lambda: pytest.fail("CPU layout must be independent")) + layouts = manager.get_cpu_transfer_layouts() + assert min(region.byte_start for layout in layouts.values() for region in layout.regions) == 0 + assert not {"state_token_len", "format_version"} & { + region.field for layout in layouts.values() for region in layout.regions + } + expected_page_bytes = (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 + assert all(layout.tokens_per_page == page_tokens for layout in layouts.values()) + for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): + monkeypatch.setattr(state, "get_cpu_transfer_regions", lambda: pytest.fail("cached layout must be reused")) + assert manager.get_cpu_page_bytes() == expected_page_bytes + assert manager.get_cpu_transfer_layouts() is layouts + + +def test_cpu_bounded_capacity_includes_state_field_padding(cpu_args: EasyDict): + class PaddedGDNState(GDNState): + def get_cpu_transfer_regions(self): + return tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + if region.field == "ssm" else region + for region in super().get_cpu_transfer_regions() + ) + + manager = AttResourceManager(request_capacity=1) + state = PaddedGDNState( + ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + ), + request_capacity=1, device="meta", + ) + manager.set_att_state("main", state) + manager.finish_state_creation() + # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 + assert manager.get_cpu_page_bytes() == 80 + layout = manager.get_cpu_transfer_layouts()[state.resource_id] + assert (layout.byte_start, layout.byte_end) == (0, 80) + assert layout.regions[-1].byte_end == 76 + + +@pytest.mark.parametrize("kind", ["packed", "packed_fp8", "packed_glm"]) +def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: EasyDict): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cpu_args.cpu_cache_token_page_size = 3 + states = [] + for encoding in ("gqa", kind): + template = definition(encoding) + state = rebuild_full_state(template, config=template.config, token_capacity=8, device="cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.view(torch.uint8).fill_(211) + states.append(state) + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", states[0]) + manager.set_att_state("draft0", states[1]) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_creation() + layouts = manager.get_cpu_transfer_layouts() + packed, gqa = manager.get_full_attention_states() + packed_layout, gqa_layout = layouts[packed.resource_id], layouts[gqa.resource_id] + data_end = max(r.byte_end for r in packed_layout.regions) + gqa_start = gqa_layout.byte_start + # 后一个 State 紧跟实际字段终点，仅保留对齐空间，不为 GPU packed 行 padding 预留位置。 + assert gqa_start == (data_end + 7) // 8 * 8 + assert gqa_start < read_kv_storage(packed)[:, :3].numel() * read_kv_storage(packed).element_size() + for value, state in enumerate((packed, gqa), 3): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor[:, :3].fill_(value) + field.tensor[:, 4:7].zero_() + pages = torch.full((3, manager.get_cpu_page_bytes()), 165, dtype=torch.uint8, pin_memory=True) + manager.write_req_to_cpu(torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), pages, [1], layouts=layouts) + torch.cuda.synchronize() + assert (pages[1, data_end:gqa_start] == 165).all() + assert (pages[[0, 2]] == 165).all() + manager.read_req_from_cpu(pages, [1], torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), layouts=layouts) + for value, state in enumerate((packed, gqa), 3): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 4:7] == value).all() + + +@pytest.mark.parametrize("page_tokens", [1, 9]) +def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: EasyDict, page_tokens: int): + cpu_args.cpu_cache_token_page_size = page_tokens + manager = AttResourceManager(request_capacity=1) + manager.bounded_att_big_page_allocator = StatePageAllocator(2) + manager.bounded_att_small_page_allocator = StatePageAllocator(1) + for state_cls, role, value_dim in ((KDAState, "draft0", 1), (GDNState, "main", 3)): + config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=value_dim, conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + ) + state = state_cls(config, request_capacity=1, mtp_step=2, device="cpu") + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=False) + manager.set_att_state(role, state) + manager.finish_state_creation() + states = manager.get_bounded_attention_states() + layouts = manager.get_cpu_transfer_layouts() + # 两份快照分别含 10/22 字节数据，预留 16/24 字节；页 token 数不放大快照。 + assert [min(r.byte_start for r in layouts[state.resource_id].regions) for state in states] == [0, 16] + assert manager.get_cpu_page_bytes() == 48 + assert [(layout.byte_start, layout.byte_end) for layout in layouts.values()] == [(0, 16), (16, 40)] + for i, state in enumerate(states): + state.conv_state.fill_(i + 3) + state.ssm_state.fill_(i + 7) + source_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.save_bounded_att_pages([0], page_kind="big", slot_ids=source_slots) + pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8) + manager.write_req_to_cpu( + torch.tensor([], dtype=torch.long), pages, [1], layouts=layouts, big_page_slot_ids=source_slots + ) + assert (pages[0] == 87).all() + assert (pages[1, 12:16] == 87).all() + manager.reset_bounded_att_state([0]) + target_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.read_req_from_cpu( + pages, [1], torch.tensor([], dtype=torch.long), layouts=layouts, big_page_slot_ids=target_slots, req_idx=0 + ) + for i, state in enumerate(states): + assert (state.conv_state[:, 0, :, :1] == i + 3).all() + assert (state.ssm_state[:, 0] == i + 7).all() + + +@pytest.mark.parametrize("transfer", ["cpu", "pd_full_att", "pd_bounded_att_or_tail"]) +def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation( + monkeypatch: pytest.MonkeyPatch, + transfer: str, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = 3 + pd_args.cpu_cache_token_page_size = 3 + manager = AttResourceManager(request_capacity=1) + config = FullAttentionConfig((0,), 1, 3, torch.bfloat16) + manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + manager.set_att_state("main", GDNState( + ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.bfloat16, + ), + request_capacity=1, device="cpu", + )) + manager.finish_state_creation() + getter = getattr(manager, f"get_{transfer}_transfer_layouts") + before = getter() + assert getter() is before + state = manager.main_gdn_state if transfer == "pd_bounded_att_or_tail" else manager.main_full_att_bf16_state + before_layout = before[state.resource_id] + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.initialize_storage(8, 4, big_page_tokens=8, device="cpu") + after = getter() + after_layout = after[state.resource_id] + assert after is not before and getter() is after + assert after_layout is not before_layout + full_state = manager.main_full_att_bf16_state + assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta + page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() + if transfer == "pd_bounded_att_or_tail": + assert after_layout.byte_end == 8 and page_bytes == 40 + elif transfer == "cpu": + layout_end = max(layout.byte_end for layout in after.values()) + assert (layout_end + 15) // 16 * 16 == page_bytes + else: + assert after_layout.byte_end == page_bytes + assert not { + "cpu_page_tokens", "cpu_page_bytes", "cpu_transfer_layouts", "pd_page_bytes", "pd_transfer_layouts" + } & manager.__dict__.keys() + + +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_manager_uses_state_regions_and_adds_each_resource_offset(transfer: str, pd_args: EasyDict): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + class ReversedCpuRegionsState(GQAFullAttentionState): + def get_cpu_transfer_regions(self): + regions = super().get_cpu_transfer_regions() + # 本实现的 CPU 页先保存 V 再保存 K，PD 仍沿用 K/V 顺序。 + return tuple( + replace(region, byte_start=other.byte_start, byte_end=other.byte_end) for region, other in zip(regions, reversed(regions)) + ) + + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + config = FullAttentionConfig( + layer_ids=(0,), global_head_num=1, head_dim=2, dtype=torch.float32 + ) + manager = AttResourceManager(request_capacity=1) + prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") + prepared_full_state.initialize_storage(8) + manager.set_att_state("main", prepared_full_state) + prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") + prepared_full_state.initialize_storage(8) + manager.set_att_state("draft0", prepared_full_state) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_creation() + layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() + for index, state in enumerate(manager.get_full_attention_states()): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + state.k_storage[:, :4].fill_(index + 3) + state.v_storage[:, :4].fill_(index + 7) + # 1 层、4 tokens，K/V 各 1 head、每 head 2 个 FP32 元素。 + state_bytes = 4 * 2 * 2 * torch.float32.itemsize + k_offset = state_bytes // 2 if transfer == "cpu" else 0 + k_region, v_region = layouts[state.resource_id].regions + assert k_region.field == "k" and v_region.field == "v" + assert k_region.byte_start == index * state_bytes + k_offset + assert v_region.byte_start == index * state_bytes + state_bytes // 2 - k_offset + + page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() + pages = torch.zeros( + (1, page_bytes), dtype=torch.uint8, device="cpu" if transfer == "cpu" else "cuda", pin_memory=transfer == "cpu" + ) + if transfer == "cpu": + manager.write_req_to_cpu(torch.arange(4, dtype=torch.long, device="cuda"), pages, [0], layouts=layouts) + manager.read_req_from_cpu(pages, [0], torch.arange(4, 8, dtype=torch.long, device="cuda"), layouts=layouts) + else: + manager.write_to_pd_buffer(pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, dtype=torch.long, device="cuda")) + manager.read_from_pd_buffer(pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, 8, dtype=torch.long, device="cuda")) + for state in manager.get_full_attention_states(): + torch.testing.assert_close(read_kv_storage(state)[:, 4:8], read_kv_storage(state)[:, :4]) + assert not read_kv_storage(state)[:, 8:].any() diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py new file mode 100644 index 0000000000..1a5eadd06b --- /dev/null +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl +from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, +) +from unit_tests.common.att_resources.test_capacity import definition + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kind", ["gqa", "static_head", "packed_glm"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("capture", [False, True]) +def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, capture: bool): + template = definition(kind) + state = rebuild_full_state(template, config=replace(template.config, layer_ids=(3, 8)), token_capacity=8, device="cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + manager = AttResourceManager(request_capacity=2) + manager.set_att_state(role, state) + manager.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) + manager.finish_state_creation() + state = getattr(manager, f"{role}_{template.get_state_type()}_state") + infer = InferStateInfo() + infer.att_resource_manager = manager + infer.is_draft_model = role == "draft0" + infer.att_resource_role = role + infer.mem_index = torch.tensor([1, 5], dtype=torch.int32, device="cuda") + heads = 1 if kind == "packed_glm" else 2 * state.config.head_num + dim = state.config.kv_dim if kind == "packed_glm" else state.config.head_dim + kv = torch.full((2, heads, dim), 0.5, dtype=torch.bfloat16, device="cuda") + layer = SimpleNamespace(layer_num_=8, full_att_state_name=state.resource_id) + TransformerLayerInferTpl._post_cache_kv(layer, kv, infer, None) + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + TransformerLayerInferTpl._post_cache_kv(layer, kv, infer, None) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + graph.replay() + assert not hasattr(infer, "mem_manager") + assert not hasattr(infer, "full_att_state") + assert not read_kv_storage(state)[0].view(torch.uint8).any() + view = state.get_layer_view(8) + if kind == "packed_glm": + assert torch.equal(view.kv[infer.mem_index], kv) + assert not view.indexer_k.any() + else: + k, v = view + actual = torch.cat((k[infer.mem_index], v[infer.mem_index]), dim=1) + if kind == "static_head": + actual = actual.view(torch.float8_e4m3fn).float() * state.scales[1][None, :, None] + # Backends index fixed scales by physical row, even when model layer + # IDs are sparse or include appended draft layers. + assert BaseAttBackend._find_layer_index(None, k, v, state) == 1 + assert state.cpu_scales.shape == (2, 4) + assert torch.equal(actual, kv) + + +def test_packed_fp8_prefill_combines_history_and_fresh_kv_with_sparse_remapping(): + template = definition("packed_fp8") + state = rebuild_full_state(template, config=template.config, token_capacity=8, device="cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + history_ids = torch.tensor([3, 7], dtype=torch.int32, device="cuda") + history = torch.randn(2, 1, 576, dtype=torch.bfloat16, device="cuda") + state.write_full_att_kv(0, history_ids, history) + fresh = torch.full((1, 1, 576), 123, dtype=torch.bfloat16, device="cuda") + topk = torch.tensor([[7, 5, 3, 7, -1], [5, -1, -1, -1, -1]], dtype=torch.int32, device="cuda") + compact, mapped = get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=state.get_layer_view(0).kv, + topk_mem_indices=topk, + prefill_mem_index=torch.tensor([5], dtype=torch.int32, device="cuda"), + prefill_cache_kv=fresh, + prefill_dtype=fresh.dtype, + ) + assert compact.shape == (3, 1, 576) + assert torch.equal(mapped == -1, topk == -1) + assert mapped[0, 0] == mapped[0, 3] + assert torch.equal(compact[mapped[0, 1]], fresh[0]) + assert torch.equal(compact[mapped[1, 0]], fresh[0]) + torch.testing.assert_close(compact[mapped[0, 0], :, :512], history[1, :, :512], atol=0.15, rtol=0.08) + torch.testing.assert_close(compact[mapped[0, 2], :, :512], history[0, :, :512], atol=0.15, rtol=0.08) + assert torch.equal(compact[mapped[0, 0], :, 512:], history[1, :, 512:]) + assert torch.equal(compact[mapped[0, 2], :, 512:], history[0, :, 512:]) diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py new file mode 100644 index 0000000000..dc8f428d97 --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -0,0 +1,150 @@ +"""FP8 backends obtain calibration from the role's matching storage scheme.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.basemodel.attention.fa3 import fp8 as fa3 +from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer +from lightllm.common.basemodel.infer_struct import InferStateInfo +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from collections.abc import Sequence + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kind", ["static_head", "static_tensor"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( + monkeypatch: pytest.MonkeyPatch, + kind: str, + role: str, + prefill: bool, +): + template = definition(kind) + states = [] + for value in (2, 5): + scales = template.scales + value + scales[1] += 10 + q_scales = None if kind == "static_tensor" else template.q_scales + value + states.append(rebuild_full_state( + template, config=replace(template.config, layer_ids=(3, 8)), token_capacity=8, device="cuda", + scales=scales, q_scales=q_scales, + )) + manager = AttResourceManager(request_capacity=2) + manager.set_att_state("main", states[0]) + manager.set_att_state("draft0", states[1]) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + # Another encoding in the same role must not supply the FP8 calibration. + plain = definition("gqa") + plain = rebuild_full_state(plain, config=plain.config, token_capacity=8, device="cuda") + manager.set_att_state(role, plain) + manager.finish_state_creation() + selected = getattr(manager, f"{role}_{template.get_state_type()}_state") + k, v = selected.get_layer_view(8) + infer = InferStateInfo() + infer.att_resource_manager = manager + infer.is_draft_model = role == "draft0" + infer.att_resource_role = role + infer.batch_size, infer.max_q_seq_len = 2, 1 + infer.b_seq_len = torch.ones(2, dtype=torch.int32, device="cuda") + backend = object.__new__(BaseAttBackend) + observed = {} + + if kind == "static_head": + state_cls = fa3.Fp8Fa3PrefillAttState if prefill else fa3.Fp8Fa3DecodeAttState + parent_cls = fa3.Fa3PrefillAttState if prefill else fa3.Fa3DecodeAttState + monkeypatch.setattr(parent_cls, "init_state", lambda self: None) + att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + att.b_att_seq_len = infer.b_seq_len + att.init_state() + + def quantize(q: torch.Tensor, scales: torch.Tensor): + observed["q_scales"] = scales + return q + + def attention(**kwargs: object): + observed.update(kwargs) + return kwargs["q"] + + monkeypatch.setattr(fa3, "q_per_head_static_fp8_quant", quantize) + monkeypatch.setattr(fa3, "flash_attn_with_kvcache", attention) + else: + state_cls = flashinfer.Fp8FlashInferPrefillAttState if prefill else flashinfer.Fp8FlashInferDecodeAttState + att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + wrapper = SimpleNamespace(run=lambda *args, **kwargs: observed.update(kwargs)) + att.prefill_wrapper = att.decode_wrapper = wrapper + + q = torch.zeros((2, selected.config.head_num, selected.config.head_dim), dtype=torch.bfloat16, device="cuda") + + def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): + return torch.empty(shape, dtype=dtype, device=device) + + if prefill: + att.prefill_att(q, k, v, alloc_func=allocate) + else: + att.decode_att(q, k, v, alloc_func=allocate) + if kind == "static_head": + heads = selected.config.head_num + torch.testing.assert_close(observed["q_scales"], selected.q_scales[1]) + torch.testing.assert_close(observed["k_descale"], selected.scales[1, :heads].expand(2, heads)) + torch.testing.assert_close(observed["v_descale"], selected.scales[1, heads:].expand(2, heads)) + else: + assert observed["k_scale"] == selected.cpu_scales[1, 0] + assert observed["v_scale"] == selected.cpu_scales[1, 1] + assert not hasattr(infer, "full_att_state") + + +@pytest.mark.parametrize("head_num", [1, 3]) +@pytest.mark.parametrize("scale_mode", ["per_tensor", "per_head"]) +def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_num: int, scale_mode: str): + from lightllm.common.att_resources.full_att import StaticFP8AttentionConfig, StaticFP8FullAttentionState + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 + + scale_count = 2 if scale_mode == "per_tensor" else 2 * head_num + scales = torch.arange(1, scale_count + 1, dtype=torch.float32, device="cuda").unsqueeze(0) / 4 + state = StaticFP8FullAttentionState( + layout_config=StaticFP8AttentionConfig( + layer_ids=(5,), global_head_num=head_num, head_dim=128, dtype=torch.bfloat16, + scale_mode=scale_mode, calibration_id="split-storage", + ), + page_size=4, device="cuda", scales=scales, + q_scales=torch.ones((1, head_num), dtype=torch.float32, device="cuda") if scale_mode == "per_head" else None, + ) + state.initialize_storage(8) + state.k_storage.zero_() + state.v_storage.zero_() + indices = torch.tensor([1, 6], dtype=torch.int32, device="cuda") + kv = torch.randn((2, 2 * head_num, 128), dtype=torch.bfloat16, device="cuda") + state.write_full_att_kv(layer_id=5, token_indices=indices, kv=kv) + k, v = state.get_layer_view(5) + assert k.untyped_storage().data_ptr() != v.untyped_storage().data_ptr() + actual = torch.cat((k[indices], v[indices]), dim=1) + head_scales = scales[0].repeat_interleave(head_num) if scale_mode == "per_tensor" else scales[0] + expected = (kv.float() / head_scales[None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn) + assert torch.equal(actual, expected.view(torch.uint8)) + assert not k[[0, 2, 3, 4, 5, 7, 8, 9, 10, 11]].any() + assert not v[[0, 2, 3, 4, 5, 7, 8, 9, 10, 11]].any() + + # 算子的原合并输出入口仍应保持相同的量化结果。 + combined = torch.zeros((12, 2 * head_num, 128), dtype=torch.uint8, device="cuda") + destindex_copy_kv_fp8( + KV=kv, DestLoc=indices, scales=scales[0], Out=combined.view(torch.float8_e4m3fn), + is_per_tensor_quant=scale_mode == "per_tensor", + ) + assert torch.equal(combined[indices], actual) diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py new file mode 100644 index 0000000000..8ae0feb50e --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -0,0 +1,216 @@ +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +"""CPU 页直接拷贝：真实索引、跨页布局、原始编码与未参与区域的保持。""" + +from dataclasses import replace + +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers +from unit_tests.common.att_resources.test_state_identity import allocated_state +from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout +from lightllm.common.att_resources.types import AxisSlice + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", [ + "gqa", "mla", "int4", "int8", "static_head", "static_tensor", "packed", "packed_fp8", "packed_glm", +]) +@pytest.mark.parametrize("token_major", [False, True]) +def test_direct_cpu_pages_preserve_bytes_without_gather(kind: str, token_major: bool, monkeypatch: pytest.MonkeyPatch): + source = allocated_state(kind, device="cuda") + target = allocated_state(kind, device="cuda") + for tensor in full_state_test_buffers(source).values(): + tensor.view(torch.uint8).random_(0, 256) + for tensor in full_state_test_buffers(target).values(): + tensor.view(torch.uint8).fill_(173) + layout = transfer_layout(source) + if token_major: + regions = [] + for region in layout.regions: + axes = ("token", "layer", *region.axes[2:]) + regions.append(replace(region, axes=axes, shape=tuple(region.shape[region.axes.index(a)] for a in axes))) + layout = replace(layout, regions=tuple(regions)) + # 页池本身是更大联合页的视图，检验真实 page stride，而不是按字段长度推算下一页。 + pages = torch.full((4, layout.byte_end + 64), 91, dtype=torch.uint8, pin_memory=True) + page_view = pages[:, 32:-32] + src_ids = torch.tensor([6, 6, 1, 1, 7, 7, 0, 0, 3, 3], dtype=torch.int32, device="cuda")[::2] + dst_ids = torch.tensor([2, 2, 6, 6, 4, 4, 1, 1, 7, 7], dtype=torch.int64, device="cuda")[::2] + expected = {name: tensor[:, dst_ids].clone() for name, tensor in full_state_test_buffers(target).items()} + source_fields = full_state_test_fields(source, full_state_test_buffers(source)) + expected_fields = full_state_test_fields(target, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(field.tensor[:, src_ids]) + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("CPU transfer must not allocate gathered KV or scatter it back") + + tensor_to = torch.Tensor.to + tensor_contiguous = torch.Tensor.contiguous + + def reject_index_to(tensor: torch.Tensor, *args: object, **kwargs: object): + if tensor is src_ids or tensor is dst_ids: + pytest.fail("token indices must be passed directly to the kernel") + return tensor_to(tensor, *args, **kwargs) + + def reject_index_contiguous(tensor: torch.Tensor, *args: object, **kwargs: object): + if tensor is src_ids or tensor is dst_ids: + pytest.fail("token index views must not be materialized") + return tensor_contiguous(tensor, *args, **kwargs) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "to", reject_index_to) + patch.setattr(torch.Tensor, "contiguous", reject_index_contiguous) + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + with torch.cuda.stream(stream): + source.write_to_cpu(src_ids, page_view, [2, 0], layout=layout) + target.read_from_cpu(page_view, [2, 0], dst_ids, layout=layout) + stream.synchronize() + for name, tensor in full_state_test_buffers(target).items(): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + assert (tensor[:, [0, 3, 5, 8, 9, 10, 11]].contiguous().view(torch.uint8) == 173).all() + assert (pages[:, :32] == 91).all() and (pages[:, -32:] == 91).all() + assert (pages[[1, 3]] == 91).all() + for region in layout.regions: + external = page_view[0, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + # 最后一页只有一个有效 token，其余槽位仍保留填充值。 + unused = [slice(None)] * len(region.shape) + unused[region.axes.index("token")] = slice(1, None) + assert (external[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_direct_cpu_pages_only_update_selected_head_slice_and_skip_replica_writes(): + state = allocated_state("gqa", device="cuda") + state.k_storage.fill_(3) + state.v_storage.fill_(7) + layout = transfer_layout(state) + regions = [] + for region in layout.regions: + slices = tuple( + AxisSlice(axis="head", src_start=1, dst_start=0, length=1) if item.axis == "head" else item + for item in region.axis_slices + ) + regions.append(replace(region, axis_slices=slices)) + layout = replace(layout, regions=tuple(regions)) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) + ids = torch.tensor([1, 5], dtype=torch.int64, device="cuda") + state.write_to_cpu(ids, pages, [0], layout=layout) + torch.cuda.synchronize() + snapshot = pages.clone() + state.k_storage.fill_(19) + state.v_storage.fill_(19) + read_only = replace(layout, regions=tuple(replace(region, write_enabled=False) for region in layout.regions)) + state.write_to_cpu(ids, pages, [0], layout=read_only) + state.read_from_cpu(pages, [0], ids, layout=read_only) + torch.cuda.synchronize() + assert torch.equal(pages, snapshot) + assert (state.k_storage[:, ids, 1] == 3).all() and (state.v_storage[:, ids, 1] == 7).all() + assert (state.k_storage[:, :, 0] == 19).all() and (state.v_storage[:, :, 0] == 19).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_direct_cpu_pages_accept_registered_external_memory(): + """共享缓存通过 cudaHostRegister 注册，而非 torch 的 pinned allocator 分配。""" + import ctypes + import mmap + + state = allocated_state("gqa", device="cuda") + state.k_storage.fill_(3) + state.v_storage.fill_(7) + layout = transfer_layout(state) + memory = mmap.mmap(-1, layout.byte_end) + pages = torch.frombuffer(memory, dtype=torch.uint8).view(1, -1) + cuda = ctypes.CDLL("/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so") + cuda.cudaHostRegister.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint] + cuda.cudaHostUnregister.argtypes = [ctypes.c_void_p] + pointer = ctypes.c_void_p(pages.data_ptr()) + assert cuda.cudaHostRegister(pointer, layout.byte_end, 3) == 0 + try: + assert pages.is_pinned() + ids = torch.tensor([1, 5], dtype=torch.int32, device="cuda") + state.write_to_cpu(ids, pages, [0], layout=layout) + state.k_storage.zero_() + state.v_storage.zero_() + state.read_from_cpu(pages, [0], ids, layout=layout) + torch.cuda.synchronize() + assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() + finally: + torch.cuda.synchronize() + assert cuda.cudaHostUnregister(pointer) == 0 + del pages + memory.close() + + +@pytest.mark.parametrize("kind", ["gqa", "int8", "static_head", "mla", "packed_fp8"]) +@pytest.mark.parametrize("write", [True, False]) +def test_cpu_cache_transfer_rejects_cpu_state(kind: str, write: bool): + state = allocated_state(kind, device="cpu") + layout = transfer_layout(state) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8) + ids = torch.tensor([0], dtype=torch.int64) + with pytest.raises(AssertionError, match="CUDA State storage"): + if write: + state.write_to_cpu(ids, pages, [0], layout=layout) + else: + state.read_from_cpu(pages, [0], ids, layout=layout) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("token_count", [0, 1]) +def test_gqa_single_cpu_page_and_empty_transfer(token_count: int) -> None: + source = allocated_state("gqa", device="cuda") + target = allocated_state("gqa", device="cuda") + source.k_storage.fill_(3) + source.v_storage.fill_(7) + target.k_storage.fill_(19) + target.v_storage.fill_(19) + layout = transfer_layout(source) + pages = torch.full(size=(3, layout.byte_end), fill_value=91, dtype=torch.uint8, pin_memory=True) + src_ids = torch.tensor(data=[5], dtype=torch.int32, device="cuda")[:token_count] + dst_ids = torch.tensor(data=[2], dtype=torch.int64, device="cuda")[:token_count] + page_indices = [1] if token_count else [] + + source.write_to_cpu(src_ids, pages, page_indices, layout=layout) + target.read_from_cpu(pages, page_indices, dst_ids, layout=layout) + torch.cuda.synchronize() + if token_count: + assert (target.k_storage[:, 2] == 3).all() and (target.v_storage[:, 2] == 7).all() + assert (target.k_storage[:, [0, 1, 3, 4, 5, 6, 7]] == 19).all() + assert (target.v_storage[:, [0, 1, 3, 4, 5, 6, 7]] == 19).all() + else: + assert (target.k_storage == 19).all() and (target.v_storage == 19).all() + assert (pages[1] == 91).all() + assert (pages[[0, 2]] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", [ + "gqa", "mla", "int4", "int8", "static_head", "static_tensor", "packed", "packed_fp8", "packed_glm", +]) +def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: + state = allocated_state(kind, device="cuda") + layout = transfer_layout(state) + cpu_ids = torch.tensor([0], dtype=torch.int32) + cuda_ids = torch.tensor([1], dtype=torch.int64, device="cuda") + cpu_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) + pd_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_full_att_kv(state.config.layer_ids[0], cpu_ids, torch.empty(0, device="cuda")) + for src, dst in ((cpu_ids, cuda_ids), (cuda_ids, cpu_ids)): + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_full_att_tokens(src, dst) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_tokens_from_other_dp(src, dst, state) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_to_cpu(cpu_ids, cpu_pages, [0], layout=layout) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.read_from_cpu(cpu_pages, [0], cpu_ids, layout=layout) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_to_pd_buffer(cpu_ids, pd_pages, pd_page_index=0, layout=layout) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.read_from_pd_buffer(pd_pages, cpu_ids, pd_page_index=0, layout=layout) diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py new file mode 100644 index 0000000000..68610c1825 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -0,0 +1,281 @@ +"""各 full-attention 编码的统一 CPU/PD kernel：无中间 KV、跳步循环与字段完整性。""" + +from __future__ import annotations + +from dataclasses import replace +from functools import partial +from importlib import import_module +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState +from lightllm.common.att_resources.types import StateTransferLayout +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, full_state_test_fields, rebuild_full_state +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +SCHEMES = { + "int4": "group_quantized", + "int8": "group_quantized", + "static_head": "static_fp8", + "static_tensor": "static_fp8", + "mla": "mla", + "packed": "packed_mla", + "packed_fp8": "packed_mla", + "packed_glm": "packed_mla", +} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", list(SCHEMES)) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +@pytest.mark.parametrize("grid_num", [1, 36]) +@pytest.mark.parametrize("token_count", [1, 513]) +def test_full_att_pages_copy_all_fields_with_bounded_grid( + kind: str, + transfer: str, + grid_num: int, + token_count: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + page_tokens = 128 if transfer == "cpu" else 1024 + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = page_tokens + template = definition(kind) + source = rebuild_full_state(template, token_capacity=1024, device="cuda") + target = rebuild_full_state(template, token_capacity=1024, device="cuda") + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + for buffer in source_buffers.values(): + buffer.view(torch.uint8).random_(0, 256) + for buffer in target_buffers.values(): + buffer.view(torch.uint8).fill_(173) + calibration = [] + if isinstance(source, StaticFP8FullAttentionState): + for state in (source, target): + for tensor in (state.scales, state.q_scales): + if tensor is not None: + calibration.append((tensor, tensor.clone())) + + regions = source.get_cpu_transfer_regions() if transfer == "cpu" else source.get_pd_transfer_regions() + if grid_num == 1: + # 调换层/token 轴，验证寻址使用真实 stride；region 顺序也不要求与存储成员相同。 + regions = tuple(replace( + region, + axes=("token", "layer", *region.axes[2:]), + shape=(region.shape[1], region.shape[0], *region.shape[2:]), + ) for region in reversed(regions)) + layout = StateTransferLayout( + resource_id=source.resource_id, + tokens_per_page=page_tokens, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + page_count = (token_count + page_tokens - 1) // page_tokens + page_indices = [3, 0, 4, 1, 6][:page_count] + pages = torch.full( + size=(8, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + page_view = pages[:, 16:-16] + src_ids = torch.randperm(n=1024, device="cuda")[:token_count].to(dtype=torch.int32) + dst_order = torch.randperm(n=1024, device="cuda") + dst_ids = dst_order[:token_count] + expected = {name: tensor[:, dst_ids].clone() for name, tensor in target_buffers.items()} + source_fields = full_state_test_fields(source, source_buffers) + expected_fields = full_state_test_fields(target, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(other=field.tensor[:, src_ids]) + + # 直接覆盖 State 调用的入口参数，读写仍通过生产 State 接口提交。 + scheme = SCHEMES[kind] + state_module = import_module(name=type(source).__module__) + transfer_name = f"transfer_{scheme}_pages" + monkeypatch.setattr(state_module, transfer_name, partial(getattr(state_module, transfer_name), grid_num=grid_num)) + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("page transfers must not gather/scatter KV or copy token indices to host") + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + patch.setattr(torch.Tensor, "tolist", reject_gather) + with torch.cuda.stream(stream): + if transfer == "cpu": + source.write_to_cpu(token_indices=src_ids, cpu_page_buffers=page_view, cpu_page_indices=page_indices, layout=layout) + target.read_from_cpu(cpu_page_buffers=page_view, cpu_page_indices=page_indices, token_indices=dst_ids, layout=layout) + else: + source.write_to_pd_buffer(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout) + target.read_from_pd_buffer(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, layout=layout) + stream.synchronize() + + # 比较整行原始字节：packed padding 沿用目标填充值，字段数据来自源存储。 + for name, tensor in target_buffers.items(): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + assert (tensor[:, dst_order[token_count:]].contiguous().view(torch.uint8) == 173).all() + assert (tensor[:, 1024:].view(torch.uint8) == 173).all() + for tensor, before in calibration: + assert torch.equal(tensor, before) + assert (pages[[index for index in range(8) if index not in page_indices]] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() + last_count = token_count - (page_count - 1) * page_tokens + for region in regions: + last_page = page_view[page_indices[-1], region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + unused = [slice(None)] * 4 + unused[region.axes.index("token")] = slice(last_count, None) + assert (last_page[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_all_full_att_schemes_transfer_through_peer_ipc( + pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +) -> None: + """各编码字段通过生产 IPC peer 映射跨 GPU 搬运，padding 保留目标值。""" + from multiprocessing.reduction import ForkingPickler + import torch.multiprocessing as mp + + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + pd_args.pd_kv_page_size = 4 + states, layouts, expected_buffers = [], [], [] + for kind in SCHEMES: + state = rebuild_full_state(definition(kind), token_capacity=8, device="cuda:1") + buffers = full_state_test_buffers(state) + for tensor in buffers.values(): + tensor.view(torch.uint8).random_(0, 256) + tensor[:, 5].view(torch.uint8).fill_(173) + expected = {name: tensor[:, 5:6].clone() for name, tensor in buffers.items()} + source_fields = full_state_test_fields(state, buffers) + expected_fields = full_state_test_fields(state, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(other=field.tensor[:, 1:2]) + regions = state.get_pd_transfer_regions() + layouts.append(StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + )) + states.append(state) + expected_buffers.append(expected) + torch.cuda.synchronize(device=1) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps(states)) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_peer_transfer_worker, args=(child, payload, layouts)) + process.start() + child.close() + try: + assert parent.poll(60), "PD peer worker did not finish" + error = parent.recv() + assert error is None, error + for state, expected in zip(states, expected_buffers): + for name, tensor in full_state_test_buffers(state).items(): + assert torch.equal(tensor[:, 5:6].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _peer_transfer_worker( + connection: Connection, payload: bytes, layouts: list[StateTransferLayout], +) -> None: + from multiprocessing.reduction import ForkingPickler + import traceback + + try: + torch.cuda.set_device(device=0) + states = ForkingPickler.loads(payload) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + source_ids = torch.tensor(data=[1], dtype=torch.int32, device="cuda:0") + target_ids = torch.tensor(data=[5], dtype=torch.int64, device="cuda:0") + for state, layout in zip(states, layouts): + pages = torch.empty(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda:0") + state.write_to_pd_buffer( + token_indices=source_ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_buffer( + pd_page_buffers=pages, + token_indices=target_ids, + pd_page_index=1, + layout=layout, + ) + stream.synchronize() + connection.send(None) + connection.recv() + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["gqa", *SCHEMES]) +@pytest.mark.parametrize("readies", [[False, True, False], [True, False, True], [True, True, True], [False, False, False]]) +@pytest.mark.parametrize("grid_num", [1, 36]) +def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( + kind: str, readies: list[bool], grid_num: int, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +) -> None: + pd_args.cpu_cache_token_page_size = 5 + state = rebuild_full_state(definition(kind), token_capacity=32, device="cuda") + # 单 program 验证跳过 ready 块后仍会处理后续未就绪块；默认 grid 保持生产配置。 + state_module = import_module(name=type(state).__module__) + scheme = "gqa" if kind == "gqa" else SCHEMES[kind] + transfer_name = f"transfer_{scheme}_pages" + monkeypatch.setattr(state_module, transfer_name, partial(getattr(state_module, transfer_name), grid_num=grid_num)) + for buffer in full_state_test_buffers(state).values(): + buffer.view(torch.uint8).random_(0, 256) + regions = state.get_cpu_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, tokens_per_page=5, byte_start=0, + byte_end=max(region.byte_end for region in regions), regions=regions, + ) + # 页编号不连续，token 索引带 stride；最后一页只写三个 token。 + page_indices = [3, 0, 4] + token_indices = torch.arange(26, dtype=torch.int32, device="cuda")[::2] + baseline = torch.full((6, layout.byte_end + 32), 173, dtype=torch.uint8, pin_memory=True) + pages = torch.full_like(baseline, 173, pin_memory=True) + state.write_to_cpu(token_indices, baseline[:, 16:-16], page_indices, layout=layout) + torch.cuda.synchronize() + expected = pages.clone() + for page, ready in zip(page_indices, readies): + if not ready: + expected[page].copy_(baseline[page]) + + def reject_allocation(*args: object, **kwargs: object) -> None: + pytest.fail("all-ready writes must not allocate index or ready tensors") + + with monkeypatch.context() as patch: + if all(readies): + patch.setattr(torch, "tensor", reject_allocation) + state.write_to_cpu( + token_indices, pages[:, 16:-16], page_indices, layout=layout, cpu_page_readies=readies, + ) + torch.cuda.synchronize() + # 整页比较同时覆盖 KV、量化 scale、packed indexer 字段及前后 padding。 + assert torch.equal(pages, expected) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py new file mode 100644 index 0000000000..f59c6e8341 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_state.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + + +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources.full_att import ( + FullAttentionState, + GQAFullAttentionState, + MLAFullAttentionState, + FullAttentionConfig, +) +from lightllm.common.att_resources.types import ( + AxisSlice, StateTransferLayout, TransferRegion, +) + + +def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): + prepared_full_state = cls(FullAttentionConfig((2, 5), 1, 4, torch.float32), page_size=4, device=device) + prepared_full_state.initialize_storage(12) + return prepared_full_state + + +def layout_for(state: FullAttentionState): + regions = [] + offset = 16 + for name in full_state_test_fields(state, full_state_test_buffers(state)): + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=offset + 128, + dtype=torch.float32, + axes=("token", "layer", "head", "dim"), + shape=(4, 2, 1, 4), + axis_slices=( + AxisSlice(axis="head", src_start=0, dst_start=0, length=1), + AxisSlice(axis="dim", src_start=0, dst_start=0, length=4), + ), + ) + ) + offset += 128 + return StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=0, + byte_end=offset + 16, + regions=tuple(regions), + ) + + +@pytest.mark.parametrize("cls", [GQAFullAttentionState, MLAFullAttentionState]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_full_attention_write_copy_and_pd( + cls: type[FullAttentionState], device: str | torch.device | int, monkeypatch: pytest.MonkeyPatch +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = full_state(cls, device) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + ids = torch.tensor([1, 3], dtype=torch.int32, device=device) + data = torch.arange(2 * read_kv_storage(state).shape[2] * 4, dtype=torch.float32, device=device).reshape( + 2, read_kv_storage(state).shape[2], 4 + ) + if device == "cuda": + state.write_full_att_kv(2, ids, data) + state.write_full_att_kv(5, ids, data + 100) + else: + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_full_att_kv(2, ids, data) + # CPU 仅构造复制和传输测试的数据，不执行推理 KV 写入。 + for layer, values in ((0, data), (1, data + 100)): + head_start = 0 + for buffer in full_state_test_buffers(state).values(): + head_end = head_start + buffer.shape[2] + buffer[layer, ids.long()] = values[:, head_start:head_end] + head_start = head_end + def reject_host_list(tensor: torch.Tensor): + raise AssertionError("token copies must not convert indices to Python lists") + + if device == "cuda": + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) + assert torch.equal(read_kv_storage(state)[:, [1, 3]], read_kv_storage(state)[:, [8, 9]]) + else: + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) + layout = layout_for(state) + pages = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, device=device) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + state.write_to_pd_buffer(ids, pages, pd_page_index=1, layout=layout) + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + state.read_from_pd_buffer(pages, ids, pd_page_index=1, layout=layout) + return + state.write_to_pd_buffer(ids, pages, pd_page_index=1, layout=layout) + target = full_state(cls, device) + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.zero_() + target.read_from_pd_buffer(pages, torch.tensor([5, 6], dtype=torch.long, device=device), pd_page_index=1, layout=layout) + assert torch.equal(read_kv_storage(target)[:, [5, 6]], read_kv_storage(state)[:, [1, 3]]) + assert (pages[0] == 255).all() and (pages[2] == 255).all() + assert not read_kv_storage(target)[:, 0].any() + + +@pytest.mark.parametrize("state_offset, extra_page_bytes", [(0, 0), (32, 64)]) +def test_cpu_partial_last_page(state_offset: int, extra_page_bytes: int): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = full_state(device="cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape_as(storage_buffer)) + layout = layout_for(state) + layout = replace( + layout, + byte_start=layout.byte_start + state_offset, + byte_end=layout.byte_end + state_offset, + regions=tuple( + replace(r, byte_start=r.byte_start + state_offset, byte_end=r.byte_end + state_offset) + for r in layout.regions + ), + ) + buffer = torch.full((4, layout.byte_end + extra_page_bytes), 252, dtype=torch.uint8, pin_memory=True) + state.write_to_cpu(torch.tensor([0, 1, 2, 3, 7], dtype=torch.long, device="cuda"), buffer, [2, 0], layout=layout) + target = full_state(device="cuda") + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.zero_() + target.read_from_cpu(buffer, [2, 0], torch.tensor([6, 7, 8, 9, 10], dtype=torch.long, device="cuda"), layout=layout) + assert torch.equal(read_kv_storage(target)[:, [6, 7, 8, 9, 10]], read_kv_storage(state)[:, [0, 1, 2, 3, 7]]) + assert (buffer[[1, 3]] == 252).all() + assert (buffer[:, :layout.byte_start] == 252).all() and (buffer[:, layout.byte_end:] == 252).all() + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_invalid_late_region_shape_cannot_partially_write_pd_page(): + state = full_state(device="cuda") + layout = layout_for(state) + region = layout.regions[-1] + region = replace(region, byte_end=region.byte_end - region.dtype.itemsize) + bad = replace(layout, regions=(*layout.regions[:-1], region)) + # 直接构造 Tensor 视图时，字节数与 shape 不匹配仍会在提交复制前失败。 + buffer = torch.full((layout.byte_end + 256,), 53, dtype=torch.uint8, device="cuda") + with pytest.raises(RuntimeError, match="shape"): + state.write_to_pd_buffer(torch.tensor([0], dtype=torch.long, device="cuda"), buffer.unsqueeze(0), pd_page_index=0, layout=bad) + assert (buffer == 53).all() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + config = FullAttentionConfig((1, 3, 5, 7, 9, 11), 1, 4, torch.float32) + source = GQAFullAttentionState(config, page_size=4, device=device) + source.initialize_storage(12) + target = GQAFullAttentionState(config, page_size=4, device=device) + target.initialize_storage(12) + for storage_buffer in full_state_test_buffers(source).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel(), device=device).reshape_as(storage_buffer)) + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.fill_(-1) + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + target.copy_tokens_from_other_dp( + src_token_indices=torch.tensor([2, 0]), dst_token_indices=torch.tensor([6, 7]), + remote_dp_full_att_state=source, + ) + return + target.copy_tokens_from_other_dp( + src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), + dst_token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + remote_dp_full_att_state=source, + ) + torch.testing.assert_close(read_kv_storage(target)[:, [6, 7]], read_kv_storage(source)[:, [2, 0]]) + assert (read_kv_storage(target)[:, 8:] == -1).all() + assert (read_kv_storage(target)[:, :6] == -1).all() + + # 拒绝将单层数据广播到多层，失败时目标 token 保持原值。 + single_layer_source = GQAFullAttentionState( + replace(config, layer_ids=config.layer_ids[:1]), page_size=4, device=device + ) + single_layer_source.initialize_storage(12) + for storage_buffer in full_state_test_buffers(single_layer_source).values(): + storage_buffer.fill_(9) + with pytest.raises(AssertionError, match="equivalent TP state configurations"): + target.copy_tokens_from_other_dp( + src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), + dst_token_indices=torch.tensor([8, 9], dtype=torch.long, device=device), + remote_dp_full_att_state=single_layer_source, + ) + assert (read_kv_storage(target)[:, 8:] == -1).all() diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py new file mode 100644 index 0000000000..e75a38790b --- /dev/null +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -0,0 +1,96 @@ +"""Real model initialization, compute writes, and profiled allocation.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +import json + +from unit_tests.common.att_resources.model_helpers import prepare_model_states + +import pytest +import torch + +from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState +from lightllm.utils.envs_utils import get_env_start_args +from unit_tests.common.att_resources.test_capacity import allocated_bytes +from unit_tests.server.router.model_infer.test_cpu_resource_layout import CASES, make_runtime, runtime + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +def test_model_storage_write_and_copy_preserve_all_fields( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], +): + _, (resources,) = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, 1) + (state,) = resources.get_full_attention_states("main") + config = state.config + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") + dim = config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else config.head_dim + kv = torch.randn(2, (1 if "kv" in full_state_test_buffers(state) else 2) * config.head_num, dim, device="cuda", dtype=torch.bfloat16) + state.write_full_att_kv(config.layer_ids[-1], ids, kv) + if len(config.layer_ids) > 1: + assert not read_kv_storage(state)[0].any() + resources.copy_full_att_tokens(ids, torch.tensor([1, 5], dtype=torch.int32, device="cuda")) + for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): + if "token" in field.axes: + assert torch.equal( + field.tensor[:, ids].contiguous().view(torch.uint8), + field.tensor[:, [1, 5]].contiguous().view(torch.uint8), + ), name + assert state.capacity == resources.full_att_allocator.size == 16 + assert read_kv_storage(state).shape[1] == 20 # Includes the whole HOLD page. + + +@pytest.mark.parametrize("mode", ["normal", "prefill"]) +@pytest.mark.parametrize("kind,encoding,model_cls", CASES[:8]) +def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + mode: str, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], +): + _, (reference,) = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, 1) + (state,) = reference.get_full_attention_states("main") + args = get_env_start_args() + args.run_mode, args.pd_kv_page_size, args.pd_kv_page_num = mode, 8, 3 + transfer_bytes = reference.get_pd_page_bytes() * 3 if mode == "prefill" else 0 + storage_bytes = allocated_bytes(state) + budget = storage_bytes + transfer_bytes + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (budget, 1 << 30)) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + model = model_cls.__new__(model_cls) + model.args, model.config = args, json.loads((tmp_path / "config.json").read_text()) + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.max_total_token_num, model.max_req_num, model.mem_fraction = None, 1, 1.0 + prepare_model_states(model) + model.att_resource_manager.initialize_resources( + args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + (state,) = model.att_resource_manager.get_full_attention_states("main") + assert not hasattr(model, "full_att_state") + assert state.capacity == 16 + assert allocated_bytes(state) == storage_bytes + if mode == "prefill": + assert model.att_resource_manager.get_pd_page_bytes() == reference.get_pd_page_bytes() diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py new file mode 100644 index 0000000000..76b698b619 --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -0,0 +1,223 @@ +"""GQA PD 直接搬运：原始位模式、物理页选择、TP 切片及有效 token 范围。""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.att_resources.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.att_resources.types import StateTransferLayout + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("token_major", [False, True]) +@pytest.mark.parametrize("token_count", [1, 3]) +def test_gqa_pd_direct_copy_between_tp_layouts( + dtype: torch.dtype, + token_major: bool, + token_count: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pd_args.pd_kv_page_size = 4 + # P 保存完整 heads，D 只恢复 rank 1 的一半 heads；两侧外部页格式相同。 + source = GQAFullAttentionState( + layout_config=FullAttentionConfig(layer_ids=(2, 5), global_head_num=4, head_dim=128, dtype=dtype), + page_size=4, + device="cuda", + ) + target = GQAFullAttentionState( + layout_config=replace(source.config, tp_world_size=2, tp_rank=1), + page_size=4, + device="cuda", + ) + for state in (source, target): + state.initialize_storage(token_capacity=8) + for tensor in (source.k_storage, source.v_storage): + tensor.view(torch.uint8).random_(0, 256) + for tensor in (target.k_storage, target.v_storage): + tensor.view(torch.uint8).fill_(173) + layouts = [] + for state in (source, target): + regions = [] + for region in state.get_pd_transfer_regions(): + axes = ("token", "layer", "head", "dim") if token_major else region.axes + regions.append(replace( + region, + byte_start=region.byte_start + 16, + byte_end=region.byte_end + 16, + axes=axes, + shape=tuple(region.shape[region.axes.index(axis)] for axis in axes), + )) + layouts.append(StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=16, + byte_end=max(region.byte_end for region in regions), + regions=tuple(regions), + )) + layout, target_layout = layouts + pages = torch.full(size=(3, layout.byte_end + 64), fill_value=91, dtype=torch.uint8, device="cuda") + page_view = pages[:, 32:-32] # 页池视图的真实跨度大于可见宽度。 + src_ids = torch.tensor(data=[6, 9, 1, 9, 7, 9], dtype=torch.int32, device="cuda")[::2][:token_count] + dst_ids = torch.tensor(data=[2, 6, 4], dtype=torch.int64, device="cuda")[:token_count] + expected = [tensor[:, src_ids, 2:].contiguous().view(torch.uint8) for tensor in (source.k_storage, source.v_storage)] + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("PD transfer must not gather or scatter intermediate KV tensors") + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + patch.setattr(torch.Tensor, "tolist", reject_gather) + with torch.cuda.stream(stream): + source.write_to_pd_buffer(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=1, layout=layout) + target.read_from_pd_buffer(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=1, layout=target_layout) + stream.synchronize() + for tensor, expected_bytes in zip((target.k_storage, target.v_storage), expected): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected_bytes) + assert (tensor[:, [index for index in range(12) if index not in [2, 6, 4][:token_count]]].contiguous().view(torch.uint8) == 173).all() + assert (pages[[0, 2]] == 91).all() + assert (pages[:, :48] == 91).all() and (pages[:, -32:] == 91).all() + for region in layout.regions: + external = page_view[1, region.byte_start:region.byte_end].view(dtype).view(region.shape) + unused = [slice(None)] * 4 + unused[region.axes.index("token")] = slice(token_count, None) + assert (external[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + # 反向验证分片写出：rank 1 只修改全局 heads [2, 4)，rank 0 区域保持原始字节。 + pages.fill_(91) + target.write_to_pd_buffer( + token_indices=dst_ids, pd_page_buffers=page_view, pd_page_index=2, layout=target_layout, + ) + for region, expected_bytes in zip(target_layout.regions, expected): + external = page_view[2, region.byte_start:region.byte_end].view(dtype).view(region.shape) + external = external.permute(tuple(region.axes.index(axis) for axis in ("layer", "token", "head", "dim"))) + assert torch.equal(external[:, :token_count, 2:].contiguous().view(torch.uint8), expected_bytes) + assert (external[:, :, :2].contiguous().view(torch.uint8) == 91).all() + assert (external[:, token_count:].contiguous().view(torch.uint8) == 91).all() + assert (pages[:2] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: + pd_args.pd_kv_page_size = 4 + # MQA 的 KV head 在 TP ranks 间复制，非首 rank 禁止写出但仍能读回。 + state = GQAFullAttentionState( + layout_config=FullAttentionConfig( + layer_ids=(0, 2), global_head_num=1, head_dim=128, dtype=torch.bfloat16, tp_world_size=2, tp_rank=1, + ), + page_size=4, + device="cuda", + ) + state.initialize_storage(token_capacity=8) + state.k_storage.fill_(19) + state.v_storage.fill_(19) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, tokens_per_page=4, byte_start=0, + byte_end=max(region.byte_end for region in regions), regions=regions, + ) + pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda") + for region, value in zip(regions, (3, 7)): + pages[1, region.byte_start:region.byte_end].view(region.dtype).fill_(value) + snapshot = pages.clone() + ids = torch.tensor(data=[5, 1], dtype=torch.int32, device="cuda") + state.write_to_pd_buffer(token_indices=ids, pd_page_buffers=pages, pd_page_index=1, layout=layout) + state.read_from_pd_buffer(pd_page_buffers=pages, token_indices=ids, pd_page_index=1, layout=layout) + assert torch.equal(pages, snapshot) + assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() + assert (state.k_storage[:, [0, 2, 3, 4, 6, 7]] == 19).all() + assert (state.v_storage[:, [0, 2, 3, 4, 6, 7]] == 19).all() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch) -> None: + """使用生产的 IPC peer 映射，从传输 GPU 直接读写另一块 GPU 上的 K/V。""" + from multiprocessing.reduction import ForkingPickler + import torch.multiprocessing as mp + + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + pd_args.pd_kv_page_size = 4 + state = GQAFullAttentionState( + layout_config=FullAttentionConfig(layer_ids=(0, 2), global_head_num=2, head_dim=128, dtype=torch.bfloat16), + page_size=4, + device="cuda:1", + ) + state.initialize_storage(token_capacity=8) + state.k_storage.zero_() + state.v_storage.zero_() + state.k_storage[:, 1].fill_(3) + state.v_storage[:, 1].fill_(7) + torch.cuda.synchronize(device=1) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, tokens_per_page=4, byte_start=0, + byte_end=max(region.byte_end for region in regions), regions=regions, + ) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps(state)) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_peer_pd_worker, args=(child, payload, layout)) + process.start() + child.close() + try: + assert parent.poll(45), "PD peer worker did not finish" + error = parent.recv() + assert error is None, error + assert (state.k_storage[:, 5] == 3).all() and (state.v_storage[:, 5] == 7).all() + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _peer_pd_worker(connection: Connection, payload: bytes, layout: StateTransferLayout) -> None: + from multiprocessing.reduction import ForkingPickler + import traceback + + try: + torch.cuda.set_device(device=0) + state = ForkingPickler.loads(payload) + pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda:0") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.write_to_pd_buffer( + token_indices=torch.tensor(data=[1], dtype=torch.int32, device="cuda:0"), + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_buffer( + pd_page_buffers=pages, + token_indices=torch.tensor(data=[5], dtype=torch.int64, device="cuda:0"), + pd_page_index=1, + layout=layout, + ) + stream.synchronize() + assert not pages[0].any() + connection.send(None) + connection.recv() # 保持导入的存储存活，直到父进程检查原 GPU 上的数据。 + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py new file mode 100644 index 0000000000..b8cd1e170f --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -0,0 +1,81 @@ +"""GQA 共用搬运 kernel 在受限 grid 下完整覆盖多轮、跨页和尾块数据。""" + +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.att_resources.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.att_resources.full_att.gqa_transfer import transfer_gqa_pages +from lightllm.common.att_resources.types import StateTransferLayout + +if TYPE_CHECKING: + from easydict import EasyDict + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("grid_num", [1, 7, 36]) +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_args: "EasyDict") -> None: + pd_args.pd_kv_page_size = 32 + config = FullAttentionConfig(layer_ids=(0, 2, 4), global_head_num=3, head_dim=128, dtype=torch.bfloat16) + source = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") + source.initialize_storage(token_capacity=128) + target.initialize_storage(token_capacity=128) + for tensor in (source.k_storage, source.v_storage): + tensor.view(torch.uint8).random_(0, 256) + for tensor in (target.k_storage, target.v_storage): + tensor.view(torch.uint8).fill_(173) + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=source.resource_id, + tokens_per_page=32, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + # 3 layers × 97 tokens × 3 heads × 128 dims，超过 36×256，且最后一块只有 128 个元素。 + token_count = 97 + src_ids = torch.randperm(n=128, device="cuda")[:token_count].to(dtype=torch.int32) + dst_order = torch.randperm(n=128, device="cuda") + dst_ids = dst_order[:token_count] + pages = torch.full( + size=(6, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, 16:-16] + page_indices = [3, 0, 4, 1] + transfer_gqa_pages( + k_storage=source.k_storage, + v_storage=source.v_storage, + token_indices=src_ids, + page_buffers=page_view, + page_indices=page_indices, + layout=layout, + write=True, + grid_num=grid_num, + ) + transfer_gqa_pages( + k_storage=target.k_storage, + v_storage=target.v_storage, + token_indices=dst_ids, + page_buffers=page_view, + page_indices=page_indices, + layout=layout, + write=False, + grid_num=grid_num, + ) + torch.cuda.synchronize() + for src, dst in ((source.k_storage, target.k_storage), (source.v_storage, target.v_storage)): + assert torch.equal(dst[:, dst_ids].contiguous().view(torch.uint8), src[:, src_ids].contiguous().view(torch.uint8)) + assert (dst[:, dst_order[token_count:]].contiguous().view(torch.uint8) == 173).all() + assert (dst[:, 128:].view(torch.uint8) == 173).all() + assert (pages[[2, 5]] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() + for region in regions: + last_page = page_view[1, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + assert (last_page[:, 1:].contiguous().view(torch.uint8) == 91).all() diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py new file mode 100644 index 0000000000..84480681a8 --- /dev/null +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.att_resources import AttResourceManager +from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager +from lightllm.common.att_resources.att_tail import KPoolAttentionTailState +from unit_tests.common.att_resources.test_full_state import full_state, layout_for +from unit_tests.common.att_resources.test_bounded_state import page_layout +from unit_tests.common.att_resources.test_tail_state import make_tail, tail_layout + +if TYPE_CHECKING: + from easydict import EasyDict + + +def composition(device: str | torch.device | int = "cpu", *, with_tail: bool = True, tail_pool_size: int = 4): + return MainDraftTestResourceManager(device, with_tail=with_tail, tail_pool_size=tail_pool_size) + + +def layouts_for(manager: AttResourceManager, *, kind: str = "cpu"): + layouts = [] + if kind != "att_state": + layouts += [layout_for(state) for state in (manager.main_full_att_fp32_state, manager.draft0_mla_fp32_state)] + if kind != "kv": + layouts += [page_layout(state) for state in (manager.main_kda_state, manager.draft0_gdn_state)] + if kind == "att_state": + layouts += [tail_layout(state) for state in manager.get_attention_tail_states("main")] + offset = 0 + result = {} + for layout in layouts: + state_end = offset + layout.byte_end + result[layout.resource_id] = replace( + layout, + byte_start=offset + layout.byte_start, + byte_end=state_end, + tokens_per_page=None if kind == "att_state" else 4, + regions=tuple( + replace(r, byte_start=r.byte_start + offset, byte_end=r.byte_end + offset) for r in layout.regions + ), + ) + offset = state_end + return result + + +def init_source(manager: AttResourceManager): + for i, state in enumerate((manager.main_full_att_fp32_state, manager.draft0_mla_fp32_state)): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.copy_( + torch.arange(storage_buffer.numel(), device=storage_buffer.device).reshape_as(storage_buffer) + i + ) + for i, state in enumerate((manager.main_kda_state, manager.draft0_gdn_state)): + state.conv_state[:, 0].fill_(11 + i) + state.conv_state[:, 1].fill_(23 + i) + state.ssm_state[:, :3].fill_(31 + i) + state.ssm_state[:, 3:6].fill_(47 + i) + for state in manager.get_attention_tail_states("main"): + state.tail_buffer[:, 0].fill_(3) + state.tail_buffer[:, 1].fill_(9) + slots = manager.bounded_att_big_page_allocator.alloc(2) + manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) + if manager.main_kda_state.conv_state.is_cuda: + torch.cuda.synchronize() + return slots + + +def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.MonkeyPatch): + manager = composition() + slots = init_source(manager) + manager.restore_bounded_att_pages([2], page_kind="big", slot_ids=[slots[0]]) + for i, state in enumerate((manager.main_kda_state, manager.draft0_gdn_state)): + assert (state.conv_state[:, 2, :, :2] == 11 + i).all() + assert (state.ssm_state[:, 6] == 31 + i).all() + + def fail(*args: object, **kwargs: object): + raise RuntimeError("simulated copy failure") + + monkeypatch.setattr(manager.draft0_gdn_state, "save_bounded_att_pages", fail) + with pytest.raises(RuntimeError): + manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) + assert manager.bounded_att_big_page_allocator.free_count == 1 + manager.bounded_att_big_page_allocator.free(slots) + assert manager.bounded_att_big_page_allocator.free_count == 3 + + +def test_disabled_domains(): + manager = AttResourceManager(request_capacity=3) + manager.set_att_state("main", full_state()) + manager.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) + manager.finish_state_creation() + manager.reset_bounded_att_state([0]) + manager.reset_attention_tail_state([0]) + with pytest.raises(AssertionError, match="disabled"): + manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) + empty_manager = AttResourceManager(request_capacity=3) + empty_manager.finish_state_creation() + with pytest.raises(AssertionError, match="no full-attention"): + empty_manager.copy_full_att_tokens(torch.tensor([0], dtype=torch.long), torch.tensor([1], dtype=torch.long)) + + +def test_reset_tail_clears_all_tail_states_and_preserves_linear(): + first = make_tail() + second = KPoolAttentionTailState( + replace(first.config, pool_size=3), + request_capacity=first.request_capacity, mtp_step=first.mtp_step, device=first.tail_buffer.device, + ) + first.tail_buffer.fill_(5) + second.tail_buffer.fill_(7) + tail_manager = AttResourceManager(request_capacity=3) + tail_manager.set_att_state("main", first) + tail_manager.set_att_state("main", second) + tail_manager.finish_state_creation() + tail_manager.reset_attention_tail_state([0, 1]) + assert not first.tail_buffer[:, :2].any() and not second.tail_buffer[:, :2].any() + assert (first.tail_buffer[:, 2] == 5).all() and (second.tail_buffer[:, 2] == 7).all() + manager = composition() + init_source(manager) + # Preserve accepted-state cursors as well as both algorithms' runtime and + # cached pages. A full reset here would silently corrupt a prefix restore. + for state in (manager.main_kda_state, manager.draft0_gdn_state): + state.mtp_state_index[1] = 2 + tensors = [ + *full_state_test_buffers(manager.main_full_att_fp32_state).values(), + *full_state_test_buffers(manager.draft0_mla_fp32_state).values(), + ] + for state in (manager.main_kda_state, manager.draft0_gdn_state): + tensors.extend(( + state.conv_state, state.ssm_state, state.mtp_state_index, + state.big_pages_pool.conv_state, state.big_pages_pool.ssm_state, + state.small_pages_pool.conv_state, state.small_pages_pool.ssm_state, + )) + preserved = [(tensor, tensor.contiguous().view(torch.uint8).clone()) for tensor in tensors] + manager.reset_attention_tail_state([1]) + assert not getattr(manager, f"main_{manager.att_tail_type}_state").tail_buffer[:, 1].any() + for tensor, snapshot in preserved: + assert torch.equal(tensor.contiguous().view(torch.uint8), snapshot) + + +@pytest.mark.parametrize("device", ["cuda"]) +def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | torch.device | int): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source, target = composition(device), composition(device) + slots = init_source(source) + layouts = layouts_for(source) + pages = torch.zeros( + (3, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=device == "cuda" + ) + source.write_req_to_cpu( + torch.arange(8, dtype=torch.long, device=device), + pages, + [2, 0], + layouts=layouts, + big_page_slot_ids=slots, + ) + if device == "cuda": + torch.cuda.synchronize() + tail = getattr(target, f"main_{target.att_tail_type}_state") + tail.tail_buffer.fill_(37) + assert tail.resource_id not in layouts + target_slots = target.bounded_att_big_page_allocator.alloc(2) + target.read_req_from_cpu( + pages, + [2, 0], + torch.arange(4, 12, dtype=torch.long, device=device), + layouts=layouts, + big_page_slot_ids=target_slots, + req_idx=2, + ) + for src, dst in zip( + (source.main_kda_state, source.draft0_gdn_state), (target.main_kda_state, target.draft0_gdn_state) + ): + for src_slot, dst_slot in zip(slots, target_slots): + assert torch.equal(src.big_pages_pool.conv_state[src_slot], dst.big_pages_pool.conv_state[dst_slot]) + assert torch.equal(src.big_pages_pool.ssm_state[src_slot], dst.big_pages_pool.ssm_state[dst_slot]) + assert (tail.tail_buffer[:, :2] == 37).all() + assert not tail.tail_buffer[:, 2].any() + for i, state in enumerate((target.main_kda_state, target.draft0_gdn_state)): + assert (state.conv_state[:, 2, :, :2] == 23 + i).all() + assert (state.ssm_state[:, 6] == 47 + i).all() + for src, dst in zip( + (source.main_full_att_fp32_state, source.draft0_mla_fp32_state), + (target.main_full_att_fp32_state, target.draft0_mla_fp32_state), + ): + assert torch.equal(read_kv_storage(src)[:, :8], read_kv_storage(dst)[:, 4:12]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("token_count", [2, 4, 6, 8]) +def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: int): + source, target = composition("cuda", with_tail=False), composition("cuda", with_tail=False) + source_slots = init_source(source) + layouts = layouts_for(source) + page_count = (token_count + 3) // 4 + page_ids = [2, 0][:page_count] + pages = torch.zeros((3, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) + small_slots = [] + if token_count % 4: + small_slots = source.bounded_att_small_page_allocator.alloc(1) + source.save_bounded_att_pages([page_count - 1], page_kind="small", slot_ids=small_slots) + torch.cuda.synchronize() + source.write_req_to_cpu( + torch.arange(token_count, device="cuda"), pages, page_ids, + layouts=layouts, big_page_slot_ids=source_slots[:token_count // 4], small_page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + + # 保留末槽给 LOAD；调用方只分配完整页，碎页不占普通页额度。 + target.bounded_att_big_page_allocator = StatePageAllocator(3, reserved_slots=1) + target.bounded_att_cpu_load_big_page_slot = 2 + allocator = target.bounded_att_big_page_allocator + slots = allocator.alloc(token_count // 4) + original_slots = slots.copy() + free_count = allocator.free_count + target.read_req_from_cpu( + pages, page_ids, torch.arange(4, 4 + token_count, device="cuda"), + layouts=layouts, big_page_slot_ids=slots, req_idx=2, + ) + torch.cuda.synchronize() + assert slots == original_slots + assert allocator.free_count == free_count + for src, dst in zip(source.get_bounded_attention_states(), target.get_bounded_attention_states()): + pool = dst.big_pages_pool + for src_slot, dst_slot in zip(source_slots, slots): + torch.testing.assert_close(pool.conv_state[dst_slot], src.big_pages_pool.conv_state[src_slot]) + torch.testing.assert_close(pool.ssm_state[dst_slot], src.big_pages_pool.ssm_state[src_slot]) + final_slot = source_slots[page_count - 1] + torch.testing.assert_close(dst.conv_state[:, 2, :, :2].cpu(), src.big_pages_pool.conv_state[final_slot]) + torch.testing.assert_close(dst.ssm_state[:, 6].cpu(), src.big_pages_pool.ssm_state[final_slot]) + assert not dst.conv_state[:, :2].any() + if token_count % 4: + torch.testing.assert_close(dst.big_pages_pool.conv_state[2], src.big_pages_pool.conv_state[final_slot]) + torch.testing.assert_close(dst.big_pages_pool.ssm_state[2], src.big_pages_pool.ssm_state[final_slot]) + for src, dst in zip(source.get_full_attention_states(), target.get_full_attention_states()): + torch.testing.assert_close(read_kv_storage(dst)[:, 4:4 + token_count], read_kv_storage(src)[:, :token_count]) + + +@pytest.mark.parametrize("tail_pool_size", [4, 8]) +def test_cpu_restore_resets_selected_tail_after_all_pages(tail_pool_size: int): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + source, target = composition(device="cuda", tail_pool_size=tail_pool_size), composition(device="cuda", tail_pool_size=tail_pool_size) + slots = init_source(source) + layouts = layouts_for(source) + pages = torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) + source.write_req_to_cpu( + torch.arange(8, dtype=torch.long, device="cuda"), + pages, + [0, 1], + layouts=layouts, + big_page_slot_ids=slots, + ) + torch.cuda.synchronize() + target_slots = target.bounded_att_big_page_allocator.alloc(2) + tail = getattr(target, f"main_{target.att_tail_type}_state") + tail.tail_buffer.fill_(37) + # 缓存流程保证最终前缀 L=8 是可恢复边界；中间页 L=4 不单独重置 tail。 + target.read_req_from_cpu( + pages, + [0, 1], + torch.arange(8, dtype=torch.long, device="cuda"), + layouts=layouts, + big_page_slot_ids=target_slots, + req_idx=0, + ) + assert not getattr(target, f"main_{target.att_tail_type}_state").tail_buffer[:, 0].any() + assert (getattr(target, f"main_{target.att_tail_type}_state").tail_buffer[:, 1:] == 37).all() + + +@pytest.mark.parametrize("device", ["cuda"]) +def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.device | int, pd_args: EasyDict): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source, target, dp_target = composition(device), composition(device), composition(device) + init_source(source) + for kind in ("kv", "att_state"): + layouts = layouts_for(source, kind=kind) + page = torch.empty(max(layout.byte_end for layout in layouts.values()), dtype=torch.uint8, device=device) + src_args = {"token_indices": torch.tensor([0, 2], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 1, "state_token_len": 8} + dst_args = {"token_indices": torch.tensor([5, 7], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 2, "state_token_len": 8} + source.write_to_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=layouts, **src_args) + target.read_from_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=layouts, **dst_args) + dp_tail = dp_target.get_attention_tail_states()[0] + dp_tail.tail_buffer.fill_(19) + dp_target.copy_from_other_dp( + remote_att_manager=source, + src_token_indexes=torch.tensor([0, 2], dtype=torch.int32, device=device), + dst_token_indexes=torch.tensor([5, 7], dtype=torch.long, device=device), + src_big_page_ids=[], dst_big_page_ids=[], + src_req_idx=1, dst_req_idx=2, + ) + for a, b in zip( + (target.main_full_att_fp32_state, target.draft0_mla_fp32_state), + (dp_target.main_full_att_fp32_state, dp_target.draft0_mla_fp32_state), + ): + assert torch.equal(read_kv_storage(a)[:, [5, 7]], read_kv_storage(b)[:, [5, 7]]) + for a, b in zip( + (target.main_kda_state, target.draft0_gdn_state), (dp_target.main_kda_state, dp_target.draft0_gdn_state) + ): + assert torch.equal(a.conv_state[:, 2], b.conv_state[:, 2]) + assert torch.equal(a.ssm_state[:, 6], b.ssm_state[:, 6]) + pd_tail = target.get_attention_tail_states()[0].tail_buffer[:, 2] + expected_tail = torch.zeros_like(pd_tail) + for position in range(4, 8): + expected_tail[:, position % pd_tail.shape[1]].fill_(9) + assert torch.equal(pd_tail, expected_tail) + assert not dp_tail.tail_buffer[:, 2].any() + assert (dp_tail.tail_buffer[:, :2] == 19).all() + assert (source.get_attention_tail_states()[0].tail_buffer[:, 1] == 9).all() + + +def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): + source, target = composition(), composition() + init_source(source) + tails = target.get_attention_tail_states() + for state in tails: + state.tail_buffer.fill_(19) + target.copy_from_other_dp( + remote_att_manager=source, + src_token_indexes=torch.empty(0, dtype=torch.long), + dst_token_indexes=torch.empty(0, dtype=torch.long), + src_big_page_ids=[], dst_big_page_ids=[], + src_req_idx=1, dst_req_idx=2, + ) + assert all(not state.tail_buffer[:, 2].any() for state in tails) + assert all((state.tail_buffer[:, :2] == 19).all() for state in tails) + + +@pytest.mark.parametrize("copy_tokens,copy_pages", [(True, True), (True, False), (False, True)]) +def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: bool, copy_pages: bool): + if copy_tokens and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + device = "cuda" if copy_tokens else "cpu" + source, target = composition(device), composition(device) + init_source(source) + for state in target.get_full_attention_states(): + for buffer in full_state_test_buffers(state).values(): + buffer.fill_(-1) + for index, state in enumerate(source.get_bounded_attention_states()): + for page in range(3): + state.big_pages_pool.conv_state[page].fill_(10 * index + page) + state.big_pages_pool.ssm_state[page].fill_(20 * index + page) + for state in target.get_bounded_attention_states(): + state.big_pages_pool.conv_state.fill_(-1) + state.big_pages_pool.ssm_state.fill_(-2) + state.small_pages_pool.conv_state.fill_(31) + state.small_pages_pool.ssm_state.fill_(32) + state.conv_state.fill_(41) + state.ssm_state.fill_(42) + state.mtp_state_index.fill_(1) + for state in target.get_attention_tail_states(): + state.tail_buffer.fill_(51) + + target.copy_from_other_dp( + remote_att_manager=source, + src_token_indexes=torch.tensor([0, 2] if copy_tokens else [], dtype=torch.long, device=device), + dst_token_indexes=torch.tensor([5, 7] if copy_tokens else [], dtype=torch.long, device=device), + src_big_page_ids=[2, 0] if copy_pages else [], + dst_big_page_ids=[0, 1] if copy_pages else [], + src_req_idx=1, dst_req_idx=2, + ) + + for state in target.get_full_attention_states(): + src_state = getattr(source, state.resource_id) + values = read_kv_storage(state) + if copy_tokens: + assert torch.equal(values[:, [5, 7]], read_kv_storage(src_state)[:, [0, 2]]) + else: + assert (values == -1).all() + assert (values[:, :5] == -1).all() + for state in target.get_bounded_attention_states(): + src_pool = getattr(source, state.resource_id).big_pages_pool + dst_pool = state.big_pages_pool + if copy_pages: + for src_page, dst_page in ((2, 0), (0, 1)): + assert torch.equal(dst_pool.conv_state[dst_page], src_pool.conv_state[src_page]) + assert torch.equal(dst_pool.ssm_state[dst_page], src_pool.ssm_state[src_page]) + else: + assert (dst_pool.conv_state == -1).all() and (dst_pool.ssm_state == -2).all() + assert (dst_pool.conv_state[2] == -1).all() and (dst_pool.ssm_state[2] == -2).all() + assert (state.small_pages_pool.conv_state == 31).all() + assert (state.small_pages_pool.ssm_state == 32).all() + source_state = getattr(source, state.resource_id) + assert torch.equal( + state.conv_state[:, 2, :, :state.config.conv_width], + source_state.conv_state[:, 1, :, :source_state.config.conv_width], + ) + dst_slot = 2 * (state.mtp_step + 1) + src_slot = source_state.mtp_step + 1 + assert torch.equal(state.ssm_state[:, dst_slot], source_state.ssm_state[:, src_slot]) + assert (state.conv_state[:, :2] == 41).all() + assert (state.conv_state[:, 2, :, state.config.conv_width:] == 41).all() + assert (state.ssm_state[:, :dst_slot] == 42).all() + assert (state.ssm_state[:, dst_slot + 1:] == 42).all() + assert state.mtp_state_index[2] == 0 + assert (state.mtp_state_index[:2] == 1).all() + for state in target.get_attention_tail_states(): + assert not state.tail_buffer[:, 2].any() + assert (state.tail_buffer[:, :2] == 51).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("readies", [[True, False, False], [False, True, False], [False, False, True], [True, True, True]]) +def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bool]): + source = composition("cuda", with_tail=False) + big_slots = init_source(source) + small_slots = source.bounded_att_small_page_allocator.alloc(1) + source.save_bounded_att_pages([1], page_kind="small", slot_ids=small_slots) + torch.cuda.synchronize() + layouts = layouts_for(source) + page_ids = [2, 0, 3] + token_ids = torch.arange(10, device="cuda") + baseline = torch.full((4, max(layout.byte_end for layout in layouts.values())), 139, dtype=torch.uint8, pin_memory=True) + pages = torch.full_like(baseline, 139, pin_memory=True) + source.write_req_to_cpu( + token_ids, baseline, page_ids, layouts=layouts, big_page_slot_ids=big_slots, small_page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + expected = pages.clone() + for page, ready in zip(page_ids, readies): + if not ready: + expected[page].copy_(baseline[page]) + # ready 页不能访问其页池快照；负的越界槽位可检测是否在构造视图前跳过。 + masked_big_slots = [-999 if readies[i] else slot for i, slot in enumerate(big_slots)] + masked_small_slots = [-999] if readies[2] else small_slots + source.write_req_to_cpu( + token_ids, pages, page_ids, layouts=layouts, cpu_page_readies=readies, + big_page_slot_ids=masked_big_slots, small_page_slot_ids=masked_small_slots, + ) + torch.cuda.synchronize() + assert torch.equal(pages, expected) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py new file mode 100644 index 0000000000..e109cbe0aa --- /dev/null +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -0,0 +1,116 @@ +"""Model-owned calibration slices must not depend on State initialization order.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from dataclasses import replace +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.full_att import FullAttentionState +from unit_tests.common.att_resources.model_helpers import owned_states +from lightllm.common.basemodel.basemodel_utils import init_model_full_attention +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.mark.parametrize("encoding", ["fp8kv_sph", "fp8kv_spt"]) +def test_calibration_slices_follow_model_configs_before_other_states_exist(tmp_path: Path, encoding: str): + config = dict( + architectures=["Qwen3NextForCausalLM"], + num_hidden_layers=8, + num_attention_heads=2, + num_key_value_heads=1, + hidden_size=32, + head_dim=16, + full_attention_interval=4, + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_conv_kernel_dim=3, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + scales = torch.arange(1, 9, dtype=torch.float32) * 0.125 + calibration = dict( + architectures="Qwen3NextForCausalLM", + num_head=1, + num_layers=8, + scales_shape=[8, 2], + scales=scales[:, None].expand(-1, 2).tolist(), + q_calibration=dict(num_head=1, scales_shape=[8, 1], scales=(scales[:, None] * 2).tolist()), + quant_type="per_head" if encoding == "fp8kv_sph" else "per_tensor", + ) + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text(json.dumps(calibration)) + args = StartArgs( + model_dir=str(tmp_path), + data_type="bfloat16", + llm_kv_type=encoding, + kv_quant_calibration_config_path=str(calibration_path), + page_size=4, + mtp_mode="vanilla_with_att", + ) + definitions = ( + (Qwen3NextTpPartModel, config, 8, args, (3, 7), (0, 2)), + ( + Qwen3NextTpPartModel, + dict(config, num_hidden_layers=6, full_attention_interval=2), + 6, + args, + (9, 11, 13), + (2, 5), + ), + (Qwen3NextTpPartModel, dict(config, num_hidden_layers=3), 3, args, (), (5, 5)), + ( + LlamaTpPartModel, + dict(config, num_hidden_layers=3), + 3, + replace(args, mtp_mode="vanilla_no_att"), + (), + (5, 5), + ), + # Native MTP keeps the main config but its resource hook describes one full-attention layer. + (Qwen3_5MTPModel, config, 1, args, (20,), (5, 6)), + (LlamaTpPartModel, dict(config, num_hidden_layers=2), 2, args, (21, 22), (6, 8)), + ) + models, layer_start = [], 0 + for cls, model_config, layer_count, model_args, _, _ in definitions: + model = cls.__new__(cls) + model.is_mtp_draft_model = bool(models) + model.config, model.args = model_config.copy(), model_args + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(layer_start, layer_start + layer_count)] + if models: + model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() + models.append(model) + layer_start += layer_count + + # Prepare KV definitions in reverse: no preceding model has a manager or State yet. + for index in reversed(range(len(models))): + model = models[index] + assert all(not hasattr(previous, "att_resource_manager") for previous in models[:index]) + model.att_resource_manager = AttResourceManager(request_capacity=1) + init_model_full_attention(model) + states = owned_states(model.att_resource_manager.__dict__, FullAttentionState) + state = states[0] if states else None + expected_layers, (start, end) = definitions[index][-2:] + if not expected_layers: + assert state is None + continue + assert state.config.layer_ids == expected_layers + assert read_kv_storage(state) is None + assert torch.equal(state.scales[:, 0], scales[start:end]) + if encoding == "fp8kv_sph": + assert torch.equal(state.q_scales[:, 0], scales[start:end] * 2) diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py new file mode 100644 index 0000000000..78a6bbe228 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states +from lightllm.common.att_resources.full_att import FullAttentionState +from unit_tests.common.att_resources.model_helpers import owned_states +import dataclasses + +import pytest +import torch + +from lightllm.common.req_manager import ReqManager +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import ( + get_added_mtp_kv_layer_num, + get_env_start_args, + get_unique_server_name, + set_env_start_args, + set_unique_server_name, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def clear_mtp_capacity_cache(): + get_added_mtp_kv_layer_num.cache_clear() + yield + get_added_mtp_kv_layer_num.cache_clear() + + +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("mtp", [0, 2]) +@pytest.mark.parametrize("mode", ["normal", "prefill"]) +def test_real_model_resource_hooks_and_shared_draft(monkeypatch: pytest.MonkeyPatch, glm: bool, mtp: int, mode: str): + args = StartArgs( + run_mode=mode, + data_type="bfloat16", + page_size=4, + pd_kv_page_size=128, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + linear_att_cache_size=2, + cpu_cache_token_page_size=8, + mtp_step=mtp, + mtp_mode="eagle_with_att" if mtp else None, + ) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) + set_env_start_args(dataclasses.asdict(args)) + get_env_start_args.cache_clear() + set_unique_server_name(args) + get_unique_server_name.cache_clear() + model_type = Glm5NextTpPartModel if glm else Qwen3NextTpPartModel + model = model_type.__new__(model_type) + model.args, model.data_type = args, torch.bfloat16 + model.tp_world_size_, model.tp_rank_ = 1, 0 + model.max_req_num, model.max_total_token_num = 3, None + model.max_seq_length, model.batch_max_tokens = 32, 16 + model.mem_fraction, model.num_kv_heads = 0.9, 2 + model.config = { + "num_attention_heads": 4, + "hidden_size": 64, + "num_key_value_heads": 2, + "head_dim": 16, + "n_layer": 4, + "num_hidden_layers": 4, + "full_attention_interval": 2, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_conv_kernel_dim": 3, + } + if glm: + model.config.update( + num_key_value_heads=1, + kv_lora_rank=512, + index_kpool=4, + index_head_dim=128, + layer_types=["linear_attention"] * 3 + ["deepseek_sparse_attention"], + linear_attn_config={"num_heads": 4, "head_dim": 8, "short_conv_kernel_size": 3, "kda_layers": [0, 1, 2]}, + ) + observed = [] + + def check_runtime_before_profile(device: str | torch.device | int): + resources = model.att_resource_manager + states = resources.get_full_attention_states() + assert main_linear_state(resources).conv_state.is_cuda + assert main_linear_state(resources).big_pages_pool is None + assert resources.full_att_allocator is None + assert not hasattr(resources, "pd_page_bytes") + assert not hasattr(resources, "pd_transfer_layouts") + assert read_kv_storage(states[0]) is None + observed.append(main_linear_state(resources).conv_state.data_ptr()) + reserve = 0 + if mode == "prefill": + reserve = resources.get_pd_page_bytes() * args.pd_kv_page_num + return resources.get_full_att_token_bytes() * (32 + args.page_size) + reserve, 0 + + monkeypatch.setattr(torch.cuda, "mem_get_info", check_runtime_before_profile) + model._init_req_manager() + prepare_model_states(model) + drafts = [] + if mtp: + draft_type = Glm5NextMTPModel if glm else Qwen3_5MTPModel + draft = draft_type.__new__(draft_type) + draft.main_model, draft.mtp_previous_draft_models = model, [] + draft.args, draft.data_type = args, torch.bfloat16 + draft.tp_rank_, draft.tp_world_size_ = 0, 1 + draft.max_req_num, draft.max_total_token_num = model.max_req_num, None + draft.config = dict(model.config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft._init_req_manager() + prepare_model_states(draft, layer_ids=(4,)) + assert draft.att_resource_manager is model.att_resource_manager + assert read_kv_storage(owned_states(draft.att_resource_manager.__dict__, FullAttentionState)[0]) is None + drafts.append(draft) + model.att_resource_manager.initialize_resources( + args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + model.req_manager.bind_hold_page(model.att_resource_manager.full_att_allocator.hold_token_indices) + resources = model.att_resource_manager + (full_state,) = resources.get_full_attention_states("main") + assert type(model.req_manager) is ReqManager + assert not hasattr(model.req_manager, "req_to_ssm_state") + assert resources.full_att_allocator.size == full_state.capacity == 32 + assert not hasattr(model, "mem_manager") + assert not hasattr(model, "full_att_state") + assert full_state.config.layer_ids == tuple( + layer for layer in model.linear_config.get_full_att_layer_ids() if layer < model.config["num_hidden_layers"] + ) + assert observed == [main_linear_state(resources).conv_state.data_ptr()] + assert resources.bounded_att_big_page_allocator.size == 6 + assert resources.bounded_att_small_page_allocator.size == 2 + assert not hasattr(model, "mtp_state_cursors") + assert (main_linear_state(resources).mtp_state_index is not None) == bool(mtp) + if mode == "prefill": + assert resources.get_pd_full_att_transfer_layouts() + assert resources.get_pd_bounded_att_or_tail_transfer_layouts() + if mtp: + draft = drafts[0] + assert draft.req_manager is model.req_manager + assert draft.att_resource_manager is resources + assert draft.att_resource_manager.full_att_allocator is resources.full_att_allocator + assert not hasattr(draft, "full_att_state") + (draft_state,) = resources.get_full_attention_states("draft0") + assert draft_state is not full_state + assert draft_state.capacity == full_state.capacity + assert set(draft_state.config.layer_ids).isdisjoint(full_state.config.layer_ids) + assert not hasattr(draft, "mem_manager") + assert not hasattr(draft, "mtp_state_cursors") + assert observed == [main_linear_state(resources).conv_state.data_ptr()] diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py new file mode 100644 index 0000000000..4781c32919 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state +import dataclasses +import json + +import pytest +import torch + +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles +from lightllm.common.att_resources.full_att import ( + GQAFullAttentionState, + FullAttentionConfig, +) +from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig +from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager +from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from easydict import EasyDict + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int = 0): + gk, gv = (4, 4) if glm else (2, 4) + local_k, local_v = max(1, gk // tp), max(1, gv // tp) + heads = 1 if glm else max(1, 2 // tp) + cls = Glm5NextResourceConfig if glm else Qwen3NextResourceConfig + config = cls( + tp_world_size=tp, + full_att_all_num_kv_heads=1 if glm else 2, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=heads, + full_att_head_dim=584 if glm else 16, + global_linear_k_heads=gk, + global_linear_v_heads=gv, + num_linear_k_heads=local_k, + num_linear_v_heads=local_v, + head_linear_k_dim=8, + head_linear_v_dim=8, + conv_kernel_size=3, + linear_layer_num=3 if glm else 2, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4 if glm else 2, + all_layer_num=4, + draft_full_att_kv_layer_num=draft_layers, + ) + manager_cls = Glm5NextAttResourceManager if glm else Qwen3NextAttResourceManager + resources = manager_cls(config, 3, mtp, tp_rank=rank) + if glm: + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + (3, *range(4, 4 + draft_layers)), 1, 584, torch.bfloat16, rank, tp, kv_dim=512 + ), + page_size=4, + ) + full.initialize_storage(16) + else: + full = GQAFullAttentionState( + FullAttentionConfig( + (1, 3, *range(4, 4 + draft_layers)), 2, 16, torch.bfloat16, rank, tp + ), + page_size=4, + ) + full.initialize_storage(16) + main, draft = create_full_attention_roles(full, 4) + resources.bind_storage( + main, + FullAttKVAllocator(16, 4), + draft_full_state=draft, + big_page_capacity=4, + small_page_capacity=2, + ) + return resources, resources.get_full_attention_states("main")[0] + + +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("tp", [1, 2, 4]) +@pytest.mark.parametrize("mtp", [0, 2]) +def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: int, mtp: int, cpu_args: EasyDict): + sources = [make_resources(glm, tp, rank, mtp) for rank in range(tp)] + targets = [make_resources(glm, tp, rank, mtp) for rank in range(tp)] + expected = [] + for rank, (resources, full) in enumerate(sources): + cfg = resources.cache_config + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.fill_(1 if glm else rank // max(1, tp // 2) + 1) + state = main_linear_state(resources) + local_start = 0 + for local_heads, global_heads in ( + (cfg.num_linear_k_heads, cfg.global_linear_k_heads), + (cfg.num_linear_k_heads, cfg.global_linear_k_heads), + (cfg.num_linear_v_heads, cfg.global_linear_v_heads), + ): + start = rank * local_heads // max(1, tp * local_heads // global_heads) + for h in range(local_heads): + state.conv_state[:, 0, local_start + h * 8 : local_start + (h + 1) * 8].fill_(start + h + 3) + local_start += local_heads * 8 + for h in range(cfg.num_linear_v_heads): + state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_linear_v_heads + h + 7) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + torch.cuda.synchronize() + expected.append((state.big_pages_pool.conv_state[slot].clone(), state.big_pages_pool.ssm_state[slot].clone())) + size = sources[0][0].get_cpu_page_bytes() + buffer = torch.zeros((2, size), dtype=torch.uint8, pin_memory=True) + for resources, full in sources: + resources.write_req_to_cpu( + torch.arange(8, dtype=torch.long, device="cuda"), + buffer, + [1], + layouts=resources.get_cpu_transfer_layouts(), + big_page_slot_ids=[0], + ) + torch.cuda.synchronize() + for rank, (resources, full) in enumerate(targets): + if glm: + resources.main_indexer_kpool_c4_tail_state.tail_buffer.fill_(13) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.read_req_from_cpu( + buffer, + [1], + torch.arange(8, 16, dtype=torch.long, device="cuda"), + layouts=resources.get_cpu_transfer_layouts(), + big_page_slot_ids=[slot], + req_idx=1, + ) + torch.testing.assert_close(main_linear_state(resources).conv_state[:, 1, :, :2].cpu(), expected[rank][0]) + torch.testing.assert_close(main_linear_state(resources).ssm_state[:, mtp + 1].cpu(), expected[rank][1]) + src_fields = full_state_test_fields(sources[rank][1], full_state_test_buffers(sources[rank][1])) + dst_fields = full_state_test_fields(full, full_state_test_buffers(full)) + for name in src_fields: + assert torch.equal( + src_fields[name].tensor[:, :8].contiguous().view(torch.uint8), + dst_fields[name].tensor[:, 8:16].contiguous().view(torch.uint8), + ) + if glm: + assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].any() + assert (resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, [0, 2]] == 13).all() + assert not buffer[0].any() + + +def test_cpu_layout_query_does_not_require_cpu_cache_enabled(monkeypatch: pytest.MonkeyPatch): + args = StartArgs(enable_cpu_cache=False, cpu_cache_token_page_size=8) # 未启用缓存时也可查询全局页配置。 + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + resources, _ = make_resources(True, 1, 0, 0) + assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_transfer_layouts().values()) + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py new file mode 100644 index 0000000000..609acafdb0 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -0,0 +1,165 @@ +"""Model-specific storage geometry must also describe external cache pages.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +import dataclasses +import json +from functools import partial + +from unit_tests.common.att_resources.model_helpers import prepare_model_states + +import pytest +import torch + +from lightllm.models.bloom.model import BloomTpPartModel +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen2.model import Qwen2TpPartModel +from lightllm.models.starcoder.model import StarcoderTpPartModel +from lightllm.models.starcoder2.model import Starcoder2TpPartModel +from lightllm.models.mistral.model import MistralTpPartModel +from lightllm.models.mixtral.model import MixtralTpPartModel +from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel +from lightllm.models.gemma3.model import Gemma3TpPartModel +from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.att_resource_config import build_cpu_cache_layouts +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +@pytest.fixture +def args(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + value = StartArgs( + model_dir=str(tmp_path), + data_type="float16", + tp=4, + dp=1, + page_size=4, + enable_cpu_cache=True, + cpu_cache_token_page_size=8, + run_mode="prefill", + pd_kv_page_size=8, + ) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(value))) + get_env_start_args.cache_clear() + yield value + get_env_start_args.cache_clear() + + +CASES = [ + (BloomTpPartModel, {}, 8, 16, 2), + (LlamaTpPartModel, {"num_key_value_heads": 4, "head_dim": 32}, 4, 32, 1), + (Qwen2TpPartModel, {"num_key_value_heads": 2}, 2, 16, 1), + (StarcoderTpPartModel, {}, 1, 16, 1), + (Starcoder2TpPartModel, {"num_key_value_heads": 4}, 4, 16, 1), + (MistralTpPartModel, {"num_key_value_heads": 4, "head_dim": 32}, 4, 32, 1), + (MixtralTpPartModel, {"num_key_value_heads": 4}, 4, 16, 1), + (Gemma_2bTpPartModel, {"num_key_value_heads": 2}, 2, 16, 2), + (Gemma3TpPartModel, {"num_key_value_heads": 4}, 4, 256, 1), + ( + Gemma4TpPartModel, + {"num_key_value_heads": 4, "head_dim": 256, "num_global_key_value_heads": None, "global_head_dim": 512}, + 8, + 256, + 2, + ), + ( + Gemma4TpPartModel, + {"num_key_value_heads": 16, "head_dim": 256, "num_global_key_value_heads": 4, "global_head_dim": 512}, + 16, + 256, + 4, + ), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("model_cls,overrides,global_heads,dim,local_heads", CASES) +def test_model_and_cpu_service_use_identical_storage_geometry( + args: StartArgs | EasyDict, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + model_cls: type[TpPartBaseModel], + overrides: dict[str, object], + global_heads: int, + dim: int, + local_heads: int, +): + # Exercise raw HF aliases as well as a wrapped multimodal config. + config = dict(model_type="llama", n_head=8, n_embd=128, n_layer=2, **overrides) + wrapped = {"model_type": "gemma4", "text_config": config} if model_cls is Gemma4TpPartModel else config + (tmp_path / "config.json").write_text(json.dumps(wrapped)) + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: model_cls) + service_layout = next(iter(build_cpu_cache_layouts().values())) + for rank in range(args.tp): + model = model_cls.__new__(model_cls) + model.config, model.args = wrapped, args + model.tp_rank_, model.tp_world_size_ = rank, args.tp + model.max_total_token_num, model.max_req_num = 19, 1 + prepare_model_states(model) + resources = model.att_resource_manager + resources.initialize_resources(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) + (state,) = resources.get_full_attention_states("main") + assert state.config.head_num == local_heads and state.config.head_dim == dim + assert state.k_storage.shape == state.v_storage.shape == (2, 20, local_heads, dim) + assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) + assert resources.full_att_allocator.size == 16 + layout = resources.get_cpu_transfer_layouts()[state.resource_id] + assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) + # 两层 GQA，K/V 各 global_heads 个 head，按实际 KV 精度计算全局传输页。 + expected_pd_bytes = 2 * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() + assert resources.get_pd_page_bytes() == expected_pd_bytes + fields = layout.regions + assert all(region.shape[2] == global_heads for region in fields) + assert not hasattr(model, "mem_manager") + + +@pytest.mark.parametrize("kind", KINDS) +def test_gpu_profile_counts_token_fields_hold_and_transfer_reserve( + monkeypatch: pytest.MonkeyPatch, + kind: str, + pd_args: EasyDict, +): + state = make_quant_state(kind, "cpu") + reference = AttResourceManager(request_capacity=1) + reference.set_att_state("main", state) + reference.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) + reference.finish_state_creation() + args = StartArgs(page_size=4, run_mode="prefill", pd_kv_page_size=8, pd_kv_page_num=3) + reserve = reference.get_pd_page_bytes() * args.pd_kv_page_num + required = state.get_token_bytes() * (32 + args.page_size) + # One byte below the full 32-token budget must drop a complete page. + total = 1 << 30 + free = required - 1 + reserve + total // 4 + observed_devices = [] + + def memory_info(device: str | torch.device | int): + observed_devices.append(device) + return free, total + + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.get_current_device_id", lambda: 3) + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + template = rebuild_full_state(state, config=state.config, token_capacity=0, device="meta") + resources = AttResourceManager(request_capacity=1) + resources.set_att_state("main", rebuild_full_state(template, config=template.config)) + resources.initialize_resources(args, token_capacity=None, mem_fraction=0.75, device="cpu") + assert resources.full_att_allocator.size == 28 + assert resources.get_pd_page_bytes() == reference.get_pd_page_bytes() + assert observed_devices == [torch.device("cuda", 3)] diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py new file mode 100644 index 0000000000..f72d7dc3cd --- /dev/null +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -0,0 +1,109 @@ +"""Verification commits State-owned recurrent data without model cursor aliases.""" + +from __future__ import annotations + + +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.bounded_att import GDNState +from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens +from unit_tests.common.att_resources.test_bounded_state import make_state + + +def make_manager(device: str | torch.device | int): + manager = AttResourceManager(request_capacity=3) + manager.set_att_state("main", make_state(device=device)) + manager.set_att_state("main", make_state(GDNState, device=device)) + manager.set_att_state("draft0", make_state(GDNState, device=device)) + manager.bounded_att_big_page_allocator = StatePageAllocator(3) + manager.bounded_att_small_page_allocator = StatePageAllocator(2) + manager.finish_state_creation() + return manager + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_commit_selects_state_for_one_role(): + manager = make_manager("cuda") + for state in manager.get_bounded_attention_states(): + for slot in range(3): + state.ssm_state[:, slot].fill_(slot + 10) + state.mtp_state_index.fill_(1) + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="main", + ) + for state in (manager.main_kda_state, manager.main_gdn_state): + assert state.mtp_state_index.tolist() == [2, 1, 1] + selected = state.mtp_state_index[0].item() + assert (state.ssm_state[:, selected] == 12).all() + assert manager.draft0_gdn_state.mtp_state_index.tolist() == [1, 1, 1] + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([1], dtype=torch.int32, device="cuda"), + role="draft0", + ) + assert manager.draft0_gdn_state.mtp_state_index.tolist() == [0, 1, 1] + assert manager.main_kda_state.mtp_state_index.tolist() == [2, 1, 1] + assert manager.main_gdn_state.mtp_state_index.tolist() == [2, 1, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_commit_without_recurrent_state_and_invalid_role(): + manager = AttResourceManager(request_capacity=3) + manager.finish_state_creation() + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="main", + ) + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="draft0", + ) + with pytest.raises(AssertionError, match="unknown model role"): + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="target", + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("widths", [(3, 3, 3), (3, 1, 2)]) +def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): + manager = make_manager("cuda") + backend = SimpleNamespace( + model=SimpleNamespace( + att_resource_manager=manager, + req_manager=SimpleNamespace( + req_sampling_params_manager=SimpleNamespace( + req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") + ) + ), + ) + ) + reqs, starts, tokens, expected_mask = [], [], [], [] + for req, width, accepted in zip((2, 0, 1), widths, (3, 1, 2)): + starts.append(len(reqs)) + reqs.extend([req] * width) + tokens.extend(99 if offset == accepted - 1 else 10 + offset for offset in range(width)) + expected_mask.extend(int(offset < accepted) for offset in range(width)) + accept_lengths, accepted_index = verify_mtp_tokens( + backend, + torch.tensor(tokens, device="cuda"), + torch.tensor(reqs, dtype=torch.int32, device="cuda"), + torch.tensor(starts, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), + ) + assert accept_lengths.tolist() == [3, 1, 2] + assert accepted_index.tolist() == expected_mask + assert manager.main_kda_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.main_gdn_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.draft0_gdn_state.mtp_state_index.tolist() == [0, 0, 0] diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py new file mode 100644 index 0000000000..4689a3d29e --- /dev/null +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -0,0 +1,278 @@ +"""Exercise production PD transporters between independent CUDA processes.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state + +import importlib.util +import os +import pickle +import time +import traceback + +import numpy as np +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.pd_io_struct import PDAgentMetadata, PDChunckedTransTask +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args +from unit_tests.common.att_resources.test_model_resources import make_resources + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from lightllm.common.att_resources.full_att.packed_mla_state import PackedMLAFullAttentionState + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter + from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter + from multiprocessing.connection import Connection + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +@pytest.mark.parametrize("backend", ["nccl", "nixl"]) +def test_pd_network_restores_packed_kv_accepted_linear_state_and_full_tail(backend: str): + if backend == "nixl" and importlib.util.find_spec("nixl") is None: + pytest.skip("NIXL unavailable") + context = mp.get_context("spawn") + pipes = [context.Pipe() for _ in range(2)] + processes = [ + context.Process(target=_worker, args=(child, backend, sender)) + for (_, child), sender in zip(pipes, (True, False)) + ] + try: + for process in processes: + process.start() + for _, child in pipes: + child.close() + metadata = [_receive(parent, "metadata") for parent, _ in pipes] + for (parent, _), peer in zip(pipes, reversed(metadata)): + parent.send(peer) + source, target = [_receive(parent, "result") for parent, _ in pipes] + assert source.keys() == target.keys() + for field in source: + np.testing.assert_array_equal(source[field], target[field], err_msg=field) + for parent, _ in pipes: + parent.send(None) + for process in processes: + process.join(timeout=30) + assert process.exitcode == 0 + finally: + for process in processes: + if process.is_alive(): + process.terminate() + if process.pid is not None: + process.join(timeout=10) + if process.is_alive(): + process.kill() + process.join(timeout=10) + for parent, child in pipes: + parent.close() + child.close() + + +def _notification( + transporter: NcclKVTransporter | NixlKVTransporter, + stage: str, + pending: list[PDChunckedTransTask], + key: str | None = None, +): + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + for messages in transporter.get_new_notifs().values(): + pending.extend(pickle.loads(message) for message in messages) + # Notifications are grouped by peer, so a later page's request can be + # returned before the current page's done. Match tasks as production does. + for index, task in enumerate(pending): + if task.error_info is not None: + raise AssertionError(task.error_info) + if task.write_stage == stage and (key is None or task.get_key() == key): + return pending.pop(index) + time.sleep(0.002) + raise TimeoutError(f"PD {stage} notification did not arrive") + + +def _task(peer: PDAgentMetadata, kind: str): + return PDChunckedTransTask( + request_id=123, + start_kv_index=0 if kind == "kv" else 5, + end_kv_index=5, + time_out_secs=60, + pd_master_node_id=0, + prefill_dp_index=0, + decode_dp_index=0, + src_device_id=0, + dst_device_id=1, + mem_indexes=[3, 0, 7, 1, 9] if kind == "kv" else [], + prefill_agent_name=None, + prefill_agent_metadata=None, + prefill_num_pages=None, + prefill_page_reg_desc=None, + decode_agent_name=peer.agent_name, + decode_agent_metadata=peer.agent_metadata, + decode_num_pages=peer.num_pages, + decode_page_reg_desc=peer.page_reg_desc, + first_gen_token_id=42, + first_gen_token_logprob=-0.5, + src_page_index=1 if kind == "kv" else 2, + page_kind=kind, + req_idx=1, + ) + + +def _snapshot( + resources: AttResourceManager, + full: PackedMLAFullAttentionState, + indices: Sequence[int] | torch.Tensor, + req_idx: int, + cursor: int, +): + result = {} + for name, field in full_state_test_fields(full, full_state_test_buffers(full)).items(): + if "token" in field.axes: + ids = torch.tensor(indices, device=field.tensor.device) + result[name] = field.tensor.index_select(field.axes.index("token"), ids).float().cpu().numpy() + state = main_linear_state(resources) + result["conv"] = state.conv_state[:, req_idx, ..., cursor : cursor + state.config.conv_width].float().cpu().numpy() + result["ssm"] = state.ssm_state[:, req_idx * 3 + cursor].cpu().numpy() + tail = resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, req_idx] + result["tail"] = torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() + return result + + +def _worker(connection: Connection, backend: str, sender: bool): + transporter = None + pending = [] + try: + device = 0 if sender else 1 + torch.cuda.set_device(device) + torch.set_num_threads(2) + set_env_start_args(StartArgs(run_mode="prefill" if sender else "decode", pd_kv_page_size=32)) + get_env_start_args.cache_clear() + resources, full = make_resources(True, 1, 0, 2) + buffer = allocate_transfer_buffer(resources, 3, device=f"cuda:{device}") + buffer.fill_(93 if sender else 11) + copy_stream = torch.cuda.Stream() + if backend == "nccl": + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter + + transporter = NcclKVTransporter(os.getpid(), device, buffer, host_ip="127.0.0.1") + else: + from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter + + transporter = NixlKVTransporter(os.getpid(), device, buffer) + connection.send( + ( + "metadata", + PDAgentMetadata( + transporter.agent_name, + transporter.agent_metadata, + transporter.num_pages, + transporter.local_page_mem_desc, + ), + ) + ) + assert connection.poll(60), "peer metadata did not arrive" + peer = connection.recv() + if sender: + for i, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + if "token" in field.axes: + values = torch.arange(field.tensor.numel(), device=device).reshape(field.tensor.shape) % 17 + i + 1 + field.tensor.copy_(values) + state = main_linear_state(resources) + for i, tensor in enumerate( + (state.conv_state, state.ssm_state, resources.main_indexer_kpool_c4_tail_state.tail_buffer) + ): + tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) + expected = _snapshot(resources, full, [3, 0, 7, 1, 9], 1, 0) + for kind in ("kv", "att_state"): + task = _task(peer, kind) + with torch.cuda.stream(copy_stream): + copy_page( + [resources], + buffer, + page_index=task.src_page_index, + dp_index=0, + tp_size=1, + token_indices=torch.tensor(task.mem_indexes, dtype=torch.long, device=device), + kind=kind, + req_idx=1, + state_token_len=task.end_kv_index, + write=True, + ) + # The production sender publishes a page only after its copy event completes. + copy_stream.synchronize() + transporter.send_write_request_task_to_decode_node(task) + ready = _notification(transporter, "ready", pending, task.get_key()) + assert ready.get_key() == task.get_key() + task.dst_page_index = ready.dst_page_index + task.xfer_handle = transporter.write_blocks_paged(task) + deadline = time.monotonic() + 60 + while transporter.check_task_status(task) == "PROC": + assert time.monotonic() < deadline, "PD transfer timed out" + time.sleep(0.002) + assert transporter.check_task_status(task) == "DONE" + transporter.send_write_done_task_to_decode_node(task) + transporter.release_xfer_handle(task.xfer_handle) + else: + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.zero_() + pool = main_linear_state(resources).big_pages_pool + pool.conv_state.fill_(41) + pool.ssm_state.fill_(43) + for kind in ("kv", "att_state"): + request = _notification(transporter, "request", pending) + assert request.page_kind == kind + request.dst_page_index = 2 if kind == "kv" else 1 + transporter.send_write_ready_task_to_prefill_node(request) + done = _notification(transporter, "done", pending, request.get_key()) + assert done.get_key() == request.get_key() + # Match production: unpack on a separate stream after the done notification. + with torch.cuda.stream(copy_stream): + copy_page( + [resources], + buffer, + page_index=request.dst_page_index, + dp_index=0, + tp_size=1, + token_indices=torch.tensor([8, 12, 10, 15, 11], dtype=torch.long, device=device) if kind == "kv" else torch.tensor([], dtype=torch.long, device=device), + kind=kind, + req_idx=2, + state_token_len=request.end_kv_index, + write=False, + ) + copy_stream.synchronize() + expected = _snapshot(resources, full, [8, 12, 10, 15, 11], 2, 0) + assert main_linear_state(resources).mtp_state_index[2] == 0 + assert not main_linear_state(resources).conv_state[:, :2].any() + assert not main_linear_state(resources).ssm_state[:, :6].any() + assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, :2].any() + assert not read_kv_storage(full)[:, [0, 1, 2, 3]].any() + assert (pool.conv_state == 41).all() + assert (pool.ssm_state == 43).all() + assert (buffer[0] == (93 if sender else 11)).all() + connection.send(("result", expected)) + assert connection.poll(60), "parent did not acknowledge completion" + connection.recv() + except BaseException: + connection.send(("error", traceback.format_exc())) + raise + finally: + if transporter is not None: + transporter.shutdown() + connection.close() + + +def _receive(connection: Connection, stage: str): + assert connection.poll(120), f"PD worker did not report {stage}" + kind, result = connection.recv() + assert kind == stage, result + return result diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py new file mode 100644 index 0000000000..f5c50fb4cc --- /dev/null +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.basemodel.test_overlap_utils import _make_decode_input, _make_prefill_input + + +def draft_model(mode: str): + resources = AttResourceManager(request_capacity=3) + resources.set_att_state("main", definition("gqa")) + draft = LlamaTpPartModel.__new__(LlamaTpPartModel) + draft.is_mtp_draft_model = True + draft.args = StartArgs(mtp_mode=mode, data_type="bfloat16") + draft.max_req_num = 2 + draft.layers_infer = [SimpleNamespace(layer_num_=2)] + draft.config = dict(num_attention_heads=2, num_key_value_heads=1, hidden_size=8, num_hidden_layers=1) + draft.tp_rank_, draft.tp_world_size_ = 0, 1 + main = LlamaTpPartModel.__new__(LlamaTpPartModel) + main.args, main.config = draft.args, dict(draft.config, num_hidden_layers=2) + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(2)] + main.att_resource_manager = resources + draft.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + draft.main_model, draft.mtp_previous_draft_models = main, [] + return draft, resources + + +@pytest.mark.parametrize("mode", ["vanilla_no_att", "eagle_no_att"]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.MonkeyPatch, mode: str, prefill: bool): + draft, resources = draft_model(mode) + draft.prepare_att_resources() + resources.finish_state_creation() + draft._init_att_backend() + draft._init_att_backend1() + assert draft.att_resource_manager is resources + assert not hasattr(draft, "full_att_state") + assert not resources.get_full_attention_states("draft0") + assert not any(name.startswith("draft0_") and name.endswith("_state") for name in vars(resources)) + + draft.infer_state_class = InferStateInfo + draft.return_all_prompt_logics = False + draft.hidden_collector_prototype = SimpleNamespace(new_instance=lambda: None) + draft.req_manager = object() + draft._select_mem_indexes = lambda value: torch.zeros(value.batch_size, dtype=torch.int32) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.dist_group_manager.get_group", lambda index: None) + model_input = _make_prefill_input([10, 11], req_idx=0, is_decode_req=False) if prefill else _make_decode_input() + state = draft._create_inferstate(model_input) + assert not hasattr(state, "full_att_state") + assert state.att_resource_manager is resources + assert state.is_draft_model + assert state.prefill_att_state is None and state.decode_att_state is None + assert state.req_manager is draft.req_manager + + +@pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) +def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: str): + draft, resources = draft_model(mode) + draft.prepare_att_resources() + assert draft.att_resource_manager is resources + resources.finish_state_creation() + (state,) = resources.get_full_attention_states("draft0") + assert not hasattr(draft, "full_att_state") + assert state.resource_id == "draft0_full_att_bf16_state" + assert state.config.layer_ids == (2,) + assert read_kv_storage(state) is None diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py new file mode 100644 index 0000000000..c95475996a --- /dev/null +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNPagesPool, KDAPagesPool +from lightllm.common.att_resources.bounded_att import utils + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.queues import Queue + + +@pytest.mark.parametrize("pool_cls", [KDAPagesPool, GDNPagesPool]) +@pytest.mark.parametrize("dp_fetch", [False, True]) +@pytest.mark.parametrize("pin_memory", [False, True]) +def test_page_pool_allocation_strategy( + cpu_args: EasyDict, + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], + dp_fetch: bool, + pin_memory: bool, +) -> None: + if pin_memory and not torch.cuda.is_available(): + pytest.skip("CUDA required for pinned memory") + cpu_args.enable_dp_prompt_cache_fetch = dp_fetch + config = ConvSSMStateConfig( + layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, + key_dim=4, value_dim=4, conv_width=3, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + ) + pool = pool_cls(size=2, config=config, pin_memory=pin_memory) + assert pool.conv_state.shape == (2, 2, 12, 3) + assert pool.ssm_state.shape == (2, 2, 1, 4, 4) + assert pool.conv_state.dtype == torch.bfloat16 + assert pool.ssm_state.dtype == torch.float32 + for tensor in (pool.conv_state, pool.ssm_state): + assert tensor.is_shared() == dp_fetch + assert tensor.is_pinned() == pin_memory + tensor.fill_(7) + assert torch.all(tensor == 7) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_shared_registration_lives_until_last_view( + cpu_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + unregistered = [] + unregister = utils._unregister_page_pool_storage + + def record_unregister(ptr: int, device_id: int) -> None: + unregistered.append(ptr) + unregister(ptr=ptr, device_id=device_id) + + monkeypatch.setattr(utils, "_unregister_page_pool_storage", record_unregister) + tensor = utils.create_page_pool_tensor(size=(4096,), dtype=torch.float32) + ptr = tensor.data_ptr() + view = tensor[1:] + del tensor + gc.collect() + assert not unregistered + assert view.is_pinned() + del view + gc.collect() + assert unregistered == [ptr] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_gpu_page_pool_keeps_device_allocation(cpu_args: EasyDict) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + tensor = utils.create_page_pool_tensor(size=(8,), dtype=torch.float32, device="cuda") + assert tensor.device.type == "cuda" + assert not tensor.is_pinned() + + +def test_empty_shared_pool_skips_cuda_registration( + cpu_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + + def unexpected_cuda_init() -> None: + pytest.fail("empty shared storage must not initialize CUDA") + + monkeypatch.setattr(torch.cuda, "current_device", unexpected_cuda_init) + tensor = utils.create_page_pool_tensor(size=(0, 4), dtype=torch.float32) + assert tensor.shape == (0, 4) + assert tensor.is_shared() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +@pytest.mark.parametrize("pool_cls", [KDAPagesPool, GDNPagesPool]) +def test_pinned_page_pool_survives_queue_on_another_gpu( + cpu_args: EasyDict, pool_cls: type[KDAPagesPool] | type[GDNPagesPool], +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + torch.cuda.set_device(0) + config = ConvSSMStateConfig( + layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, + key_dim=4, value_dim=4, conv_width=3, + conv_dtype=torch.float32, ssm_dtype=torch.float32, + ) + pool = pool_cls(size=2, config=config) + tensors = (pool.conv_state, pool.ssm_state) + pointers = [tensor.data_ptr() for tensor in tensors] + for tensor in tensors: + values = torch.arange(tensor.numel(), device="cuda", dtype=tensor.dtype).reshape(tensor.shape) + tensor.copy_(values, non_blocking=True) + torch.cuda.synchronize() + # 创建端也允许重复调用，不得重复注册或复制存储。 + pool.register_pin_memory() + ctx = mp.get_context("spawn") + queue = ctx.Queue() + result = ctx.Queue() + worker = ctx.Process(target=_read_page_pool_on_other_gpu, args=(queue, result)) + worker.start() + try: + queue.put(pool) + assert result.get(timeout=45) == [(True, False, True), (True, False, True)] + worker.join(timeout=15) + assert worker.exitcode == 0 + for tensor, ptr in zip(tensors, pointers): + assert tensor.data_ptr() == ptr + assert tensor.is_shared() and tensor.is_pinned() + assert tensor.flatten()[0].item() == 1234 + finally: + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + queue.close() + result.close() + + +def _read_page_pool_on_other_gpu(queue: Queue, result: Queue) -> None: + torch.cuda.set_device(1) + pool = queue.get(timeout=30) + tensors = (pool.conv_state, pool.ssm_state) + was_pinned = [tensor.is_pinned() for tensor in tensors] + pointers = [tensor.data_ptr() for tensor in tensors] + pool.register_pin_memory() + pool.register_pin_memory() + states = [] + for tensor, ptr, before in zip(tensors, pointers, was_pinned): + assert tensor.data_ptr() == ptr + values = tensor.to(device="cuda", non_blocking=True) + torch.cuda.synchronize() + expected = torch.arange(tensor.numel(), device="cuda", dtype=tensor.dtype).reshape(tensor.shape) + assert torch.equal(values, expected) + tensor.flatten()[0] = 1234 + states.append((tensor.is_shared(), before, tensor.is_pinned())) + result.put(states) diff --git a/unit_tests/common/att_resources/test_paged_storage.py b/unit_tests/common/att_resources/test_paged_storage.py new file mode 100644 index 0000000000..fae8ecf18a --- /dev/null +++ b/unit_tests/common/att_resources/test_paged_storage.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import pytest + +from lightllm.common.att_resources.allocators import FullAttKVAllocator + + +def test_kv_cache_allocation_requires_complete_pages(): + allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) + with pytest.raises(AssertionError, match="complete KV pages"): + allocator.alloc(5) + allocated = allocator.alloc(8) + assert len(allocated) == 8 + assert not set(allocator.hold_token_indices) & set(allocated.tolist()) + allocator.free(allocated) + assert allocator.can_use_mem_size == 16 + + +@pytest.mark.parametrize("size,page_size", [(5, 4), (-4, 4), (4, 0)]) +def test_kv_allocator_capacity_requires_complete_pages(size: int, page_size: int): + with pytest.raises(AssertionError, match="complete pages"): + FullAttKVAllocator(size, page_size=page_size, pin_memory=False) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py new file mode 100644 index 0000000000..40433b2820 --- /dev/null +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +from dataclasses import replace +from math import prod + +import pytest +import torch + +from lightllm.common.att_resources.full_att import ( + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, + PackedMLAAttentionConfig, + PackedMLAFullAttentionState, +) +from lightllm.common.att_resources.types import ( + AxisSlice, StateTransferLayout, TransferRegion, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.att_resources.full_att.base_state import FullAttentionState + + +KINDS = ("int4", "int8", "fp8_tensor", "fp8_head", "packed", "packed_fp8") + + +def make_quant_state(kind: str, device: str | torch.device | int): + if kind in ("int4", "int8"): + config = GroupQuantizedAttentionConfig((1, 3), 2, 16, torch.float16, bits=int(kind[-1])) + prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) + prepared_full_state.initialize_storage(8) + return prepared_full_state + if kind.startswith("fp8_"): + per_head = kind == "fp8_head" + config = StaticFP8AttentionConfig( + (1, 3), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor", calibration_id="test" + ) + scales = torch.arange(1, 9 if per_head else 5, dtype=torch.float32, device=device).reshape(2, -1) / 16 + prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device, scales=scales) + prepared_full_state.initialize_storage(8) + return prepared_full_state + fp8 = kind == "packed_fp8" + config = PackedMLAAttentionConfig( + (1, 3), 1, 800 if fp8 else 648, torch.uint8 if fp8 else torch.bfloat16, fp8_nope_dim=512 if fp8 else 0 + ) + prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) + prepared_full_state.initialize_storage(8) + return prepared_full_state + + +def transfer_layout(state: FullAttentionState): + offset, regions = 0, [] + for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): + shape = tuple(4 if axis == "token" else dim for axis, dim in zip(field.axes, field.tensor.shape)) + itemsize = field.tensor.element_size() + offset = (offset + itemsize - 1) // itemsize * itemsize + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=offset + prod(shape) * itemsize, + dtype=field.tensor.dtype, + axes=field.axes, + shape=shape, + axis_slices=tuple( + AxisSlice(axis=axis, src_start=0, dst_start=0, length=dim) + for axis, dim in zip(field.axes, shape) if axis not in ("layer", "token") + ), + ) + ) + offset += prod(shape) * itemsize + return StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=0, + byte_end=(offset + 15) // 16 * 16, + regions=tuple(regions), + ) + + +def assert_token_fields_equal( + source: FullAttentionState, + source_ids: Sequence[int] | torch.Tensor, + target: FullAttentionState, + target_ids: Sequence[int] | torch.Tensor, +): + src = full_state_test_fields(source, full_state_test_buffers(source)) + dst = full_state_test_fields(target, full_state_test_buffers(target)) + for name in src: + assert torch.equal( + src[name].tensor[:, source_ids].contiguous().view(torch.uint8), + dst[name].tensor[:, target_ids].contiguous().view(torch.uint8).to(src[name].tensor.device), + ), name + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( + kind: str, device: str | torch.device | int, monkeypatch: pytest.MonkeyPatch, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source, target = make_quant_state(kind, device), make_quant_state(kind, device) + for i, field in enumerate(full_state_test_fields(source, full_state_test_buffers(source)).values()): + field.tensor.fill_(i + 1) + field.tensor[:, 0].fill_(i + 11) + if device == "cuda": + source.copy_full_att_tokens(torch.tensor([0, 2], dtype=torch.long, device=device), torch.tensor([4, 5], dtype=torch.long, device=device)) + assert_token_fields_equal(source, [0, 2], source, [4, 5]) + layout = transfer_layout(source) + cpu = torch.zeros((2, layout.byte_end), dtype=torch.uint8, pin_memory=device == "cuda") + if device == "cpu": + with pytest.raises(AssertionError, match="CUDA State storage"): + source.write_to_cpu(torch.tensor([0, 1, 2, 3, 4], device=device), cpu, [1, 0], layout=layout) + else: + source.write_to_cpu(torch.tensor([0, 1, 2, 3, 4], dtype=torch.long, device=device), cpu, [1, 0], layout=layout) + if device == "cuda": + torch.cuda.synchronize() + target.read_from_cpu(cpu, [1, 0], torch.tensor([1, 2, 3, 4, 5], dtype=torch.long, device=device), layout=layout) + assert_token_fields_equal(source, [0, 1, 2, 3, 4], target, [1, 2, 3, 4, 5]) + page = torch.empty(layout.byte_end, dtype=torch.uint8, device=device) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + source.write_to_pd_buffer(torch.tensor([0, 3], dtype=torch.long, device=device), page.unsqueeze(0), pd_page_index=0, layout=layout) + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + target.read_from_pd_buffer(page.unsqueeze(0), torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, layout=layout) + else: + source.write_to_pd_buffer(torch.tensor([0, 3], dtype=torch.long, device=device), page.unsqueeze(0), pd_page_index=0, layout=layout) + target.read_from_pd_buffer(page.unsqueeze(0), torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, layout=layout) + assert_token_fields_equal(source, [0, 3], target, [6, 7]) + src_indices = torch.tensor([0, 1], dtype=torch.int32, device=device) + dst_indices = torch.tensor([0, 2], dtype=torch.int64, device=device) + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("DP token indices must not be converted to Python lists") + + tensor_to = torch.Tensor.to + + def reject_index_conversion(tensor: torch.Tensor, *args, **kwargs): + if tensor is src_indices or tensor is dst_indices: + pytest.fail("DP State must use the supplied CUDA indices directly") + return tensor_to(tensor, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + patch.setattr(torch.Tensor, "to", reject_index_conversion) + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + target.copy_tokens_from_other_dp( + src_token_indices=src_indices, dst_token_indices=dst_indices, remote_dp_full_att_state=source, + ) + return + target.copy_tokens_from_other_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, + ) + assert_token_fields_equal(source, [0, 1], target, [0, 2]) + + +@pytest.mark.parametrize("kind", KINDS) +def test_cuda_compute_write_preserves_unrelated_fields(kind: str): + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = make_quant_state(kind, "cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + if isinstance(state, PackedMLAFullAttentionState): + state.get_layer_view(1).indexer_k.fill_(29) + dims = state.config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else state.config.head_dim + heads = 1 if isinstance(state, PackedMLAFullAttentionState) else 2 * state.config.head_num + kv = torch.randn((2, heads, dims), dtype=torch.bfloat16, device="cuda") + ids = torch.tensor([1, 6], dtype=torch.int32, device="cuda") + state.write_full_att_kv(1, ids, kv) + if isinstance(state, GroupQuantizedFullAttentionState): + quant = read_kv_storage(state)[0, ids].reshape(2, heads, -1, state.config.group_size * state.config.bits // 8) + if state.config.bits == 4: + raw = quant.to(torch.uint8) + quant = torch.stack(((raw & 15).to(torch.int16) - 7, (raw >> 4).to(torch.int16) - 7), -1).flatten(-2) + dequant = (quant.float() * torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids, ..., None]).reshape_as(kv) + assert torch.all((dequant - kv.float()).abs() <= torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids].amax() * 2) + elif isinstance(state, StaticFP8FullAttentionState): + scales = state.scales[0] + if state.config.scale_mode == "per_tensor": + scales = scales.repeat_interleave(state.config.head_num) + dequant = read_kv_storage(state)[0, ids].view(torch.float8_e4m3fn).float() * scales[None, :, None] + torch.testing.assert_close(dequant, kv.float(), rtol=0.13, atol=0.03) + else: + view = state.get_layer_view(1) + assert (view.indexer_k == 29).all() + if not state.config.fp8_nope_dim: + assert torch.equal(view.kv[ids], kv) + else: + fields = full_state_test_fields(state) + assert torch.isfinite(fields["kv_scale"].tensor[0, ids]).all() + assert torch.equal(fields["kv_rope"].tensor[0, ids], kv[..., 512:]) + assert not read_kv_storage(state)[1].any() + + +def test_fp8_calibration_mismatch_fails(): + src = make_quant_state("fp8_head", "cpu") + dst = rebuild_full_state(src, config=replace(src.config, calibration_id="other")) + with pytest.raises(AssertionError, match="calibration"): + dst.copy_tokens_from_other_dp( + src_token_indices=torch.tensor([0], dtype=torch.long), + dst_token_indices=torch.tensor([1], dtype=torch.long), + remote_dp_full_att_state=src, + ) + + +@pytest.mark.parametrize("kind", ["fp8_tensor", "fp8_head"]) +def test_deferred_storage_preserves_real_fp8_calibration(kind: str): + source = make_quant_state(kind, "cpu") + state = StaticFP8FullAttentionState( + source.config, page_size=source.page_size, device="meta", + scales=source.scales, q_scales=source.q_scales, + ) + assert read_kv_storage(state) is None and not state.scales.is_meta + assert state.scales is source.scales + state.device = torch.device("cpu") + state.initialize_storage(source.capacity) + assert torch.equal(state.scales, source.scales) + assert torch.equal(torch.from_numpy(state.cpu_scales), source.scales) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py new file mode 100644 index 0000000000..96c3cc67d2 --- /dev/null +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -0,0 +1,30 @@ +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +import torch + +from lightllm.common.att_resources.full_att import ( + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, +) + + +def test_quantized_kv_buffers_include_the_full_hold_page(): + allocations = [] + for bits in (8, 4): + state = GroupQuantizedFullAttentionState( + GroupQuantizedAttentionConfig((0, 1, 2), 2, 16, torch.float16, bits=bits), + page_size=4, device="cpu", + ) + state.initialize_storage(32) + allocations.extend(tuple(buffer.shape) for buffer in full_state_test_buffers(state).values()) + + assert allocations == [ + (3, 36, 2, 16), + (3, 36, 2, 16), + (3, 36, 2, 2), + (3, 36, 2, 2), + (3, 36, 2, 8), + (3, 36, 2, 8), + (3, 36, 2, 2), + (3, 36, 2, 2), + ] diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py new file mode 100644 index 0000000000..5d40e8dbca --- /dev/null +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import main_linear_state +from unit_tests.common.att_resources.test_model_resources import make_resources + +if TYPE_CHECKING: + from easydict import EasyDict + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("glm", [False, True]) +def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( + glm: bool, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +): + pd_args.pd_kv_page_size = 128 + source, _ = make_resources(glm, 1, 0, 2, draft_layers=2) + target, _ = make_resources(glm, 1, 0, 2, draft_layers=2) + for index, role in enumerate(("main", "draft0")): + state = getattr(source, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + field.tensor.fill_(11 + index) + if glm: + getattr(source, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 0].fill_(7 + index) + getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 2].fill_(19) + main_linear_state(source).conv_state.fill_(3) + main_linear_state(source).ssm_state.fill_(5) + page = torch.empty(source.get_pd_page_bytes(), dtype=torch.uint8, device="cuda") + for kind, getter in ( + ("kv", "get_pd_full_att_transfer_layouts"), + ("att_state", "get_pd_bounded_att_or_tail_transfer_layouts"), + ): + source.write_to_pd_buffer( + page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=getattr(source, getter)(), + **({"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 0, "state_token_len": 8}), + ) + target.read_from_pd_buffer( + page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=getattr(target, getter)(), + **({"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 1, "state_token_len": 8}), + ) + src_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + dst_indices = torch.tensor([10, 11], dtype=torch.int64, device="cuda") + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("DP manager and States must preserve Tensor indices") + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + target.copy_from_other_dp( + remote_att_manager=source, + src_token_indexes=src_indices, + dst_token_indexes=dst_indices, + src_big_page_ids=[], dst_big_page_ids=[], + src_req_idx=0, dst_req_idx=2, + ) + for index, role in enumerate(("main", "draft0")): + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + assert (field.tensor[:, 8:12] == 11 + index).all() + if glm: + tail = getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer + expected_tail = torch.zeros_like(tail[:, 1]) + for position in range(4, 8): + expected_tail[:, position % tail.shape[2]].fill_(7 + index) + assert torch.equal(tail[:, 1], expected_tail) # PD restores only the prefill logical window. + assert not tail[:, 2].any() # DP fetch resets the aligned prefix. + + slot = source.bounded_att_big_page_allocator.alloc(1)[0] + source.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + torch.cuda.synchronize() + cpu = torch.empty((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.write_req_to_cpu( + torch.arange(8, dtype=torch.long, device="cuda"), + cpu, + [0], + layouts=source.get_cpu_transfer_layouts(), + big_page_slot_ids=[slot], + ) + torch.cuda.synchronize() + dst_slot = target.bounded_att_big_page_allocator.alloc(1)[0] + target.read_req_from_cpu( + cpu, + [0], + torch.arange(8, 16, dtype=torch.long, device="cuda"), + layouts=target.get_cpu_transfer_layouts(), + big_page_slot_ids=[dst_slot], + req_idx=1, + ) + for index, role in enumerate(("main", "draft0")): + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + assert (field.tensor[:, 8:16] == 11 + index).all() + if glm: + assert not getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 1].any() + assert not getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 2].any() diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py new file mode 100644 index 0000000000..5cdbea0018 --- /dev/null +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state +from types import SimpleNamespace +import traceback +import uuid + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.server.router.model_infer import shared_att_resources_utils +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from unit_tests.common.att_resources.test_model_resources import make_resources +from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state, assert_token_fields_equal + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cpu_pd_and_dp_model_layouts_cover_quantized_fields( + kind: str, device: str, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + src, dst = make_quant_state(kind, device), make_quant_state(kind, device) + managers = [] + for state in (src, dst): + manager = AttResourceManager(request_capacity=2) + manager.set_att_state("main", state) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_creation() + managers.append(manager) + src, dst = [manager.get_full_attention_states("main")[0] for manager in managers] + for i, field in enumerate(full_state_test_fields(src, full_state_test_buffers(src)).values()): + if "token" in field.axes: + field.tensor.fill_(i + 3) + source_indices = torch.tensor([1, 3], dtype=torch.int32, device=device) + pd_indices = torch.tensor([4, 6], dtype=torch.int64, device=device) + cpu_indices = torch.tensor([5, 7], dtype=torch.int64, device=device) + dp_indices = torch.tensor([0, 2], dtype=torch.int64, device=device) + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("transfer paths must not convert Tensor indices to Python lists") + + monkeypatch.setattr(torch.Tensor, "tolist", reject_host_list) + page = allocate_transfer_buffer(managers[0], 1, device=device) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + copy_page([managers[0]], page, page_index=0, dp_index=0, tp_size=1, token_indices=source_indices, write=True) + else: + copy_page([managers[0]], page, page_index=0, dp_index=0, tp_size=1, token_indices=source_indices, write=True) + copy_page([managers[1]], page, page_index=0, dp_index=0, tp_size=1, token_indices=pd_indices, write=False) + assert_token_fields_equal(src, [1, 3], dst, [4, 6]) + + cpu_pages = torch.empty((1, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") + if device == "cpu": + with pytest.raises(AssertionError, match="CUDA State storage"): + managers[0].write_req_to_cpu(source_indices, cpu_pages, [0], layouts=managers[0].get_cpu_transfer_layouts()) + else: + managers[0].write_req_to_cpu(source_indices, cpu_pages, [0], layouts=managers[0].get_cpu_transfer_layouts()) + if device == "cuda": + torch.cuda.synchronize() + managers[1].read_req_from_cpu(cpu_pages, [0], cpu_indices, layouts=managers[1].get_cpu_transfer_layouts()) + assert_token_fields_equal(src, [1, 3], dst, [5, 7]) + + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + managers[1].copy_from_other_dp( + remote_att_manager=managers[0], src_token_indexes=source_indices, dst_token_indexes=dp_indices, + src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, dst_req_idx=0, + ) + return + managers[1].copy_from_other_dp( + remote_att_manager=managers[0], + src_token_indexes=source_indices, + dst_token_indexes=dp_indices, + src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, dst_req_idx=0, + ) + assert_token_fields_equal(src, [1, 3], dst, [0, 2]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("source_device", [0, 1]) +@pytest.mark.parametrize("copy_big_pages", [False, True]) +def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( + monkeypatch: pytest.MonkeyPatch, + glm: bool, + source_device: int, + copy_big_pages: bool, +): + if source_device >= torch.cuda.device_count(): + pytest.skip("two CUDA devices required") + from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans as module + + with torch.cuda.device(source_device): + source, source_full = make_resources(glm, 1, 0, 2) + with torch.cuda.device(0): + target, target_full = make_resources(glm, 1, 0, 2) + for i, field in enumerate(full_state_test_fields(source_full, full_state_test_buffers(source_full)).values()): + if "token" in field.axes: + field.tensor.fill_(i + 5) + main_linear_state(source).conv_state.fill_(-1) + main_linear_state(source).conv_state[:, 1, ..., :2].fill_(13) + main_linear_state(source).ssm_state[:, 3].fill_(17) + if glm: + source.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].fill_(23) + target.main_indexer_kpool_c4_tail_state.tail_buffer[:, 2].fill_(29) + torch.cuda.synchronize(source_device) + # 同时验证运行态与指定大页复制，保留其他页及小页快照。 + for pool in (main_linear_state(target).big_pages_pool, main_linear_state(target).small_pages_pool): + pool.conv_state.fill_(41) + pool.ssm_state.fill_(43) + src_big_page_ids, dst_big_page_ids = [], [] + if copy_big_pages: + src_big_page_ids = source.bounded_att_big_page_allocator.alloc(1) + dst_big_page_ids = target.bounded_att_big_page_allocator.alloc(1) + source_pool = main_linear_state(source).big_pages_pool + source_pool.conv_state[src_big_page_ids[0]].fill_(51) + source_pool.ssm_state[src_big_page_ids[0]].fill_(53) + req = SimpleNamespace(req_idx=2, cur_kv_len=4, hold_kv_len=8, hybrid_len_to_big_page_id={}) + task = module.TransTask( + dst_req=req, src_dp_rank=1, src_req_idx=1, + src_token_indexes=torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=f"cuda:{source_device}"), + dst_token_indexes=torch.tensor([8, 9, 10, 11]), + src_big_page_ids=src_big_page_ids, dst_big_page_ids=dst_big_page_ids, + ) + syncs = [] + + def barrier(group: str): + torch.cuda.synchronize(0) + torch.cuda.synchronize(source_device) + assert req.cur_kv_len == 4 + assert not req.hybrid_len_to_big_page_id + syncs.append(group) + + monkeypatch.setattr(module.dist, "barrier", barrier) + worker = module.DPKVSharedMoudle.__new__(module.DPKVSharedMoudle) + worker.backend = SimpleNamespace( + model=SimpleNamespace(att_resource_manager=target), + dp_att_resource_managers={1: source}, + is_hybrid_att_model=True, + is_master_in_dp=False, + node_nccl_group="node", + ) + worker.big_page_tokens = 4 + worker.kv_trans([task]) + assert syncs == ["node"] and req.cur_kv_len == 8 + assert_token_fields_equal(source_full, [0, 1, 2, 3], target_full, [8, 9, 10, 11]) + assert (main_linear_state(target).conv_state[:, 2, ..., :2] == 13).all() + assert (main_linear_state(target).ssm_state[:, 6] == 17).all() + assert main_linear_state(target).mtp_state_index[2] == 0 + assert not main_linear_state(target).conv_state[:, :2].any() + big_pool = main_linear_state(target).big_pages_pool + if copy_big_pages: + page_id = dst_big_page_ids[0] + assert (big_pool.conv_state[page_id] == 51).all() + assert (big_pool.ssm_state[page_id] == 53).all() + assert req.hybrid_len_to_big_page_id == {8: page_id} + else: + assert not req.hybrid_len_to_big_page_id + untouched = [page for page in range(big_pool.size) if page not in dst_big_page_ids] + assert (big_pool.conv_state[untouched] == 41).all() + assert (big_pool.ssm_state[untouched] == 43).all() + assert (main_linear_state(target).small_pages_pool.conv_state == 41).all() + assert (main_linear_state(target).small_pages_pool.ssm_state == 43).all() + if glm: + assert not target.main_indexer_kpool_c4_tail_state.tail_buffer[:, 2].any() + assert (source.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1] == 23).all() + assert not target.main_indexer_kpool_c4_tail_state.tail_buffer[:, :2].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("p2p", [False, True]) +@pytest.mark.parametrize("share_pools,ipc_kinds", [ + (False, ("pd",)), + (True, ("pd", "dp")), + (True, ("dp", "pd")), +]) +def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( + monkeypatch: pytest.MonkeyPatch, p2p: bool, share_pools: bool, + ipc_kinds: tuple[str, ...], cpu_args: EasyDict, +): + cpu_args.enable_dp_prompt_cache_fetch = share_pools + from lightllm.utils.envs_utils import get_unique_server_name + from lightllm.utils.shm_utils import create_or_link_shm + + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "att-ipc-" + uuid.uuid4().hex) + get_unique_server_name.cache_clear() + monkeypatch.setattr(shared_att_resources_utils, "get_current_rank_in_node", lambda: 0) + if p2p and torch.cuda.device_count() < 2: + pytest.skip("two CUDA devices required for IPC peer mapping") + monkeypatch.setattr(shared_att_resources_utils, "kv_trans_use_p2p", lambda: p2p) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", mp.reductions.reduce_tensor.__code__) + source_device = int(p2p) + with torch.cuda.device(source_device): + resources, full = make_resources(True, 1, 0, 2) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.fill_(9) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, 4].zero_() + main_linear_state(resources).conv_state.fill_(11) + resources.main_indexer_kpool_c4_tail_state.tail_buffer.fill_(7) + for state in resources.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + pool.conv_state.zero_() + pool.ssm_state.zero_() + slot = resources.bounded_att_big_page_allocator.alloc(1) + resources.save_bounded_att_pages([0], page_kind="big", slot_ids=slot) + mapping = torch.zeros((3, 16), dtype=torch.int32, device=f"cuda:{source_device}") + mapping[1, 0] = 7 + torch.cuda.synchronize(source_device) + for kind in ipc_kinds: + shared_att_resources_utils.publish_att_resources(resources, mapping, kind=kind, reader_count=1) + ctx = mp.get_context("spawn") + parent, child = ctx.Pipe() + process = ctx.Process(target=_ipc_worker, args=(child, ipc_kinds)) + process.start() + child.close() + try: + assert parent.poll(45), "CUDA IPC child did not finish" + error = parent.recv() + assert error is None, error + assert (read_kv_storage(full)[:, 4] == 9).all() + assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].any() + assert (main_linear_state(resources).conv_state == 11).all() + assert (main_linear_state(resources).big_pages_pool.conv_state[slot[0]] == 11).all() + assert not main_linear_state(resources).big_pages_pool.ssm_state[slot[0]].any() + if share_pools: + assert main_linear_state(resources).small_pages_pool.conv_state.flatten()[0].item() == 37 + assert resources.bounded_att_big_page_allocator.free_count == 1 + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + for kind in ipc_kinds: + shm = create_or_link_shm(f"rank_storage_{kind}_0", expected_size=-1, force_mode="link") + shm.close() + shm.unlink() + get_unique_server_name.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("source_tp,target_tp", [(1, 4), (4, 1), (2, 8), (8, 2)]) +def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, target_tp: int, pd_args: EasyDict): + pd_args.pd_kv_page_size = 128 + source = [make_resources(False, source_tp, rank, 2) for rank in range(source_tp)] + target = [make_resources(False, target_tp, rank, 2) for rank in range(target_tp)] + for rank, (manager, full) in enumerate(source): + for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + for local_head in range(full.config.head_num): + head = rank * full.config.head_num // max(1, source_tp // 2) + local_head + field.tensor[:, :, local_head].fill_(100 * field_index + head) + cfg = main_linear_state(manager).config + offset = 0 + for group, (local_heads, global_heads) in enumerate( + ((cfg.num_k_heads, 2), (cfg.num_k_heads, 2), (cfg.num_v_heads, 4)) + ): + start = rank * local_heads // max(1, source_tp // global_heads) + for head in range(local_heads): + main_linear_state(manager).conv_state[:, 1, offset + head * 8 : offset + (head + 1) * 8, :2].fill_( + 10 * group + start + head + ) + offset += local_heads * 8 + start = rank * cfg.num_v_heads // max(1, source_tp // 4) + for head in range(cfg.num_v_heads): + main_linear_state(manager).ssm_state[:, 3, head].fill_(30 + start + head) + page = allocate_transfer_buffer(source[0][0], 1, device="cuda") + for kind in ("kv", "att_state"): + copy_page( + [m for m, _ in source], + page, + page_index=0, + dp_index=0, + tp_size=source_tp, + token_indices=torch.tensor([0, 3], dtype=torch.long, device="cuda"), + kind=kind, + req_idx=1, + write=True, + ) + copy_page( + [m for m, _ in target], + page, + page_index=0, + dp_index=0, + tp_size=target_tp, + token_indices=torch.tensor([4, 7], dtype=torch.long, device="cuda"), + kind=kind, + req_idx=2, + write=False, + ) + for rank, (manager, full) in enumerate(target): + for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + for local_head in range(full.config.head_num): + head = rank * full.config.head_num // max(1, target_tp // 2) + local_head + assert (field.tensor[:, [4, 7], local_head] == 100 * field_index + head).all() + cfg = main_linear_state(manager).config + offset = 0 + for group, (local_heads, global_heads) in enumerate( + ((cfg.num_k_heads, 2), (cfg.num_k_heads, 2), (cfg.num_v_heads, 4)) + ): + start = rank * local_heads // max(1, target_tp // global_heads) + for head in range(local_heads): + actual = main_linear_state(manager).conv_state[:, 2, offset + head * 8 : offset + (head + 1) * 8, :2] + assert (actual == 10 * group + start + head).all() + offset += local_heads * 8 + start = rank * cfg.num_v_heads // max(1, target_tp // 4) + for head in range(cfg.num_v_heads): + assert (main_linear_state(manager).ssm_state[:, 6, head] == 30 + start + head).all() + assert main_linear_state(manager).mtp_state_index[2] == 0 + + +def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: + try: + torch.cuda.set_device(0) + for kind in ipc_kinds: + resources, mapping = shared_att_resources_utils.load_att_resources(0, kind=kind) + assert resources._states_created + assert resources.full_att_allocator is None + assert resources.bounded_att_big_page_allocator is None + assert resources.bounded_att_small_page_allocator is None + if kind == "dp": + for state in resources.get_bounded_attention_states(): + assert state.big_pages_pool.conv_state.is_shared() + assert not state.big_pages_pool.conv_state.is_pinned() + resources.register_page_pools_pin_memory() + resources.register_page_pools_pin_memory() + for state in resources.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if kind == "dp": + assert pool is not None + assert pool.conv_state.is_pinned() and pool.ssm_state.is_pinned() + copied = pool.conv_state.to(device="cuda", non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(copied.cpu(), pool.conv_state) + else: + assert pool is None + if kind == "dp": + main_linear_state(resources).small_pages_pool.conv_state.flatten()[0] = 37 + assert not hasattr(resources, "req_to_token_indexs") + assert mapping[1, 0].item() == 7 + resources.copy_full_att_tokens(torch.tensor([0], dtype=torch.long, device="cuda"), torch.tensor([4], dtype=torch.long, device="cuda")) + resources.reset_attention_tail_state([1]) + torch.cuda.synchronize() + # 一类通道耗尽不会影响另一类通道的读取计数。 + with pytest.raises(RuntimeError, match="exhausted"): + shared_att_resources_utils.load_att_resources(0, kind=kind) + connection.send(None) + # Keep imported tensors alive until the parent has checked the writes. + connection.recv() + except BaseException: + connection.send(traceback.format_exc()) + finally: + AttResourceManager._get_states.cache_clear() + connection.close() diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py new file mode 100644 index 0000000000..b61adeef99 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from collections.abc import Callable + +import pytest +import torch + +from lightllm.common.att_resources.full_att import ( + FullAttentionConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, + StaticFP8FullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, + create_gqa_state, create_mla_state, +) +from lightllm.common.att_resources.bounded_att import ( + BoundedAttentionState, ConvSSMStateConfig, create_gdn_state, create_kda_state, +) +from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state + + +@pytest.mark.parametrize( + "kv_type, expected_type, storage_dtype, width", + [ + ("None", "full_att_bf16", torch.bfloat16, 128), + ("int4kv", "full_att_int4_group16_bf16_scale", torch.int8, 64), + ("int8kv", "full_att_int8_group16_bf16_scale", torch.int8, 128), + ("fp8kv_spt", "full_att_fp8_static_per_tensor", torch.uint8, 128), + ("fp8kv_sph", "full_att_fp8_static_per_head", torch.uint8, 128), + ], +) +def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( + kv_type: str, expected_type: str, storage_dtype: torch.dtype, width: int, +): + config = FullAttentionConfig( + layer_ids=(3, 7), global_head_num=8, head_dim=128, dtype=torch.bfloat16, tp_rank=2, tp_world_size=4, + ) + scales = torch.full((2, 4 if kv_type == "fp8kv_sph" else 2), 0.5, dtype=torch.float32) + q_scales = torch.ones((2, 2), dtype=torch.float32) if kv_type == "fp8kv_sph" else None + state = create_gqa_state( + config=config, kv_type=kv_type, page_size=4, device="meta", group_size=16, + scales=scales, q_scales=q_scales, calibration_id="test-calibration", + ) + assert state.get_state_type() == expected_type + assert state.resource_id is None + assert state.config.layer_ids == (3, 7) + assert state.config.global_head_num == 8 and state.config.head_num == 2 + assert state.config.tp_rank == 2 and state.config.tp_world_size == 4 + assert state.config.head_dim == config.head_dim and config.dtype == torch.bfloat16 + assert state.capacity == 0 and read_kv_storage(state) is None + if kv_type.startswith("fp8"): + assert state.scales is scales and not state.scales.is_meta + assert state.q_scales is q_scales + state.device = torch.device("cpu") + state.initialize_storage(8) + assert state.k_storage.shape == state.v_storage.shape == (2, 12, 2, width) + assert state.k_storage.untyped_storage().data_ptr() != state.v_storage.untyped_storage().data_ptr() + assert state.k_storage.dtype == state.v_storage.dtype == storage_dtype + if kv_type in ("int4kv", "int8kv"): + assert state.k_scales.shape == state.v_scales.shape == (2, 12, 2, 8) + assert state.k_scales.untyped_storage().data_ptr() != state.v_scales.untyped_storage().data_ptr() + + +@pytest.mark.parametrize( + "packed, kv_type, kv_dim, row_width, expected_type", + [ + (False, "None", 576, 576, "mla_bf16"), + (True, "None", 576, 648, "mla_packed_bf16"), + (True, "None", 512, 584, "mla_packed_bf16"), + (True, "fp8kv_dsa", 576, 800, "mla_packed_fp8"), + ], +) +def test_mla_creation_computes_packed_width_from_logical_dimensions( + packed: bool, kv_type: str, kv_dim: int, row_width: int, expected_type: str, +): + config = FullAttentionConfig( + layer_ids=(2, 6), global_head_num=1, head_dim=kv_dim, dtype=torch.bfloat16, tp_rank=1, tp_world_size=4, + ) + state = create_mla_state(config=config, packed=packed, kv_type=kv_type, page_size=4, device="meta") + assert state.get_state_type() == expected_type + assert state.config.replicated and not config.replicated + assert state.config.head_dim == row_width and config.head_dim == kv_dim + assert state.capacity == 0 and read_kv_storage(state) is None + state.device = torch.device("cpu") + state.initialize_storage(8) + assert read_kv_storage(state).shape == (2, 12, 1, row_width) + if packed: + assert state.config.kv_dim == kv_dim + assert state.get_layer_view(2).indexer_k.shape[0] == 12 + + +@pytest.mark.parametrize("create_state", [create_kda_state, create_gdn_state]) +def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Callable[..., BoundedAttentionState]): + config = ConvSSMStateConfig( + layer_ids=(0, 2), global_k_heads=2, global_v_heads=2, + key_dim=8, value_dim=8, conv_width=3, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + ) + state = create_state(config=config, request_capacity=3, mtp_step=2, device="cpu") + assert state.config is config and state.resource_id is None + assert state.conv_state.shape == (2, 3, 48, 5) + assert state.ssm_state.shape == (2, 9, 2, 8, 8) + assert state.mtp_state_index.shape == (3,) + assert state.big_pages_pool is None and state.small_pages_pool is None + + +def test_kpool_creation_preserves_pool_and_mtp_dimensions(): + config = KPoolStateConfig(layer_ids=(3, 7), pool_size=4, head_dim=128, dtype=torch.bfloat16) + state = create_kpool_state(config=config, request_capacity=3, mtp_step=2, device="cpu") + assert state.config is config and state.resource_id is None + assert state.tail_buffer.shape == (2, 3, 6, 256) + assert state.get_state_type() == "indexer_kpool_c4_tail" + + +def test_create_rejects_encoding_from_another_logical_family(): + config = FullAttentionConfig(layer_ids=(0,), global_head_num=1, head_dim=576, dtype=torch.bfloat16) + with pytest.raises(AssertionError, match="unsupported GQA encoding"): + create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") + with pytest.raises(AssertionError, match="unsupported packed MLA encoding"): + create_mla_state(config=config, packed=True, kv_type="int4kv", device="meta") + with pytest.raises(AssertionError, match="unsupported MLA encoding"): + create_mla_state(config=config, kv_type="fp8kv_dsa", device="meta") + + +@pytest.mark.parametrize( + "state_class", + [GQAFullAttentionState, GroupQuantizedFullAttentionState, StaticFP8FullAttentionState, + MLAFullAttentionState, PackedMLAFullAttentionState], +) +def test_full_attention_implementations_are_independent(state_class: type[FullAttentionState]): + # 具体方案只共享抽象契约，不继承其他具体方案的存储或操作实现。 + assert state_class.__bases__ == (FullAttentionState,) + assert not state_class.__abstractmethods__ + assert FullAttentionState.__abstractmethods__ <= state_class.__dict__.keys() diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py new file mode 100644 index 0000000000..34d0169666 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from easydict import EasyDict + + +KINDS = ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] + + +@pytest.mark.parametrize("kind", KINDS) +def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_args: EasyDict): + from inspect import signature + + state = definition(kind) + parameters = signature(type(state)).parameters + assert not {"token_capacity", "buffer", "scale_buffer"} & parameters.keys() + assert state.capacity == 0 and read_kv_storage(state) is None + token_bytes = state.get_token_bytes() + cpu_regions = state.get_cpu_transfer_regions() + pd_regions = state.get_pd_transfer_regions() + assert read_kv_storage(state) is None + original_config = state.config + state.device = torch.device("cpu") + state.initialize_storage(8) + assert state.config is original_config + assert state.capacity == 8 and read_kv_storage(state).shape[1] == 8 + state.page_size + assert not read_kv_storage(state).is_meta + assert state.get_token_bytes() == token_bytes + assert state.get_cpu_transfer_regions() == cpu_regions + assert state.get_pd_transfer_regions() == pd_regions + actual_bytes = sum(buffer.numel() * buffer.element_size() for buffer in full_state_test_buffers(state).values()) + assert actual_bytes == token_bytes * (8 + state.page_size) + + +def test_numbered_drafts_keep_resource_id_order_and_unique_members(): + manager = AttResourceManager(request_capacity=1) + for role in ("draft35", "draft10", "draft1", "main", "draft0", "draft2"): + manager.set_att_state(role, definition("gqa")) + with pytest.raises(AssertionError, match="already bound"): + manager.set_att_state("draft1", definition("gqa")) + manager.finish_state_creation() + assert [state.resource_id.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ + "draft0", "draft10", "draft1", "draft2", "draft35", "main", + ] + assert manager.get_full_attention_states("draft1") == (manager.draft1_full_att_bf16_state,) + assert manager.get_full_attention_states("draft10") == (manager.draft10_full_att_bf16_state,) + + +def test_custom_state_type_controls_registered_resource_name(): + from unit_tests.common.att_resources.test_full_state import full_state + from lightllm.common.att_resources.full_att import GQAFullAttentionState + + class CustomGQAState(GQAFullAttentionState): + def get_state_type(self) -> str: + return "full_att_custom_fp32" + + state = full_state(CustomGQAState) + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("draft0", state) + manager.finish_state_creation() + assert manager.draft0_full_att_custom_fp32_state is state + assert state.resource_id == "draft0_full_att_custom_fp32_state" + assert manager.get_full_attention_states("draft0") == (state,) + assert not hasattr(state, "state_type") + + +@pytest.mark.parametrize("kind", ["full_meta", "full", "bounded", "tail"]) +@pytest.mark.parametrize("model_name", ["main", "draft1"]) +def test_set_att_state_preserves_object_and_storage(kind: str, model_name: str): + from unit_tests.common.att_resources.test_bounded_state import make_state + from unit_tests.common.att_resources.test_tail_state import make_tail + + if kind == "full_meta": + state = definition("gqa") + storage = tuple(full_state_test_buffers(state).values()) + elif kind == "full": + state = allocated_state("gqa") + storage = tuple(full_state_test_buffers(state).values()) + elif kind == "bounded": + state = make_state() + storage = state.conv_state + pool = state.big_pages_pool + else: + state = make_tail() + storage = state.tail_buffer + config = state.config + manager = AttResourceManager(request_capacity=3) + manager.set_att_state(model_name, state) + name = f"{model_name}_{state.get_state_type()}_state" + assert getattr(manager, name) is state + assert state.resource_id == name + assert state.config is config and not hasattr(config, "resource_id") + if kind.startswith("full"): + assert all(actual is original for actual, original in zip(full_state_test_buffers(state).values(), storage)) + elif kind == "bounded": + assert state.conv_state is storage and state.big_pages_pool is pool + else: + assert state.tail_buffer is storage + other_role = "draft0" if model_name == "main" else "main" + with pytest.raises(AssertionError, match="multiple model roles"): + manager.set_att_state(other_role, state) + assert state.resource_id == name and getattr(manager, name) is state + manager.finish_state_creation() + with pytest.raises(AssertionError, match="creation is already finished"): + manager.set_att_state(other_role, definition("gqa")) + + +def test_model_resource_role_numbers_the_36th_draft_and_rejects_the_37th(): + from lightllm.models.llama.model import LlamaTpPartModel + + draft = LlamaTpPartModel.__new__(LlamaTpPartModel) + draft.is_mtp_draft_model = True + draft.mtp_previous_draft_models = [object()] * 35 + assert draft.att_resource_role == "draft35" + draft.mtp_previous_draft_models.append(object()) + with pytest.raises(AssertionError, match="at most 36"): + _ = draft.att_resource_role + + +@pytest.mark.parametrize("role", ["draft", "draft36", "draft-1", "draft01"]) +def test_set_att_state_rejects_invalid_draft_numbers(role: str): + manager = AttResourceManager(request_capacity=1) + with pytest.raises(AssertionError, match="unknown model role"): + manager.set_att_state(role, definition("gqa")) + + +def allocated_state(kind: str, *, device: str = "cpu"): + template = definition(kind) + return rebuild_full_state(template, config=template.config, token_capacity=8, device=device) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", KINDS) +def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + physical = allocated_state(kind, device="cuda") + main, draft = create_full_attention_roles(physical, 1) + assert main.config.layer_ids == (0,) and draft.config.layer_ids == (2,) + for role, state, layer in (("main", main, 0), ("draft0", draft, 1)): + physical_buffers = full_state_test_buffers(physical) + for name, buffer in full_state_test_buffers(state).items(): + original = physical_buffers[name] + assert buffer.untyped_storage().data_ptr() != original.untyped_storage().data_ptr() + buffer.fill_(layer + 3) + if isinstance(state, StaticFP8FullAttentionState): + for buffer, original in ((state.scales, physical.scales), (state.q_scales, physical.q_scales)): + if buffer is not None: + assert buffer.untyped_storage().data_ptr() != original.untyped_storage().data_ptr() + + source = AttResourceManager(request_capacity=2) + source.set_att_state("main", main) + source.set_att_state("draft0", draft) + source.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + source.finish_state_creation() + assert getattr(source, main.resource_id) is main + assert getattr(source, draft.resource_id) is draft + source.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) + + target_physical = allocated_state(kind, device="cuda") + target_main, target_draft = create_full_attention_roles(target_physical, 1) + target = AttResourceManager(request_capacity=2) + target.set_att_state("main", target_main) + target.set_att_state("draft0", target_draft) + target.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + target.finish_state_creation() + page = torch.zeros((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.write_req_to_cpu(torch.tensor([0, 1], dtype=torch.long, device="cuda"), page, [0], layouts=source.get_cpu_transfer_layouts()) + target.read_req_from_cpu(page, [0], torch.tensor([6, 7], dtype=torch.long, device="cuda"), layouts=source.get_cpu_transfer_layouts()) + for src, dst in ((main, target_main), (draft, target_draft)): + fields = {name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items()} + for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): + if "token" in field.axes: + assert torch.equal(field.tensor[:, :2], fields[name][:, 6:8]) + assert set(source.get_cpu_transfer_layouts()) == {main.resource_id, draft.resource_id} + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_layers(overlap: bool): + main = allocated_state("gqa") + draft = allocated_state("gqa") + manager = AttResourceManager(request_capacity=2) + manager.set_att_state("main", main) + manager.set_att_state("draft0", draft) + manager.full_att_allocator = FullAttKVAllocator(main.capacity, main.page_size, pin_memory=False) + template = allocated_state("int8") + quant = rebuild_full_state(template, config=replace(template.config, layer_ids=(4, 6))) + manager.set_att_state("main", quant) + if overlap: + overlapping = allocated_state("static_head") + manager.set_att_state("main", overlapping) + manager.finish_state_creation() + assert manager.main_full_att_bf16_state is main + assert manager.draft0_full_att_bf16_state is draft + assert manager.main_full_att_int8_group8_bf16_scale_state is quant + assert manager.get_full_attention_states("draft0") == (draft,) + expected = [main, quant] + if overlap: + assert manager.main_full_att_fp8_static_per_head_state is overlapping + expected.append(overlapping) + assert manager.get_full_attention_states("main") == tuple(sorted(expected, key=lambda state: state.resource_id)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + full = allocated_state("gqa", device="cuda") + quant = allocated_state("int8", device="cuda") + assert full.get_state_type() == "full_att_bf16" + assert quant.get_state_type() == "full_att_int8_group8_bf16_scale" + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor.fill_(5) + field.tensor[:, 4:6].zero_() + manager = AttResourceManager(request_capacity=2) + for state in (full, quant): + manager.set_att_state("main", state) + manager.finish_state_creation() + manager.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 4:6] == 5).all() + layouts = manager.get_cpu_transfer_layouts() + page = torch.zeros((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) + manager.write_req_to_cpu( + torch.tensor([0, 1], dtype=torch.long, device="cuda"), page, [0], layouts=layouts + ) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor[:, 6:8].zero_() + manager.read_req_from_cpu( + page, [0], torch.tensor([6, 7], dtype=torch.long, device="cuda"), layouts=layouts + ) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 6:8] == 5).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", KINDS) +def test_token_copy_captures_large_index_sets_and_preserves_complete_buffers(kind: str): + template = definition(kind) + state = rebuild_full_state(template, token_capacity=512, device="cuda") + # 覆盖 int32/int64 和非连续索引；128 个元素足以触发旧 isin 实现的同步路径。 + source = torch.arange(256, dtype=torch.int32, device="cuda")[::2] + target = torch.arange(256, 512, dtype=torch.int64, device="cuda")[::2] + buffers = tuple(full_state_test_buffers(state).values()) + for buffer in buffers: + buffer.view(torch.uint8).random_(0, 256) + state.copy_full_att_tokens(source, target) # 先完成 Triton 编译。 + torch.cuda.synchronize() + + expected = [] + for buffer in buffers: + buffer.view(torch.uint8).random_(0, 256) + snapshot = buffer.view(torch.uint8).clone() + snapshot[:, target] = snapshot[:, source] + expected.append(snapshot) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + state.copy_full_att_tokens(source, target) + for buffer in buffers: + buffer.index_fill_(1, target, 0) + graph.replay() + torch.cuda.synchronize() + for buffer, snapshot in zip(buffers, expected): + # 逐字节核对整块存储，包含量化 scale、packed padding 和未参与搬运的槽位。 + assert torch.equal(buffer.view(torch.uint8), snapshot) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py new file mode 100644 index 0000000000..375e58cf98 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -0,0 +1,181 @@ +"""State 对象创建完成后即可查询，容量分配保持对象身份和查询缓存不变。""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + + +import gc +import pickle +import weakref +from functools import partial + +import pytest +from easydict import EasyDict + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager +from unit_tests.common.att_resources.test_capacity import definition + + +def test_created_state_queries_keep_the_same_objects_through_storage_initialization(monkeypatch: pytest.MonkeyPatch): + args = StartArgs(page_size=4) + definitions = [] + manager = AttResourceManager(request_capacity=1) + for role, kind in (("main", "gqa"), ("draft0", "int8")): + manager.set_att_state(role, definition(kind)) + definitions.append(next( + state for state in vars(manager).values() + if isinstance(state, FullAttentionState) and state.resource_id.startswith(role + "_") + )) + for query in ( + manager.get_full_attention_states, manager.get_bounded_attention_states, manager.get_attention_tail_states, + ): + with pytest.raises(AssertionError, match="State creation"): + query() + assert manager._get_states.cache_info().currsize == 0 + + cached = [] + initialize_storage = AttResourceManager.initialize_storage + + def check_created_before_allocation(manager: AttResourceManager, *args: object, **kwargs: object): + states = manager.get_full_attention_states() + assert all(read_kv_storage(state) is None for state in states) + assert manager.full_att_allocator is None + cached.append(states) + initialize_storage(manager, *args, **kwargs) + + monkeypatch.setattr(AttResourceManager, "initialize_storage", check_created_before_allocation) + monkeypatch.setattr( + "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.initialize_resources(args, token_capacity=8, mem_fraction=1.0, device="cpu") + assert manager.get_full_attention_states() is cached[0] + for role, meta in zip(("main", "draft0"), definitions): + (state,) = manager.get_full_attention_states(role) + assert state is meta and read_kv_storage(state).device.type == "cpu" + assert manager.get_full_attention_states(role) == (state,) + + +def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): + manager = MainDraftTestResourceManager() + queries = (manager.get_full_attention_states, manager.get_bounded_attention_states, manager.get_attention_tail_states) + results = {(query.__name__, role): query(role) for query in queries for role in (None, "main", "draft0")} + assert manager._get_states.cache_info().currsize == 9 + assert manager.get_full_attention_states() == ( + manager.draft0_mla_fp32_state, manager.main_full_att_fp32_state, + ) + assert manager.get_bounded_attention_states("main") == (manager.main_kda_state,) + assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_state,) + assert manager.get_attention_tail_states("draft0") == () + + hits = manager._get_states.cache_info().hits + for query in queries: + for role in (None, "main", "draft0"): + assert query(role) is results[(query.__name__, role)] + with pytest.raises(AssertionError, match="unknown model role"): + query("target") + assert manager._get_states.cache_info().hits == hits + 9 + manager.main_kda_state.conv_state.fill_(7) + assert (manager.get_bounded_attention_states("main")[0].conv_state == 7).all() + + +def test_finished_state_creation_allows_unallocated_storage(): + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", definition("gqa")) + manager.full_att_allocator = FullAttKVAllocator(0, 4, pin_memory=False) + manager.finish_state_creation() + assert read_kv_storage(manager.get_full_attention_states()[0]) is None + + +def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): + source = MainDraftTestResourceManager() + source_states = source.get_bounded_attention_states() + owner_pools = [(state.big_pages_pool, state.small_pages_pool) for state in source_states] + + exported = source.create_pd_ipc_copy() + + assert type(exported) is type(source) and exported is not source + for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): + assert getattr(source, name) is not None + assert getattr(exported, name) is None + for state, (big_pool, small_pool) in zip(source_states, owner_pools): + copied = getattr(exported, state.resource_id) + assert copied is not state + assert copied.conv_state is state.conv_state + assert copied.ssm_state is state.ssm_state + assert copied.mtp_state_index is state.mtp_state_index + assert copied.big_pages_pool is None and copied.small_pages_pool is None + assert copied.request_capacity == state.request_capacity + assert copied.mtp_step == state.mtp_step + assert state.big_pages_pool is big_pool and state.small_pages_pool is small_pool + assert big_pool.size == source.bounded_att_big_page_allocator.size + assert small_pool.size == source.bounded_att_small_page_allocator.size + assert exported.get_bounded_attention_states() == tuple( + getattr(exported, state.resource_id) for state in source_states + ) + for state in (*source.get_full_attention_states(), *source.get_attention_tail_states()): + assert getattr(exported, state.resource_id) is state + + + +@pytest.mark.parametrize("share_pools", [False, True]) +def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share_pools: bool) -> None: + cpu_args.enable_dp_prompt_cache_fetch = share_pools + source = MainDraftTestResourceManager() + pd_copy = source.create_pd_ipc_copy() + for state in source.get_bounded_attention_states(): + exported = getattr(pd_copy, state.resource_id) + assert exported.big_pages_pool is None and exported.small_pages_pool is None + assert state.big_pages_pool is not None and state.small_pages_pool is not None + if not share_pools: + with pytest.raises(AssertionError, match="shared CPU page pools"): + source.create_dp_ipc_copy() + return + dp_copy = source.create_dp_ipc_copy() + for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): + assert getattr(dp_copy, name) is None + assert getattr(source, name) is not None + for state in source.get_bounded_attention_states(): + exported = getattr(dp_copy, state.resource_id) + assert exported.big_pages_pool is state.big_pages_pool + assert exported.small_pages_pool is state.small_pages_pool + assert exported.conv_state is state.conv_state + + +def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): + source = MainDraftTestResourceManager() + source.get_full_attention_states() + source.get_bounded_attention_states() + source.get_attention_tail_states() + imported = pickle.loads(pickle.dumps(source.create_pd_ipc_copy())) + assert imported._states_created + assert imported.full_att_allocator is None + assert imported.bounded_att_big_page_allocator is None + assert imported.bounded_att_small_page_allocator is None + misses = source._get_states.cache_info().misses + for query in ( + imported.get_full_attention_states, imported.get_bounded_attention_states, imported.get_attention_tail_states, + ): + for state in query(): + name = state.resource_id + assert state is getattr(imported, name) and state is not getattr(source, name) + assert imported._get_states.cache_info().misses == misses + 3 + + +def test_clearing_query_cache_releases_manager_and_states(): + manager = MainDraftTestResourceManager() + manager.get_full_attention_states() + manager.get_bounded_attention_states() + manager.get_attention_tail_states() + manager_ref = weakref.ref(manager) + state_ref = weakref.ref(manager.main_kda_state) + del manager + gc.collect() + assert manager_ref() is not None and state_ref() is not None + AttResourceManager._get_states.cache_clear() + gc.collect() + assert manager_ref() is None and state_ref() is None diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py new file mode 100644 index 0000000000..20ca398659 --- /dev/null +++ b/unit_tests/common/att_resources/test_tail_state.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import pytest +import torch + +from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.att_resources.types import ( + AxisSlice, StateTransferLayout, TransferRegion, +) + + +def make_tail(*, device: str | torch.device | int = "cpu", mtp_step: int = 2): + return KPoolAttentionTailState( + KPoolStateConfig(layer_ids=(3, 7), pool_size=4, head_dim=2, dtype=torch.float32), + request_capacity=3, mtp_step=mtp_step, device=device + ) + + +def tail_layout(state: KPoolAttentionTailState): + shape = (2, state.config.pool_size, 4) + nbytes = shape[0] * shape[1] * shape[2] * 4 + return StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=nbytes + 16, + regions=( + TransferRegion( + field="tail", + byte_start=16, + byte_end=16 + nbytes, + dtype=torch.float32, + axes=("layer", "window", "channel"), + shape=shape, + axis_slices=( + AxisSlice(axis="window", src_start=0, dst_start=0, length=shape[1]), + AxisSlice(axis="channel", src_start=0, dst_start=0, length=4), + ), + ), + ), + ) + + +def test_tail_reset_clears_selected_requests_only(): + state = make_tail() + state.tail_buffer.fill_(17) + state.reset_attention_tail_state([0, 1]) + assert not state.tail_buffer[:, :2].any() + assert (state.tail_buffer[:, 2] == 17).all() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) +@pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) +def test_pd_logical_window_restores_different_mtp_rings( + device: str, source_step: int, target_step: int, length: int, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source = make_tail(device=device, mtp_step=source_step) + target = make_tail(device=device, mtp_step=target_step) + source.tail_buffer.copy_(torch.arange(source.tail_buffer.numel(), device=device).reshape_as(source.tail_buffer)) + original = source.tail_buffer.clone() + target.tail_buffer.fill_(37) + layout = tail_layout(source) + assert layout == tail_layout(target) + assert source.get_pd_transfer_regions() == target.get_pd_transfer_regions() + pages = torch.full((3, layout.byte_end + 16), 253, dtype=torch.uint8, device=device) + source.write_to_pd_buffer(1, pages, pd_page_index=1, state_token_len=length, layout=layout) + target.read_from_pd_buffer(pages, 2, pd_page_index=1, state_token_len=length, layout=layout) + expected = torch.zeros_like(target.tail_buffer[:, 2]) + for position in range(max(0, length - 4), length): + expected[:, position % (4 + target_step)] = source.tail_buffer[:, 1, position % (4 + source_step)] + assert torch.equal(target.tail_buffer[:, 2], expected) + assert (target.tail_buffer[:, :2] == 37).all() + assert torch.equal(source.tail_buffer, original) + assert (pages[0] == 253).all() and (pages[2] == 253).all() + assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end:] == 253).all() + if length < 4: + external = pages[1, 16:layout.byte_end].view(torch.float32).view(2, 4, 4) + assert not external[:, :4 - length].any() diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py new file mode 100644 index 0000000000..743cb67c63 --- /dev/null +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -0,0 +1,165 @@ +"""State 配置独立保留全局布局，改变 TP 后仍可恢复同一份缓存。""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +import pickle +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.att_resources.full_att import ( + FullAttentionConfig, + GQAFullAttentionState, + MLAFullAttentionState, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from easydict import EasyDict + from lightllm.common.att_resources.bounded_att.base_state import BoundedAttentionState + from lightllm.common.att_resources.full_att.base_state import FullAttentionState + + +def cpu_layout(state: FullAttentionState | BoundedAttentionState): + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", state) + manager.finish_state_creation() + return manager.get_cpu_transfer_layouts()[state.resource_id] + + +@pytest.mark.parametrize( + "state_cls,heads,replicated,source_ranges,target_ranges", + [ + (GQAFullAttentionState, 4, False, [(0, 2), (2, 4)], [(i, i + 1) for i in range(4) for _ in range(2)]), + (GQAFullAttentionState, 2, False, [(0, 1)] * 4 + [(1, 2)] * 4, [(0, 2)]), + (GQAFullAttentionState, 2, True, [(0, 2)] * 4, [(0, 2)] * 2), + (MLAFullAttentionState, 1, True, [(0, 1)] * 4, [(0, 1)]), + ], +) +def test_serialized_full_config_restores_cpu_pages_with_different_tp( + state_cls: type[FullAttentionState], + heads: int, + replicated: bool, + source_ranges: Sequence[tuple[int, int]], + target_ranges: Sequence[tuple[int, int]], + cpu_args: EasyDict, +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cpu_args.cpu_cache_token_page_size = 4 + config = FullAttentionConfig((0,), heads, 4, torch.float32, replicated=replicated) + reference = state_cls(config, page_size=4, device="cuda") + reference.initialize_storage(4) + for storage_buffer in full_state_test_buffers(reference).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape(storage_buffer.shape)) + fields = {name: field.tensor for name, field in full_state_test_fields(reference, full_state_test_buffers(reference)).items()} + reference_layout = cpu_layout(reference) + buffer = torch.zeros((1, reference_layout.byte_end), dtype=torch.uint8, pin_memory=True) + + # 接收方只需要配置和缓存页；config 中不保存本地 head 数或 manager 映射。 + restored_config = pickle.loads(pickle.dumps(config)) + assert "head_num" not in vars(restored_config) + for rank, (start, end) in enumerate(source_ranges): + source = state_cls(replace(config, tp_rank=rank, tp_world_size=len(source_ranges)), page_size=4, device="cuda") + source.initialize_storage(4) + layout = cpu_layout(source) + assert layout.byte_end == reference_layout.byte_end + for name, field in full_state_test_fields(source, full_state_test_buffers(source)).items(): + field.tensor.copy_(fields[name][:, :, start:end]) + source.write_to_cpu(torch.arange(4, dtype=torch.long, device="cuda"), buffer, [0], layout=layout) + + for rank, (start, end) in enumerate(target_ranges): + target = state_cls(replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda") + target.initialize_storage(4) + layout = cpu_layout(target) + assert layout.byte_end == reference_layout.byte_end + target.read_from_cpu(buffer, [0], torch.arange(4, dtype=torch.long, device="cuda"), layout=layout) + for name, field in full_state_test_fields(target, full_state_test_buffers(target)).items(): + torch.testing.assert_close(field.tensor[:, :4], fields[name][:, :4, start:end]) + + +def test_bounded_config_keeps_distinct_global_heads_when_local_heads_are_replicated(cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=3, conv_width=2, + conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_rank=7, tp_world_size=8, + ) + source = GDNState(config, request_capacity=1, device="meta") + assert config.num_k_heads == config.num_v_heads == 1 + restored_config = replace(pickle.loads(pickle.dumps(config)), tp_rank=0, tp_world_size=1) + target = GDNState(restored_config, request_capacity=1, device="meta") + assert restored_config.num_k_heads == 2 and restored_config.num_v_heads == 4 + assert "num_k_heads" not in vars(config) and "num_v_heads" not in vars(config) + assert target.conv_state.shape == (1, 1, 20, 2) + assert target.ssm_state.shape == (1, 1, 4, 2, 3) + source_layout = cpu_layout(source) + target_layout = cpu_layout(target) + assert source_layout.byte_end == target_layout.byte_end + assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in source_layout.regions] == [ + (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in target_layout.regions + ] + + +@pytest.mark.parametrize("bounded", [False, True]) +@pytest.mark.parametrize("page_tokens", [4, 8]) +def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( + bounded: bool, + page_tokens: int, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = page_tokens + if bounded: + config = ConvSSMStateConfig( + layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=2, conv_width=2, + conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_world_size=8, + ) + source = GDNState(config, request_capacity=1, device="meta") + other_config = replace(config, layer_ids=(1,), global_k_heads=4) + else: + config = FullAttentionConfig((0,), 2, 4, torch.float32, tp_world_size=8) + source = GQAFullAttentionState(config, page_size=4, device="meta") + other_config = replace(config, layer_ids=(1,), global_head_num=4) + if bounded: + other = type(source)( + other_config, request_capacity=source.request_capacity, mtp_step=source.mtp_step, device="meta" + ) + else: + other = rebuild_full_state(source, config=other_config) + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("draft0", source) + manager.set_att_state("draft1", other) + manager.finish_state_creation() + query = manager.get_bounded_attention_states if bounded else manager.get_full_attention_states + assert query("draft0") == (source,) + assert query("draft1") == (other,) + assert query() == (source, other) + + page_bytes = manager.get_pd_page_bytes() + assert not hasattr(manager, "pd_page_bytes") + assert not hasattr(manager, "pd_transfer_layouts") + if not bounded: + # 两个 draft 共 6 个全局 KV head：K/V、head_dim=4、FP32，无额外页头。 + assert page_bytes == page_tokens * 192 + assert manager.get_pd_page_bytes() == page_bytes + getter = ( + manager.get_pd_bounded_att_or_tail_transfer_layouts if bounded else manager.get_pd_full_att_transfer_layouts + ) + layouts = getter() + assert getter() is layouts + assert set(layouts) == {source.resource_id, other.resource_id} + assert all(layout.tokens_per_page == (None if bounded else page_tokens) for layout in layouts.values()) + unused_getter = ( + manager.get_pd_full_att_transfer_layouts if bounded else manager.get_pd_bounded_att_or_tail_transfer_layouts + ) + assert unused_getter() == {} + assert manager.get_pd_page_bytes() == page_bytes diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index eeb2fb1438..38a5113848 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -12,7 +12,7 @@ import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig from lightllm.server.api_cli import make_argument_parser import lightllm.utils.backend_validator as backend_validator @@ -62,29 +62,6 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): assert torch.equal(ssm_states, final_state.to(cache_dtype)) -def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - infer_state = InferStateInfo() - infer_state.batch_size = b_req_idx.shape[0] - infer_state.b_req_idx = b_req_idx - infer_state.b_mtp_index = torch.zeros_like(b_req_idx) - infer_state.req_manager = SimpleNamespace( - HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, - req_to_mtp_state_index=req_to_mtp_state_index, - ) - state = gdn.LinearAttDecodeAttState( - backend=SimpleNamespace( - model=model, - uses_dynamic_spec_verify_layout=lambda: dynamic_layout, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - return state, infer_state - - def test_decode_state_initializes_normal_layout(): b_req_idx = torch.tensor([2, 4], dtype=torch.int32) state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) @@ -205,7 +182,7 @@ def test_flashqla_validation_loads_linear_cache_config(monkeypatch): ) load_calls = [] monkeypatch.setattr( - LinearAttCacheConfig, + Qwen3NextResourceConfig, "load_from_args", staticmethod(lambda: load_calls.append(True) or linear_config), ) @@ -359,3 +336,26 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ assert backend_class is FlashQlaLinearAttBackend assert validate_calls == ["flashqla2", "flashqla3"] + + +def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None): + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), + ) + infer_state = InferStateInfo() + infer_state.batch_size = b_req_idx.shape[0] + infer_state.b_req_idx = b_req_idx + infer_state.b_mtp_index = torch.zeros_like(b_req_idx) + infer_state.req_manager = SimpleNamespace( + HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, + ) + state = gdn.LinearAttDecodeAttState( + backend=SimpleNamespace( + model=model, + linear_state=SimpleNamespace(mtp_state_index=req_to_mtp_state_index), + uses_dynamic_spec_verify_layout=lambda: dynamic_layout, + ), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + return state, infer_state diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index cb0d9fef54..123f603ff2 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -80,24 +80,6 @@ def test_prefill_unpad_restores_empty_output(): assert unpadded.prompt_logics.shape == (0, 4) -def _create_empty_decode_input(): - return ModelInput( - batch_size=0, - total_token_num=0, - max_q_seq_len=1, - max_kv_seq_len=0, - input_ids=torch.empty((0,), dtype=torch.int64), - b_req_idx=torch.empty((0,), dtype=torch.int32), - b_mtp_index=torch.empty((0,), dtype=torch.int32), - b_seq_len=torch.empty((0,), dtype=torch.int32), - b_position_delta=torch.empty((0,), dtype=torch.int32), - b_shared_seq_len=torch.empty((0,), dtype=torch.int32), - b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), - is_prefill=False, - multimodal_params=[], - ) - - @torch.no_grad() def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): execution_configs = ( @@ -122,7 +104,6 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): model.tp_world_size_ = tp_world_size model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) pad_batch_sizes = [] @@ -195,7 +176,6 @@ def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(le model.tp_world_size_ = 2 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) initializations = [] executions = [] @@ -344,3 +324,21 @@ def forward(state): assert seen_sizes == [expected_size] * (2 if overlap else 1) torch.testing.assert_close(output.logits, expected) torch.testing.assert_close(torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1), expected) + + +def _create_empty_decode_input(): + return ModelInput( + batch_size=0, + total_token_num=0, + max_q_seq_len=1, + max_kv_seq_len=0, + input_ids=torch.empty((0,), dtype=torch.int64), + b_req_idx=torch.empty((0,), dtype=torch.int32), + b_mtp_index=torch.empty((0,), dtype=torch.int32), + b_seq_len=torch.empty((0,), dtype=torch.int32), + b_position_delta=torch.empty((0,), dtype=torch.int32), + b_shared_seq_len=torch.empty((0,), dtype=torch.int32), + b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), + is_prefill=False, + multimodal_params=[], + ) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index 54f2ccacf1..ac55020ee6 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -8,68 +8,6 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -def _empty_multimodal_params(batch_size: int): - return [{"images": [], "audios": []} for _ in range(batch_size)] - - -def _make_prefill_input(input_ids: list[int], req_idx: int, is_decode_req: bool) -> ModelInput: - token_num = len(input_ids) - return ModelInput( - batch_size=1, - total_token_num=token_num, - max_q_seq_len=token_num, - max_kv_seq_len=token_num, - max_cache_len=0, - input_ids=torch.tensor(input_ids, dtype=torch.int64), - b_req_idx=torch.tensor([req_idx], dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_seq_len=torch.tensor([token_num], dtype=torch.int32), - b_is_decode_req=torch.tensor([is_decode_req]), - b_ready_cache_len=torch.zeros(1, dtype=torch.int32), - b_prefill_start_loc=torch.zeros(1, dtype=torch.int32), - is_prefill=True, - b_prefill_has_output_cpu=[True], - multimodal_params=_empty_multimodal_params(1), - ) - - -def _make_decode_input() -> ModelInput: - batch_size = 6 - return ModelInput( - batch_size=batch_size, - total_token_num=39, - max_q_seq_len=1, - max_kv_seq_len=9, - input_ids=torch.arange(batch_size, dtype=torch.int64), - b_req_idx=torch.tensor([10, 11, 12, 12, 12, 12], dtype=torch.int32), - b_mtp_index=torch.tensor([0, 0, 0, 1, 2, 3], dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 7, 8, 9], dtype=torch.int32), - b_position_delta=torch.arange(batch_size, dtype=torch.int32), - b_shared_seq_len=torch.arange(10, 16, dtype=torch.int32), - b_shared_radix_node_id=torch.arange(20, 26, dtype=torch.int64), - is_prefill=False, - multimodal_params=_empty_multimodal_params(batch_size), - ) - - -def _make_empty_decode_input() -> ModelInput: - return ModelInput( - batch_size=0, - total_token_num=0, - max_q_seq_len=1, - max_kv_seq_len=0, - input_ids=torch.empty((0,), dtype=torch.int64), - b_req_idx=torch.empty((0,), dtype=torch.int32), - b_mtp_index=torch.empty((0,), dtype=torch.int32), - b_seq_len=torch.empty((0,), dtype=torch.int32), - b_position_delta=torch.empty((0,), dtype=torch.int32), - b_shared_seq_len=torch.empty((0,), dtype=torch.int32), - b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), - is_prefill=False, - multimodal_params=[], - ) - - def test_base_model_prefill_accepts_two_prebuilt_inputs(monkeypatch): model_input0 = _make_prefill_input([10, 11], req_idx=20, is_decode_req=False) model_input1 = _make_prefill_input([12], req_idx=21, is_decode_req=True) @@ -184,7 +122,6 @@ def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) model.graph = None model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(77,), page_size=1) infer_batch_sizes = [] def fake_create_inferstate(model_input, microbatch_index): @@ -207,3 +144,65 @@ def fake_create_inferstate(model_input, microbatch_index): assert infer_batch_sizes == [2, 2] assert output0.logits.shape == (1, 1) assert output1.logits.shape == (0, 1) + + +def _empty_multimodal_params(batch_size: int): + return [{"images": [], "audios": []} for _ in range(batch_size)] + + +def _make_prefill_input(input_ids: list[int], req_idx: int, is_decode_req: bool) -> ModelInput: + token_num = len(input_ids) + return ModelInput( + batch_size=1, + total_token_num=token_num, + max_q_seq_len=token_num, + max_kv_seq_len=token_num, + max_cache_len=0, + input_ids=torch.tensor(input_ids, dtype=torch.int64), + b_req_idx=torch.tensor([req_idx], dtype=torch.int32), + b_mtp_index=torch.zeros(1, dtype=torch.int32), + b_seq_len=torch.tensor([token_num], dtype=torch.int32), + b_is_decode_req=torch.tensor([is_decode_req]), + b_ready_cache_len=torch.zeros(1, dtype=torch.int32), + b_prefill_start_loc=torch.zeros(1, dtype=torch.int32), + is_prefill=True, + b_prefill_has_output_cpu=[True], + multimodal_params=_empty_multimodal_params(1), + ) + + +def _make_decode_input() -> ModelInput: + batch_size = 6 + return ModelInput( + batch_size=batch_size, + total_token_num=39, + max_q_seq_len=1, + max_kv_seq_len=9, + input_ids=torch.arange(batch_size, dtype=torch.int64), + b_req_idx=torch.tensor([10, 11, 12, 12, 12, 12], dtype=torch.int32), + b_mtp_index=torch.tensor([0, 0, 0, 1, 2, 3], dtype=torch.int32), + b_seq_len=torch.tensor([4, 5, 6, 7, 8, 9], dtype=torch.int32), + b_position_delta=torch.arange(batch_size, dtype=torch.int32), + b_shared_seq_len=torch.arange(10, 16, dtype=torch.int32), + b_shared_radix_node_id=torch.arange(20, 26, dtype=torch.int64), + is_prefill=False, + multimodal_params=_empty_multimodal_params(batch_size), + ) + + +def _make_empty_decode_input() -> ModelInput: + return ModelInput( + batch_size=0, + total_token_num=0, + max_q_seq_len=1, + max_kv_seq_len=0, + input_ids=torch.empty((0,), dtype=torch.int64), + b_req_idx=torch.empty((0,), dtype=torch.int32), + b_mtp_index=torch.empty((0,), dtype=torch.int32), + b_seq_len=torch.empty((0,), dtype=torch.int32), + b_position_delta=torch.empty((0,), dtype=torch.int32), + b_shared_seq_len=torch.empty((0,), dtype=torch.int32), + b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), + is_prefill=False, + multimodal_params=[], + ) diff --git a/unit_tests/common/test_paged_mem_manager.py b/unit_tests/common/test_paged_mem_manager.py deleted file mode 100644 index 60a7c86337..0000000000 --- a/unit_tests/common/test_paged_mem_manager.py +++ /dev/null @@ -1,59 +0,0 @@ -from types import SimpleNamespace -import pickle - -import pytest -import torch - -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_shared_memory_exports_preserve_variable_pickle_lengths(monkeypatch, tmp_path): - from lightllm.common.kv_cache_mem_manager import mem_manager as module - - monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", f"{tmp_path.parent.name}_{tmp_path.name}") - monkeypatch.setattr(module, "get_node_world_size", lambda: 4) - monkeypatch.setattr(module, "get_current_rank_in_node", lambda: 0) - monkeypatch.setattr(module, "kv_trans_use_p2p", lambda: False) - # CUDA IPC ref-counter offsets can cross pickle's one-byte integer boundary - # between exports. Each consumer still needs its own reference-counter token. - offsets = [254, 255, 256, 257, 258, 259, 260, 261] - payloads = [pickle.dumps({"ref_counter_offset": offset}) for offset in offsets] - serialized = iter(payloads) - monkeypatch.setattr(module.ForkingPickler, "dumps", lambda obj: memoryview(next(serialized))) - manager = MemoryManager.__new__(MemoryManager) - try: - manager.write_to_shm(SimpleNamespace(req_to_token_indexs=None)) - for offset in reversed(offsets): - assert manager.loads_from_shm(0) == {"ref_counter_offset": offset} - with pytest.raises(AssertionError): - manager.loads_from_shm(0) - finally: - shm = module.create_or_link_shm("mem_manager_0", -1, force_mode="link") - shm.close() - shm.unlink() - - -def test_kv_cache_allocation_requires_complete_pages(): - manager = MemoryManager.__new__(MemoryManager) - manager.page_size = 4 - manager.allocator = SimpleNamespace(alloc=lambda size: torch.arange(size, dtype=torch.int32)) - - with pytest.raises(AssertionError, match="must be a multiple of page_size 4"): - manager.alloc(5) - - assert manager.alloc(8).tolist() == list(range(8)) - - -@pytest.mark.parametrize( - "manager_class", - [MemoryManager, Deepseek2MemoryManager, PPLINT4KVMemoryManager, PPLINT8KVMemoryManager], -) -def test_kv_cache_buffer_size_requires_complete_pages(manager_class): - manager = manager_class.__new__(manager_class) - manager.page_size = 4 - - with pytest.raises(AssertionError, match="KV cache size 5 must be a multiple of page_size 4"): - manager._init_buffers(size=5, dtype=torch.float16, head_num=1, head_dim=8, layer_num=1) diff --git a/unit_tests/common/test_quantized_paged_mem_manager.py b/unit_tests/common/test_quantized_paged_mem_manager.py deleted file mode 100644 index 84fe266613..0000000000 --- a/unit_tests/common/test_quantized_paged_mem_manager.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_quantized_kv_buffers_include_the_full_hold_page(monkeypatch): - allocations = [] - - def fake_empty(shape, **kwargs): - allocations.append(shape) - return SimpleNamespace(shape=shape) - - monkeypatch.setattr(torch, "empty", fake_empty) - - int8_manager = PPLINT8KVMemoryManager.__new__(PPLINT8KVMemoryManager) - int8_manager.page_size = 4 - int8_manager.group_quant_size = 8 - int8_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - int4_manager = PPLINT4KVMemoryManager.__new__(PPLINT4KVMemoryManager) - int4_manager.page_size = 4 - int4_manager.group_quant_size = 8 - int4_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - assert allocations == [ - (3, 36, 4, 16), - (3, 36, 4, 2), - (3, 36, 4, 8), - (3, 36, 4, 2), - ] diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 224c308f28..a834b8efff 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -10,7 +10,7 @@ from lightllm.server.router.model_infer.mode_backend import base_backend -class _FakeMemManager: +class _FakeAllocator: page_size = 4 def __init__(self): @@ -24,16 +24,15 @@ def alloc(self, size): return result -def test_bind_mem_manager_initializes_hold_request_with_reserved_page(): +def test_bind_hold_page_initializes_only_padding_request(): req_manager = ReqManager.__new__(ReqManager) req_manager.HOLD_REQUEST_ID = 2 req_manager.req_list = SimpleNamespace(is_all_free=lambda: True) req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) - mem_manager = SimpleNamespace(page_size=4, HOLD_TOKEN_MEMINDEXES=(100, 101, 102, 103)) + req_manager.page_size = 4 + req_manager.bind_hold_page((100, 101, 102, 103)) - req_manager.bind_mem_manager(mem_manager) - - assert req_manager.mem_manager is mem_manager + assert not hasattr(req_manager, "mem_manager") assert req_manager.req_to_token_indexs[2].tolist() == [100, 101, 102, 103] * 2 assert req_manager.req_to_token_indexs[:2].eq(-1).all() @@ -43,31 +42,7 @@ def test_hold_request_indexes_require_all_requests_released(): req_manager.req_list = SimpleNamespace(is_all_free=lambda: False) with pytest.raises(AssertionError, match="all requests are released"): - req_manager.init_hold_request_indexs() - - -def _make_context(monkeypatch): - mem_manager = _FakeMemManager() - req_manager = SimpleNamespace( - mem_manager=mem_manager, - req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32), - ) - context = SimpleNamespace( - args=SimpleNamespace(page_size=4), - req_manager=req_manager, - radix_cache=None, - ) - monkeypatch.setattr(generic_pre_process, "g_infer_context", context) - monkeypatch.setattr(base_backend, "g_infer_context", context) - backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - backend.args = context.args - return context, backend - - -def _make_req(req_idx): - req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req + req_manager.bind_hold_page((100, 101, 102, 103)) def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypatch): @@ -77,17 +52,17 @@ def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypat mem_indexes = backend._alloc_req_kv_mem(req, alloc_token_num=4) assert mem_indexes.tolist() == [0, 1, 2, 3] assert req.hold_kv_len == 4 - assert context.req_manager.mem_manager.alloc_sizes == [4] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4] assert context.req_manager.req_to_token_indexs[0, :4].tolist() == [0, 1, 2, 3] req.cur_kv_len = 3 assert backend._alloc_req_kv_mem(req, alloc_token_num=0) is None - assert context.req_manager.mem_manager.alloc_sizes == [4] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4] req.cur_kv_len = 4 backend._alloc_req_kv_mem(req, alloc_token_num=4) assert req.hold_kv_len == 8 - assert context.req_manager.mem_manager.alloc_sizes == [4, 4] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4, 4] assert context.req_manager.req_to_token_indexs[0, :8].tolist() == list(range(8)) @@ -121,8 +96,8 @@ def __getitem__(self, key): target = _CopyTarget() context = SimpleNamespace( + att_resource_manager=SimpleNamespace(full_att_allocator=_FakeAllocator()), req_manager=SimpleNamespace( - mem_manager=_FakeMemManager(), req_to_token_indexs=_ReqToTokenIndexes(target), ), radix_cache=None, @@ -209,7 +184,7 @@ def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( context.req_manager.req_to_token_indexs[req.req_idx, :hold_kv_len] = torch.arange( start_index, start_index + hold_kv_len, dtype=torch.int32 ) - context.req_manager.mem_manager.next_index = 2 * hold_kv_len + context.att_resource_manager.full_att_allocator.next_index = 2 * hold_kv_len backend._filter_not_ready_reqs = lambda req_ids: reqs copy_modes = [] original_alloc_req_kv_mem = backend._alloc_req_kv_mem @@ -229,7 +204,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert prefill_reqs == [reqs[0]] assert decode_reqs == [] assert prefill_tokens == 3 - assert context.req_manager.mem_manager.alloc_sizes == expected_alloc_sizes + assert context.att_resource_manager.full_att_allocator.alloc_sizes == expected_alloc_sizes assert reqs[0].hold_kv_len == 8 assert not reqs[0].wait_pause assert reqs[1].hold_kv_len == hold_kv_len @@ -265,7 +240,7 @@ def test_decode_scheduler_uses_non_blocking_request_table_copy(monkeypatch): req.get_cur_total_len = lambda: 5 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 + context.att_resource_manager.full_att_allocator.next_index = 4 backend._filter_not_ready_reqs = lambda req_ids: [req] copy_modes = [] @@ -287,7 +262,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert decode_reqs == [req] assert prefill_tokens == 0 assert req.hold_kv_len == 16 - assert context.req_manager.mem_manager.alloc_sizes == [12] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [12] assert copy_modes == [True] @@ -327,7 +302,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.get_cur_total_len = lambda: req.cur_kv_len + 1 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :page_size] = torch.arange(page_size, dtype=torch.int32) - context.req_manager.mem_manager.next_index = page_size + context.att_resource_manager.full_att_allocator.next_index = page_size backend._filter_not_ready_reqs = lambda req_ids: [req] for _ in range(expected_alloc_token_num): @@ -336,7 +311,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.cur_kv_len += 1 assert req.hold_kv_len == page_size + expected_alloc_token_num - assert context.req_manager.mem_manager.alloc_sizes == [expected_alloc_token_num] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [expected_alloc_token_num] def test_decode_reserves_mtp_headroom(monkeypatch): @@ -352,7 +327,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): req.args = context.args context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 + context.att_resource_manager.full_att_allocator.next_index = 4 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == 6 assert alloc_token_num == 8 @@ -363,7 +338,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): assert run_reqs == [req, req, req] assert model_input.b_seq_len.tolist() == [4, 5, 6] assert req.hold_kv_len == 12 - assert context.req_manager.mem_manager.alloc_sizes == [8] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [8] assert context.req_manager.req_to_token_indexs[0, :12].tolist() == list(range(12)) @@ -380,7 +355,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): req.get_radix_cache_shared_len = lambda: 0 req.args = context.args context.req_manager.req_to_token_indexs[0, :3] = torch.arange(3, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 3 + context.att_resource_manager.full_att_allocator.next_index = 3 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == alloc_token_num == 6 backend._alloc_req_kv_mem(req, alloc_token_num) @@ -388,7 +363,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): model_input, _ = generic_pre_process.prepare_decode_inputs([req]) assert req.hold_kv_len == 9 - assert context.req_manager.mem_manager.alloc_sizes == [6] + assert context.att_resource_manager.full_att_allocator.alloc_sizes == [6] assert context.req_manager.req_to_token_indexs[0, :9].tolist() == list(range(9)) @@ -457,35 +432,6 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa assert free_token_indexes[0].tolist() == list(range(8)) -def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv_len=None): - if chunk_kv_len is None: - chunk_kv_len = target_kv_len - req = SimpleNamespace( - args=SimpleNamespace(page_size=page_size), - req_id=req_idx, - req_idx=req_idx, - cur_kv_len=0, - hold_kv_len=0, - paused=True, - match_call_count=0, - shared_kv_node=None, - shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), - get_cur_total_len=lambda: target_kv_len, - get_chuncked_input_token_len=lambda: chunk_kv_len, - ) - - def match_radix_cache(): - req.match_call_count += 1 - req.cur_kv_len = match_kv_len - req.hold_kv_len = match_kv_len - req.shm_req.shm_cur_kv_len = match_kv_len - - req._match_radix_cache = match_radix_cache - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) - return req - - def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) infer_context.args = SimpleNamespace(page_size=4) @@ -494,7 +440,6 @@ def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): freed_indexes = [] infer_context.req_manager = SimpleNamespace( req_to_token_indexs=torch.arange(16, dtype=torch.int32).reshape(2, 8), - free_token=lambda indexes: freed_indexes.extend(indexes.tolist()), ) infer_context.get_can_alloc_token_num = lambda: 5 large_req = _make_paused_req(req_idx=0, target_kv_len=5, page_size=4) @@ -519,7 +464,6 @@ def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch infer_context.radix_cache = None infer_context.req_manager = SimpleNamespace( req_to_token_indexs=torch.arange(8, dtype=torch.int32)[None, :], - free_token=lambda indexes: None, ) infer_context.get_can_alloc_token_num = lambda: 4 req = _make_paused_req(req_idx=0, target_kv_len=8, page_size=4, chunk_kv_len=4) @@ -547,3 +491,56 @@ def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): assert req.paused is False assert req.shm_req.is_paused is False assert req.match_call_count == 1 + + +def _make_context(monkeypatch): + mem_manager = _FakeAllocator() + req_manager = SimpleNamespace( + req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32), + ) + context = SimpleNamespace( + args=SimpleNamespace(page_size=4), + req_manager=req_manager, + att_resource_manager=SimpleNamespace(full_att_allocator=mem_manager), + radix_cache=None, + ) + monkeypatch.setattr(generic_pre_process, "g_infer_context", context) + monkeypatch.setattr(base_backend, "g_infer_context", context) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + backend.args = context.args + return context, backend + + +def _make_req(req_idx): + req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0) + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + return req + + +def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv_len=None): + if chunk_kv_len is None: + chunk_kv_len = target_kv_len + req = SimpleNamespace( + args=SimpleNamespace(page_size=page_size), + req_id=req_idx, + req_idx=req_idx, + cur_kv_len=0, + hold_kv_len=0, + paused=True, + match_call_count=0, + shared_kv_node=None, + shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), + get_cur_total_len=lambda: target_kv_len, + get_chuncked_input_token_len=lambda: chunk_kv_len, + ) + + def match_radix_cache(): + req.match_call_count += 1 + req.cur_kv_len = match_kv_len + req.hold_kv_len = match_kv_len + req.shm_req.shm_cur_kv_len = match_kv_len + + req._match_radix_cache = match_radix_cache + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) + return req diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py new file mode 100644 index 0000000000..f9293fd1df --- /dev/null +++ b/unit_tests/conftest.py @@ -0,0 +1,40 @@ +"""Shared fixtures for tests that construct runtime resources.""" + +from __future__ import annotations + + +from contextlib import suppress +from uuid import uuid4 + +import pytest + + +@pytest.fixture +def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): + """Provide a service identity and reclaim the allocator's real shared counters.""" + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.att_resources.allocators import full_att_kv_allocator + from lightllm.utils.envs_utils import get_unique_server_name + + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_kv_allocator_" + uuid4().hex) + get_unique_server_name.cache_clear() + counters = [] + shared_int = full_att_kv_allocator.SharedInt + + def track_counter(name: str): + counter = shared_int(name) + counters.append(counter) + return counter + + monkeypatch.setattr(full_att_kv_allocator, "SharedInt", track_counter) + yield + AttResourceManager._get_states.cache_clear() + names = set() + for counter in counters: + if counter.shm.name not in names: + with suppress(FileNotFoundError): + counter.shm.unlink() + names.add(counter.shm.name) + counter.shm.close() + get_unique_server_name.cache_clear() diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 2d6819e98d..e1d9131cd4 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -4,66 +4,55 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.common.att_resources.allocators import FullAttKVAllocator + +from lightllm.common.att_resources.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, +) +from lightllm.common.req_manager import ReqManager +from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager +from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture -def model_config(): - return { - "kv_lora_rank": 512, - "num_hidden_layers": 4, - "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], - "index_kpool": 4, - "index_head_dim": 128, - "linear_attn_config": { - "num_heads": 8, - "head_dim": 128, - "short_conv_kernel_size": 4, - "kda_layers": [0, 1, 2], +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +def test_nope_cache_config_uses_native_mla_width(): + config = Glm5NextResourceConfig.from_model_config( + { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], + }, }, - } - - -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -@pytest.mark.parametrize("page_size", [1, 16]) -def test_nope_cache_layout_and_hold_page(monkeypatch, model_config, kv_type, page_size): - args = StartArgs(tp=4, data_type="bfloat16", llm_kv_type=kv_type) - config = Glm5NextCacheConfig.from_model_config(model_config, args) - nested_config = Glm5NextCacheConfig.from_model_config({"text_config": model_config}, args) - assert nested_config == config - # The allocator reserves an entire HOLD page, including when KV is quantized. - monkeypatch.setattr(Glm5NextMemManager, "_init_linear_att_buffers", lambda self: None) - mem_manager_class = FP8Glm5NextMemManager if kv_type == "fp8kv_dsa" else Glm5NextMemManager - mem = object.__new__(mem_manager_class) - mem.page_size = page_size - mem._init_buffers(32, config.full_att_dtype, 1, config.full_att_head_dim, 1) - assert mem.kv_buffer.shape[1] == 32 + page_size - - assert config.full_att_head_dim == (720 if kv_type == "fp8kv_dsa" else 584) - assert config.full_att_dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else torch.bfloat16) - assert config.conv_state_dtype == torch.bfloat16 + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.full_att_head_dim == 584 + assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, kv_type): - from lightllm.models.glm5_next.model import Glm5NextTpPartModel - +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) args = StartArgs( tp=tp_world_size, - llm_kv_type=kv_type, data_type="bfloat16", linear_att_hash_page_size=4, linear_att_page_block_num=2, @@ -74,14 +63,12 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - use_fp8_kv = kv_type == "fp8kv_dsa" - dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) - config = Glm5NextCacheConfig( + config = Glm5NextResourceConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, + full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=head_dim, + full_att_head_dim=584, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, num_linear_k_heads=2, @@ -89,132 +76,95 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, + linear_layer_num=3, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=("linear_attention",) * 3 + ("deepseek_sparse_attention",), + full_attention_interval=4, + all_layer_num=4, draft_full_att_kv_layer_num=int(mtp_step > 0), ) - full_layers = config.get_full_att_kv_layer_num_with_draft_model() - model = object.__new__(Glm5NextTpPartModel) - model.args = args - model.max_total_token_num = 16 - model.linear_config = config - model.mem_fraction = 0.9 - model._init_mem_manager() - mem = model.mem_manager - req = Glm5NextReqManager(3, 16, mem, config) - att_kv = mem.get_att_input_params(3) - assert att_kv.shape == (17, 1, 584 if use_fp8_kv else 512) - assert att_kv.stride(0) == head_dim - index_bytes = mem.get_indexer_k_buffer(3) + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + config.get_full_att_layer_ids(), + 1, + 584, + torch.bfloat16, + tp_world_size=tp_world_size, + kv_dim=512, + ), + page_size=args.page_size, + ) + full.initialize_storage(16) + req = ReqManager(3, 16) + resources = Glm5NextAttResourceManager(config, 4, mtp_step) + from unit_tests.common.att_resources.model_helpers import create_full_attention_roles + + main_full, draft_full = create_full_attention_roles(full, config.all_layer_num) + resources.bind_storage( + main_full, + FullAttKVAllocator(full.capacity, full.page_size), + draft_full_state=draft_full, + big_page_capacity=4, + small_page_capacity=2, + ) + full = main_full + linear = resources.main_kda_state + tail = resources.main_indexer_kpool_c4_tail_state + assert not hasattr(req, "req_to_conv_state") + att_kv = full.get_layer_view(3).kv + assert att_kv.shape == (17, 1, 512) + assert att_kv.stride(0) == 584 + index_bytes = full.get_layer_view(3).indexer_k index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") destinations = torch.tensor([5, 9], dtype=torch.int32, device="cuda") - mem.operator.copy_kv_to_mem_manager(3, destinations, new_kv) - if use_fp8_kv: - values = att_kv[..., :448].view(torch.float8_e4m3fn).float().view(17, 1, 7, 64) - scales = att_kv[..., 576:583].view(torch.float8_e8m0fnu).float().unsqueeze(-1) - tail = att_kv[..., 448:576].view(torch.bfloat16).float() - actual = torch.cat(((values * scales).view(17, 1, 448), tail), dim=-1)[destinations] - torch.testing.assert_close(actual, new_kv.float(), atol=0.13, rtol=0.065) - assert torch.equal(tail[destinations], new_kv[..., 448:].float()) - else: - assert torch.equal(att_kv[destinations], new_kv) + full.write_full_att_kv(3, destinations, new_kv) + assert torch.equal(att_kv[destinations], new_kv) assert torch.equal(index_bytes, expected_index_bytes) - assert req.get_indexer_tail_buffer(3).shape == (4, 4 + mtp_step, 256) + assert tail.get_layer_view(3).shape == (4, 4 + mtp_step, 256) if mtp_step: - assert req.get_indexer_tail_buffer(4).shape == (4, 4 + mtp_step, 256) - assert mem.get_att_input_params(4).shape == att_kv.shape - cache = req.create_small_page_cache_manager(2) if small_page else mem.big_page_buffers - slot = cache.alloc_one_state_cache() - source_req = SimpleNamespace(req_idx=0) - req.init_hybrid_attention_state(source_req) - req.req_to_conv_state.buffer[:, 0].normal_() - req.req_to_ssm_state.buffer[:, 0].normal_() - conv = req.req_to_conv_state.buffer[:, 0, ..., :3].clone() - ssm = req.req_to_ssm_state.buffer[:, 0].clone() - if small_page: - req.save_state(source_req, slot, cache) - else: - req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [source_req], [slot]) + assert resources.draft0_indexer_kpool_c4_tail_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) + assert draft_full.get_layer_view(4).kv.shape == att_kv.shape + kind = "small" if small_page else "big" + allocator = resources.bounded_att_small_page_allocator if small_page else resources.bounded_att_big_page_allocator + slot = allocator.alloc(1)[0] + resources.reset_bounded_att_state([0]) + resources.reset_attention_tail_state([0]) + linear.conv_state[:, 0].normal_() + linear.ssm_state[:, 0].normal_() + conv = linear.conv_state[:, 0, ..., :3].clone() + ssm = linear.ssm_state[:, 0].clone() + resources.save_bounded_att_pages([0], page_kind=kind, slot_ids=[slot]) torch.cuda.synchronize() - req.req_to_conv_state.buffer[:, 0].zero_() - req.req_to_ssm_state.buffer[:, 0].zero_() - req.req_to_indexer_tail.buffer.normal_() - unchanged = req.req_to_indexer_tail.buffer[:, [0, 1, 3]].clone() - dest_req = SimpleNamespace(req_idx=2, shared_kv_node=SimpleNamespace(small_page_buffer_idx=slot)) - if small_page: - req.restore_small_page_state(dest_req) - else: - req.restore_big_page_state(slot, dest_req) + linear.conv_state[:, 0].zero_() + linear.ssm_state[:, 0].zero_() + tail.tail_buffer.normal_() + unchanged = tail.tail_buffer[:, [0, 1, 3]].clone() + resources.restore_bounded_att_pages([2], page_kind=kind, slot_ids=[slot]) + resources.reset_attention_tail_state([2]) torch.cuda.synchronize() - assert torch.equal(req.req_to_conv_state.buffer[:, 2, ..., :3], conv) - assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1)], ssm) + assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) + assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) # Both page sizes are aligned to complete pools. Restoring a prefix must # discard stale tail values from a previously allocated request slot. - assert not req.req_to_indexer_tail.buffer[:, 2].any() - assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) - req.req_to_indexer_tail.buffer[:, 2].normal_() - req.init_hybrid_attention_state(dest_req) - assert not req.req_to_conv_state.buffer[:, 2].any() - assert not req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() - assert not req.req_to_indexer_tail.buffer[:, 2].any() - assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) + assert not tail.tail_buffer[:, 2].any() + assert torch.equal(tail.tail_buffer[:, [0, 1, 3]], unchanged) + tail.tail_buffer[:, 2].normal_() + preserved_tail = tail.tail_buffer.clone() + resources.reset_bounded_att_state([2]) + assert not linear.conv_state[:, 2].any() + assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert torch.equal(tail.tail_buffer, preserved_tail) + resources.reset_attention_tail_state([2]) + assert not tail.tail_buffer[:, 2].any() + assert torch.equal(tail.tail_buffer[:, [0, 1, 3]], unchanged) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only # belong to live requests. Bytewise equality checks FP8 scale integrity. - packed_bytes = mem.kv_buffer.view(torch.uint8) + packed_bytes = full.storage.view(torch.uint8) packed_bytes[:, 0].random_(0, 256) - mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) + full.copy_full_att_tokens(torch.tensor([0]), torch.tensor([7])) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == head_dim * dtype.itemsize * full_layers - assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_kv_buffer_to_cpu_cache, - copy_cpu_cache_to_kv_buffer, - ) - - cpu_pages = torch.zeros((1, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8, pin_memory=True) - indexes = torch.arange(8, dtype=torch.int32, device="cuda") - zero = torch.zeros(1, dtype=torch.int64, device="cuda") - ready = torch.zeros(1, dtype=torch.int32, pin_memory=True) - big = mem.big_page_buffers - packed_bytes[:, :8].random_(0, 256) - big.conv_state_cache.buffer[0].normal_() - big.ssm_state_cache.buffer[0].normal_() - expected_kv = packed_bytes[:, :8].clone() - expected_conv = big.conv_state_cache.buffer[0].clone() - expected_ssm = big.ssm_state_cache.buffer[0].clone() - common = dict( - mem_indexes=indexes, - page_indexes=zero, - big_page_buffer_ids=zero, - cpu_kv_conv_state=big.conv_state_cache.buffer, - cpu_kv_ssm_state=big.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_pages, - tp_world_size=tp_world_size, - big_page_token_num=8, - linear_config=config, - ) - # Simulate TP writers to a shared CPU page: MLA/index KV is replicated, - # while each rank must retain its own KDA checkpoint region. - expected_states = [] - for rank in range(tp_world_size): - big.conv_state_cache.buffer[0].copy_(expected_conv + rank) - big.ssm_state_cache.buffer[0].copy_(expected_ssm + rank) - expected_states.append((big.conv_state_cache.buffer[0].clone(), big.ssm_state_cache.buffer[0].clone())) - copy_kv_buffer_to_cpu_cache(page_readies=ready, gpu_kv_full_att_state=mem.kv_buffer, tp_rank=rank, **common) - torch.cuda.synchronize() - for rank, (conv, ssm) in enumerate(expected_states): - packed_bytes[:, :8].zero_() - big.conv_state_cache.buffer[0].zero_() - big.ssm_state_cache.buffer[0].zero_() - copy_cpu_cache_to_kv_buffer(gpu_full_att_kv_state=mem.kv_buffer, tp_rank=rank, **common) - torch.cuda.synchronize() - assert torch.equal(packed_bytes[:, :8], expected_kv) - assert torch.equal(big.conv_state_cache.buffer[0], conv) - assert torch.equal(big.ssm_state_cache.buffer[0], ssm) @pytest.mark.parametrize("gate_type", ["silu", "sigmoid"]) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py index fce877448c..76646c5224 100644 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -12,7 +12,7 @@ @pytest.fixture def indexer(): return Glm5NextNsaInfer( - 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5, "index_kpool": 4}, 1 ) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 37485746ba..db1186b792 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -35,16 +35,6 @@ def setup(): triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) -def _reference_kda(q, k, v, gate, beta, a, bias, state): - q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) / q.shape[-1] ** 0.5 - k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) - decay = (-5 * torch.sigmoid(a.exp()[:, None] * (gate.float() + bias))).exp() - state = state * decay[..., None] - delta = (v.float() - torch.einsum("hkv,hk->hv", state, k)) * beta.float().sigmoid()[:, None] - state = state + k[..., None] * delta[:, None, :] - return torch.einsum("hkv,hk->hv", state, q), state - - @pytest.mark.parametrize("tokens", [1, 3, 65, 129]) def test_kda_chunk_and_decode_match_recurrence(tokens): heads, dim = 2, 128 @@ -252,7 +242,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -309,7 +299,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -401,7 +391,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2 ** 31 // (584 * 2) + 1 + loc = 2**31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -414,7 +404,7 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @@ -491,7 +481,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) out, final = chunk_gated_delta_rule( q, k, @@ -594,18 +584,22 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): storage = torch.zeros(tokens + 9, 1, 584, device="cuda", dtype=torch.bfloat16) tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) ragged = torch.randperm(tokens + 9, device="cuda", dtype=torch.int32)[:tokens] - manager = SimpleNamespace( - get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:], + full_att_state = SimpleNamespace( + get_layer_view=lambda _: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), ) infer = SimpleNamespace( - mem_manager=manager, + att_resource_manager=SimpleNamespace( + main_indexer_kpool_c4_tail_state=SimpleNamespace(get_layer_view=lambda _: tail), + main_mla_packed_bf16_state=full_att_state, + ), + is_draft_model=False, is_prefill=True, b_mtp_index=torch.zeros(1, device="cuda", dtype=torch.int32), mem_index=ragged, # A large capacity also exercises query chunking and uninitialized # logits beyond each query's actual pool range. max_kv_seq_len=max_kv_seq_len, - req_manager=SimpleNamespace(req_to_token_indexs=ragged[None], get_indexer_tail_buffer=lambda _: tail), + req_manager=SimpleNamespace(req_to_token_indexs=ragged[None]), b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), b_seq_len=torch.tensor([tokens], device="cuda", dtype=torch.int32), b1_cu_q_seq_len=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), @@ -617,8 +611,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): ragged_mem_index=ragged, ) indexer = Glm5NextNsaInfer( - 0, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5}, 1 + 0, + {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5, "index_kpool": 4}, + 1, ) + indexer.full_att_state_name = "main_mla_packed_bf16_state" _, full = indexer._get_indices(hidden, hidden, infer, state, weights) assert full[-1, 2048].item() == tokens - 1 # Always-selected incomplete tail. assert full[0, 0].item() == 0 and (full[0, 1:] == -1).all() @@ -646,10 +643,20 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) + scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups + + +def _reference_kda(q, k, v, gate, beta, a, bias, state): + q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) / q.shape[-1] ** 0.5 + k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) + decay = (-5 * torch.sigmoid(a.exp()[:, None] * (gate.float() + bias))).exp() + state = state * decay[..., None] + delta = (v.float() - torch.einsum("hkv,hk->hv", state, k)) * beta.float().sigmoid()[:, None] + state = state + k[..., None] * delta[:, None, :] + return torch.einsum("hkv,hk->hv", state, q), state diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 08f8579174..33a7ff73a4 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -61,7 +61,7 @@ def config(): "index_head_dim": 128, "index_topk": 2048, "index_kpool": 4, - "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4}, + "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4, "kda_layers": [0, 1, 2]}, "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], "mhc": False, "n_routed_experts": 4, @@ -170,7 +170,12 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) state = InferStateInfo() state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - state.mem_manager = SimpleNamespace(get_att_input_params=lambda layer_index: kv) + layer.full_att_state_name = "main_mla_packed_bf16_state" + state.att_resource_manager = SimpleNamespace( + main_mla_packed_bf16_state=SimpleNamespace( + get_layer_view=lambda layer_id: SimpleNamespace(kv=kv) + ) + ) state.prefill_att_state = Glm5NextSparsePrefillState() state.b1_cu_q_seq_len = torch.arange(4, device="cuda", dtype=torch.int32) state.max_q_seq_len = 1 @@ -375,21 +380,23 @@ def test_glm_rejects_tpsp(config, tp_world_size): model._verify_params() -def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): +def test_native_drafts_share_weights_and_prepare_their_own_state(config): main = object.__new__(Glm5NextTpPartModel) main.config = dict(config, mhc=True) main.tp_world_size_ = 1 - main.args, main.load_way = StartArgs(), "HF" + main.args, main.load_way = StartArgs(data_type="bfloat16", mtp_mode="vanilla_with_att", mtp_step=2), "HF" main._verify_params() main._init_some_value() - main.layers_infer = [object() for _ in range(config["num_hidden_layers"])] + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(config["num_hidden_layers"])] main.pre_post_weight = SimpleNamespace(wte_weight_=object(), lm_head_weight_=object(), final_norm_weight_=object()) - main.req_manager, main.mem_manager, main.linear_config = object(), object(), object() + main.req_manager = object() + main.max_req_num, main.tp_rank_ = 0, 0 drafts = [] for step in range(2): draft = object.__new__(Glm5NextMTPModel) draft.main_model, draft.mtp_previous_draft_models = main, list(drafts) - draft.tp_world_size_, draft.data_type = 1, torch.bfloat16 + draft.tp_world_size_, draft.tp_rank_, draft.data_type = 1, 0, torch.bfloat16 + draft.args = main.args draft.load_way = "HF" draft._init_config() draft._verify_params() @@ -398,8 +405,7 @@ def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): draft._init_infer_layer() draft._init_some_value() draft._init_req_manager() - draft._init_mem_manager() - draft._init_att_backend1() + draft.max_req_num = main.max_req_num assert main.config["mhc"] and not draft.config["mhc"] assert draft.config["layer_types"] is not main.config["layer_types"] assert draft.layers_num == 1 and draft.head_dim_ == 512 @@ -408,10 +414,30 @@ def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): assert draft.pre_post_weight.wte_weight_ is main.pre_post_weight.wte_weight_ assert draft.pre_post_weight.lm_head_weight_ is main.pre_post_weight.lm_head_weight_ assert draft.pre_post_weight.main_norm_weight_ is main.pre_post_weight.final_norm_weight_ - assert draft.req_manager is main.req_manager and draft.mem_manager is main.mem_manager - assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + assert draft.req_manager is main.req_manager drafts.append(draft) + assert all(not hasattr(model, "att_resource_manager") for model in (main, *drafts)) + assert all(not hasattr(model, "linear_config") for model in (main, *drafts)) + from contextlib import nullcontext + for model in (main, *drafts): + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + main.prepare_att_resources() + assert main.att_resource_manager.main_kda_state.conv_state.is_cuda + assert main.att_resource_manager.main_indexer_kpool_c4_tail_state.tail_buffer.is_cuda + for step, draft in enumerate(drafts): + draft.prepare_att_resources() + assert draft.att_resource_manager is main.att_resource_manager + draft._init_att_backend1() + state = getattr(draft.att_resource_manager, f"draft{step}_mla_packed_bf16_state") + assert state is not main.att_resource_manager.main_mla_packed_bf16_state + assert not hasattr(draft, "full_att_state") + assert state.storage.device.type == "meta" + assert state.config.layer_ids == (4 + step,) + assert state.resource_id == f"draft{step}_mla_packed_bf16_state" + assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + assert getattr(draft.att_resource_manager, f"draft{step}_indexer_kpool_c4_tail_state").tail_buffer.is_cuda + # NoPE state initialization must not require a model-owned RoPE cache. for model in (main, *drafts): state = model.infer_state_class() diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 00d4e326f0..9ccab62b6a 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -191,9 +191,11 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): uses_dynamic_spec_verify_layout=lambda: dynamic, split_qkv=lambda x: x.split(hidden, -1), ) - manager = SimpleNamespace( - req_to_mtp_state_index=accepted_offsets, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm) + backend.linear_state = SimpleNamespace( + mtp_state_index=accepted_offsets, + get_layer_view=lambda layer: (conv, ssm), ) + manager = SimpleNamespace(HOLD_REQUEST_ID=3) infer = InferStateInfo() infer.batch_size = tokens infer.b_req_idx = torch.tensor(request_rows, device="cuda", dtype=torch.int32) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index fab579aae1..067643b514 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,30 +1,41 @@ +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + import dataclasses +from types import SimpleNamespace import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager, MemoryManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles + +from lightllm.common.att_resources.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, +) +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager +from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] -@pytest.fixture(params=[(0, False), (2, False), (0, True), (2, True)]) +@pytest.fixture(params=[0, 2]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step, use_fp8_kv = request.param - dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) - mem_manager_class = FP8Glm5NextMemManager if use_fp8_kv else Glm5NextMemManager + mtp_step = request.param args = StartArgs( data_type="bfloat16", - llm_kv_type="fp8kv_dsa" if use_fp8_kv else "None", + pd_kv_page_size=2048, linear_att_hash_page_size=4, linear_att_page_block_num=2, mtp_step=mtp_step, @@ -34,21 +45,14 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def check_shm_refs(mem, req): - assert mem.req_to_indexer_tail is req.req_to_indexer_tail - assert mem.req_to_conv_state is req.req_to_conv_state - assert mem.req_to_ssm_state is req.req_to_ssm_state - assert mem.big_page_buffers is None - - monkeypatch.setattr(MemoryManager, "write_to_shm", check_shm_refs) - - def create(tp): - config = Glm5NextCacheConfig( + def create(tp, target_mtp_step=None): + local_mtp_step = mtp_step if target_mtp_step is None else target_mtp_step + config = Glm5NextResourceConfig( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, + full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=head_dim, + full_att_head_dim=584, global_linear_k_heads=8, global_linear_v_heads=8, num_linear_k_heads=8 // tp, @@ -56,19 +60,49 @@ def create(tp): head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, + linear_layer_num=6, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=(("linear_attention",) * 3 + ("deepseek_sparse_attention",)) * 2, + full_attention_interval=4, + all_layer_num=8, draft_full_att_kv_layer_num=int(mtp_step > 0), ) mems = [] - for _ in range(tp): - mem = mem_manager_class(32, dtype, 1, head_dim, 2 + int(mtp_step > 0), config) - req = Glm5NextReqManager(3, 32, mem, config) - mem.write_to_shm(req) - assert mem.big_page_buffers is not None - mems.append(mem) - mems[0].alloc_paged_kv_move_buffer(1, 4096) + for rank in range(tp): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + config.get_full_att_layer_ids(), + 1, + 584, + torch.bfloat16, + tp_rank=rank, + tp_world_size=tp, + kv_dim=512, + ), + page_size=args.page_size, + ) + full.initialize_storage(32) + resources = Glm5NextAttResourceManager(config, 4, local_mtp_step, tp_rank=rank) + main, draft = create_full_attention_roles(full, config.all_layer_num) + resources.bind_storage( + main, + FullAttKVAllocator(full.capacity, full.page_size), + draft_full_state=draft, + big_page_capacity=6, + small_page_capacity=2, + ) + mems.append( + SimpleNamespace( + full=main, + full_states=tuple(state for state in (main, draft) if state is not None), + linear=resources.main_kda_state, + tail=resources.main_indexer_kpool_c4_tail_state, + tails=resources.get_attention_tail_states(), + resources=resources, + ) + ) return mems return create @@ -76,16 +110,17 @@ def create(tp): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) -def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder): - source, dest = make_mems(prefill_tp), make_mems(decode_tp) +@pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder, decode_mtp_step): + source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) mtp_size = get_env_start_args().mtp_step + 1 - layers = source[0].layer_num + dst_mtp_size = decode_mtp_step + 1 + layers = sum(len(state.config.layer_ids) for state in source[0].full_states) src_req, dst_req = 0, 2 length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - token_bytes = source[0].head_dim * source[0].dtype.itemsize - packed = torch.randint(0, 256, (layers, length, 1, token_bytes), dtype=torch.uint8, device="cuda") + packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") # Non-live slots deliberately contain stale data. The sequence length, not @@ -93,49 +128,101 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem tail = torch.randn(layers, 3 + mtp_size, 256, dtype=torch.bfloat16, device="cuda") for rank, mem in enumerate(source): heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) - mem.kv_buffer.view(torch.uint8)[:, src_indexes] = packed - mem.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.req_to_ssm_state.buffer[:, src_req * mtp_size].copy_(ssm[:, heads]) - mem.req_to_indexer_tail.buffer[:, src_req].copy_(tail) + offset = 0 + for state in mem.full_states: + end = offset + len(state.config.layer_ids) + state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] + offset = end + mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) + mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) + offset = 0 + for state in mem.tails: + end = offset + len(state.config.layer_ids) + state.tail_buffer[:, src_req].copy_(tail[offset:end]) + offset = end untouched = [] for mem in dest: - mem.kv_buffer.zero_() - mem.req_to_conv_state.buffer.normal_() - mem.req_to_ssm_state.buffer.normal_() - mem.req_to_indexer_tail.buffer.normal_() + for state in mem.full_states: + state.storage.zero_() + mem.linear.conv_state.normal_() + mem.linear.ssm_state.normal_() + for state in mem.tails: + state.tail_buffer.normal_() untouched.append( [ - x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() + x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( - (mem.req_to_conv_state, 1), - (mem.req_to_ssm_state, mtp_size), - (mem.req_to_indexer_tail, 1), + (mem.linear.conv_state, 1), + (mem.linear.ssm_state, dst_mtp_size), + *((state.tail_buffer, 1) for state in mem.tails), ) ] ) - source[0].write_mem_to_page_kv_move_buffer(src_indexes, 0, 0, source, prefill_tp) - # Model a byte-preserving transport between separate P/D page buffers. - dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) - dest[0].read_page_kv_move_buffer_to_mem(dst_indexes, 0, 0, dest, decode_tp) - source[0].write_mem_to_page_kv_move_buffer([], 0, 0, source, prefill_tp, "att_state", src_req) - dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) - dest[0].read_page_kv_move_buffer_to_mem([], 0, 0, dest, decode_tp, "att_state", dst_req) + source_resources = [mem.resources for mem in source] + dest_resources = [mem.resources for mem in dest] + source_buffer = allocate_transfer_buffer(source_resources[0], 1, device="cuda") + dest_buffer = allocate_transfer_buffer(dest_resources[0], 1, device="cuda") + assert source_buffer.shape == dest_buffer.shape + assert source_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() == dest_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() + for kind in ("kv", "att_state"): + copy_page( + source_resources, + source_buffer, + page_index=0, + dp_index=0, + tp_size=prefill_tp, + token_indices=torch.tensor(src_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), + kind=kind, + req_idx=src_req, + state_token_len=length, + write=True, + ) + dest_buffer.copy_(source_buffer) + copy_page( + dest_resources, + dest_buffer, + page_index=0, + dp_index=0, + tp_size=decode_tp, + token_indices=torch.tensor(dst_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), + kind=kind, + req_idx=dst_req, + state_token_len=length, + write=False, + ) torch.cuda.synchronize() for rank, mem in enumerate(dest): heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) - assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, dst_indexes], packed) - assert torch.equal(mem.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mtp_size], ssm[:, heads]) - assert torch.equal(mem.req_to_indexer_tail.buffer[:, dst_req], tail) + for state, source_state in zip(mem.full_states, source[0].full_states): + expected = full_state_test_fields(source_state, {"kv": source_state.storage}) + for field in full_state_test_fields(state, {"kv": state.storage}).values(): + source_field = expected[field.name] + assert torch.equal( + field.tensor[:, dst_indexes].contiguous().view(torch.uint8), + source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), + ) + assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) + assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) + restored_tail = torch.cat([state.tail_buffer[:, dst_req] for state in mem.tails]) + expected_tail = torch.zeros_like(restored_tail) + for position in range(length - 4, length): + expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] + assert torch.equal(restored_tail, expected_tail) + if mem.linear.mtp_state_index is not None: + assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( - ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mtp_size), (mem.req_to_indexer_tail, 1)), + ( + (mem.linear.conv_state, 1), + (mem.linear.ssm_state, dst_mtp_size), + *((state.tail_buffer, 1) for state in mem.tails), + ), untouched[rank], ): indexes = torch.tensor([0, 1, 3], device="cuda") * stride - assert torch.equal(state.buffer[:, indexes], expected) + assert torch.equal(state[:, indexes], expected) # Continue through the next pool boundary on P and on the restored D state. # Different physical KV/request slots must yield identical pooled FP8 bytes. @@ -149,8 +236,8 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem all_indexes = indexes + new_indexes compress_pools( raw=raw, - tail=mem.req_to_indexer_tail.buffer[0], - packed_buffer=mem.get_indexer_k_buffer(3), + tail=mem.tail.tail_buffer[0], + packed_buffer=mem.full.get_layer_view(3).indexer_k, ape=ape, lengths=tensor(range(length + 1, length + count + 1)), starts=tensor([0] * count), @@ -160,16 +247,71 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem seq_lens=tensor([length + count]), max_q_len=count, ) - outputs.append(mem.get_indexer_k_buffer(3)[new_indexes[-1]].clone()) + outputs.append(mem.full.get_layer_view(3).indexer_k[new_indexes[-1]].clone()) assert torch.equal(*outputs) -def test_pd_page_capacity_includes_tail_without_resizing(make_mems): +def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch): + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - helper = mem.att_state_page_helper - # A page that fits Conv+SSM but not the appended tail must be rejected. - mem.kv_move_buffer = torch.empty((1, helper.tail_offset), dtype=torch.uint8, device="cuda") - with pytest.raises(AssertionError, match="smaller than global linear att state"): - helper.assert_page_size() - shape = mem.get_paged_kv_move_buffer_shape(2, 2048) - assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim) + with pytest.raises(AssertionError, match="must be at least the complete attention runtime"): + mem.resources.get_pd_page_bytes() + + +@pytest.mark.parametrize("source_step,target_step", [(0, 1), (0, 5), (3, 2), (5, 1)]) +@pytest.mark.parametrize("prefix", [1, 3, 4, 5, 11, 12]) +def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, prefix): + """真实 prefill 导出后，在不同大小的 ring 上验证、拒绝回退并再次验证。""" + from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.common.att_resources.types import StateTransferLayout + + torch.manual_seed(53) + config = KPoolStateConfig(layer_ids=(0,), pool_size=4, head_dim=128, dtype=torch.bfloat16) + source = create_kpool_state(config=config, request_capacity=2, mtp_step=source_step, device="cuda") + target = create_kpool_state(config=config, request_capacity=2, mtp_step=target_step, device="cuda") + raw = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) + packed = torch.zeros(64, 132, device="cuda", dtype=torch.uint8) + ragged = torch.arange(64, device="cuda", dtype=torch.int32) + + def prefill(values, destination, tail): + count = len(values) + compress_pools( + raw=values, tail=tail, packed_buffer=destination, ape=ape, + lengths=torch.arange(1, count + 1, device="cuda", dtype=torch.int32), + starts=torch.zeros(count, device="cuda", dtype=torch.int32), ragged=ragged, + req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + cu_q_lens=torch.tensor([0, count], device="cuda", dtype=torch.int32), + seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), max_q_len=count, + ) + + prefill(raw[:prefix], packed, source.get_layer_view(0)) + regions = source.get_pd_transfer_regions() + assert regions == target.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=source.resource_id, tokens_per_page=None, + byte_start=0, byte_end=regions[-1].byte_end, regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + source.write_to_pd_buffer(0, pages, pd_page_index=0, state_token_len=prefix, layout=layout) + target.tail_buffer.fill_(17) + target.read_from_pd_buffer(pages, 0, pd_page_index=0, state_token_len=prefix, layout=layout) + assert (target.tail_buffer[:, 1] == 17).all() + width = target_step + 1 + # 第一次验证后回访已接受位置；新增候选历史必须由 D 的本地 ring 正确保留。 + for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + raw[start:start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) + lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) + compress_pools( + raw=raw[start:start + count], tail=target.get_layer_view(0), packed_buffer=packed, ape=ape, + lengths=lengths, starts=torch.zeros(count, device="cuda", dtype=torch.int32), ragged=ragged, + req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), + cu_q_lens=torch.arange(count + 1, device="cuda", dtype=torch.int32), + seq_lens=lengths, max_q_len=1, + mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + ) + expected = torch.zeros_like(packed) + prefill(raw[:start + count], expected, torch.zeros_like(source.get_layer_view(0))) + closing = torch.arange(start, start + count, device="cuda") + closing = closing[(closing + 1) % 4 == 0] + assert torch.equal(packed[closing], expected[closing]) diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py index 8e2085fa00..37fb311fbc 100644 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -65,7 +65,7 @@ def test_parallel_block_decode_commits_target_hiddens_directly(model_class): model = model_class.__new__(model_class) model._cos_cached = torch.arange(24).view(6, 4) model._sin_cached = model._cos_cached + 100 - model.mem_manager = object() + model.att_resource_manager = object() model.pre_post_weight = object() target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) @@ -104,7 +104,9 @@ def context_forward(self, hidden, infer_state, layer_weight): assert isinstance(output, ModelOutput) assert output.logits.shape == (2, 1) infer_state = observed_states[0] - assert infer_state.mem_manager is model.mem_manager + assert not hasattr(infer_state, "full_att_state") + assert infer_state.att_resource_manager is model.att_resource_manager + assert infer_state.is_draft_model assert infer_state.mem_index is mem_indexes assert torch.equal(infer_state.position_cos, model._cos_cached[[2, 4]]) assert torch.equal(infer_state.position_sin, model._sin_cached[[2, 4]]) diff --git a/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py b/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py index 8816e71d6b..1b2675737a 100644 --- a/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py +++ b/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py @@ -148,7 +148,7 @@ def test_factory_creates_legacy_controller_for_legacy_strategy(): def test_factory_calculates_adaptive_cache_capacities_across_local_dp_replicas_and_applies_gpu_ratio(monkeypatch): from lightllm.utils import envs_utils, kv_cache_utils - cpu_cache_meta = SimpleNamespace(page_num=10, calcu_one_page_size=lambda: 64) + cpu_cache_meta = SimpleNamespace(page_num=10, page_bytes=64) monkeypatch.setattr(kv_cache_utils, "calcu_cpu_cache_meta", lambda: cpu_cache_meta) monkeypatch.setattr(envs_utils, "get_cache_placement_gpu_capacity_ratio", lambda: 0.5) args = SimpleNamespace( diff --git a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py index c68d7e3439..0928d0bf29 100644 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py @@ -12,7 +12,7 @@ def test_disk_roundtrip_preserves_raw_pages_before_mm_boundary_trim(cpu_cache_fa light_mem = pytest.importorskip("light_mem") from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker - client = cpu_cache_factory(CpuKVCacheMeta(3, 1, 1, 1, 4096, torch.uint8, 0, torch.uint8)) + client = cpu_cache_factory(CpuKVCacheMeta(3, 4096)) pages, _ = client.allocate_pages([11, 22, 33], False) for index in pages: client.cpu_kv_cache_tensor[index].fill_(index + 1) @@ -69,19 +69,19 @@ def allocate(req, count): return torch.arange(count, dtype=torch.int32) module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace(page_size=1) + module.args = SimpleNamespace(page_size=1, cpu_cache_token_page_size=256) module.cpu_cache_client = client module.init_sync_group = None module.need_sync_compute_stream = lambda: False module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=False, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(1024).view(1, -1)), - mem_manager=SimpleNamespace( - operator=SimpleNamespace( - load_cpu_cache_to_gpu=lambda **kwargs: copied_pages.extend(kwargs["page_indexes"].tolist()) - ) + att_resource_manager=SimpleNamespace( + get_cpu_transfer_layouts=lambda: {}, + read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: copied_pages.extend(pages), ), ), ) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index f04341a983..35c10802be 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -9,6 +9,8 @@ from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.utils import shm_utils +pytestmark = pytest.mark.usefixtures("kv_allocator_environment") + @pytest.fixture(scope="module", autouse=True) def service_name(): @@ -95,7 +97,6 @@ def test_case3(): def test_case4(): - tree = RadixCache(100, 2) ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu")) ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) @@ -245,33 +246,58 @@ def test_case9(): assert torch.equal(unmerged_node_d.token_id_key, torch.tensor([6], dtype=torch.int64)) -def test_case10(): - """ - 测试场景：测试 flush_cache 函数 - """ - print("\nTest Case 10: Testing flush_cache function\n") - tree = RadixCache(100, 0) - tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) - tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) - tree_node, size, values = tree.match_prefix( - torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node is not None - assert size == 3 +def test_flush_cache_requires_released_references_and_returns_allocator_pages(): + from lightllm.common.att_resources.allocators import FullAttKVAllocator + + allocator = FullAttKVAllocator(100, pin_memory=False) + tree = RadixCache(100, 0, att_resource_manager=SimpleNamespace(full_att_allocator=allocator)) + values = allocator.alloc(5) + key = torch.arange(5) + tree.insert(key[:3], values[:3]) + tree.insert(key, values) + node, matched_len, _ = tree.match_prefix(key[:3], update_refs=True) + assert matched_len == 3 + with pytest.raises(AssertionError, match="can not free"): + tree.flush_cache() + assert allocator.can_use_mem_size == 95 + tree.dec_node_ref_counter(node) tree.flush_cache() - tree_node, size, values = tree.match_prefix( - torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node is None - assert size == 0 - assert tree.get_tree_total_tokens_num() == 0 - assert tree.get_refed_tokens_num() == 0 - assert len(tree.root_node.children) == 0 - assert tree.root_node.token_id_key.numel() == 0 - assert tree.root_node.token_mem_index_value.numel() == 0 + assert allocator.can_use_mem_size == 100 + assert tree.get_tree_total_tokens_num() == tree.get_refed_tokens_num() == 0 + assert not tree.root_node.children assert tree.root_node.ref_counter == 1 +def test_hybrid_flush_releases_shared_kv_and_state_slots(): + from unit_tests.common.att_resources.test_manager_helpers import composition + + resources = composition(with_tail=False) + allocator = resources.full_att_allocator + big = resources.bounded_att_big_page_allocator.alloc(1)[0] + small = resources.bounded_att_small_page_allocator.alloc(1)[0] + resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[big]) + resources.save_bounded_att_pages([1], page_kind="small", slot_ids=[small]) + values = allocator.alloc(12) + tree = HybridAttPagedRadixCache(12, 106, 4, 2, resources) + tree.insert( + torch.arange(12), + values, + block_hashs=[11, 22, 33], + block_state_idxs=[None, None, small], + len_to_big_page_id=SortedDict({8: big}), + ) + node, matched_len, _ = tree.match_prefix(torch.arange(12), block_hashs=[11, 22, 33], update_refs=True) + assert matched_len == 12 + with pytest.raises(AssertionError, match="can not free"): + tree.flush_cache() + tree.dec_node_ref_counter(node) + tree.flush_cache() + assert allocator.can_use_mem_size == 12 + assert resources.bounded_att_big_page_allocator.free_count == 3 + assert resources.bounded_att_small_page_allocator.free_count == 2 + assert not tree.root_node.children + + def test_page_aligned_insert_and_match(): tree = RadixCache(100, 99, page_size=4) values = torch.arange(100, 110, dtype=torch.int64) @@ -301,11 +327,11 @@ def test_page_aligned_insert_and_match(): assert tree.get_tree_total_tokens_num() == 12 -def test_page_size_must_match_mem_manager(): +def test_page_size_must_match_allocator(): mem_manager = SimpleNamespace(page_size=8) - with pytest.raises(ValueError, match="must match mem_manager page_size 8"): - RadixCache(100, 100, mem_manager=mem_manager, page_size=4) + with pytest.raises(ValueError, match="must match allocator page_size 8"): + RadixCache(100, 100, att_resource_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) @pytest.mark.parametrize("update_refs", [False, True]) @@ -404,20 +430,25 @@ def test_page_key_bytes_does_not_share_tensor_memory(): assert key != TreeNode(page_size=4).get_child_key(token_ids) +def hybrid_resources(*, freed_big=None, freed_small=None): + return SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace( + free=lambda slots: freed_big.extend(slots) if freed_big is not None else None + ), + bounded_att_small_page_allocator=SimpleNamespace( + free_count=0, free=lambda slots: freed_small.extend(slots) if freed_small is not None else None + ), + ) + + def test_hybrid_match_refs_only_valid_prefix_after_tail_state_eviction(): freed_state_ids = [] - small_page_buffers = SimpleNamespace( - get_free_cache_num=lambda: 0, - free_state_cache=lambda free_indexes: freed_state_ids.extend(free_indexes), - ) - mem_manager = SimpleNamespace(big_page_buffers=SimpleNamespace()) tree = HybridAttPagedRadixCache( total_token_num=100, rank_in_node=101, hash_page_size=4, big_page_num=2, - kv_cache_mem_manager=mem_manager, - small_page_buffers=small_page_buffers, + att_resource_manager=hybrid_resources(freed_small=freed_state_ids), ) key = torch.arange(12, dtype=torch.int64) values = torch.arange(100, 112, dtype=torch.int64) @@ -457,10 +488,7 @@ def hybrid_tree(): rank_in_node=102, hash_page_size=4, big_page_num=2, - kv_cache_mem_manager=SimpleNamespace( - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_states.extend) - ), - small_page_buffers=SimpleNamespace(free_state_cache=lambda **kwargs: None), + att_resource_manager=hybrid_resources(freed_big=freed_big_states), ) return tree, freed_big_states @@ -526,10 +554,7 @@ def test_hybrid_match_falls_back_along_matched_branch(monkeypatch, small_state_v rank_in_node=104, hash_page_size=64, big_page_num=4, - kv_cache_mem_manager=SimpleNamespace( - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_states.extend) - ), - small_page_buffers=SimpleNamespace(get_free_cache_num=lambda: 0, free_state_cache=lambda **kwargs: None), + att_resource_manager=hybrid_resources(freed_big=freed_big_states), ) # The older request leaves a small-page state at 384, before the image. old_values = torch.arange(1000, 1384) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py new file mode 100644 index 0000000000..9464cc18a0 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -0,0 +1,186 @@ +"""DP 前缀迁移：独立共享大页列表、目标页分配和源/目标索引配对。""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans + + +@pytest.mark.parametrize("hybrid,max_req_len", [(False, 32), (True, 32), (True, 4)]) +def test_dp_page_ids_use_separate_shared_array(monkeypatch: pytest.MonkeyPatch, hybrid: bool, max_req_len: int): + arrays = {} + + def shared_array(name: str, shape: tuple[int, ...], dtype: type) -> SimpleNamespace: + array = SimpleNamespace(arr=np.full(shape, -1, dtype=dtype), create_shm=lambda: None) + arrays[name] = array + return array + + monkeypatch.setattr(dp_shared_kv_trans, "ShmArray", shared_array) + monkeypatch.setattr(dp_shared_kv_trans, "get_dp_rank_in_node", lambda: 0) + monkeypatch.setattr(dp_shared_kv_trans, "get_env_start_args", lambda: SimpleNamespace(diverse_mode=False)) + backend = SimpleNamespace( + is_hybrid_att_model=hybrid, + args=SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2, max_req_total_len=max_req_len), + ) + module = dp_shared_kv_trans.DPKVSharedMoudle(max_req_num=3, dp_size_in_node=2, backend=backend) + assert arrays["dp_shared_req_infos"].arr.shape == (3, 2, 2) + if hybrid and max_req_len >= 8: + assert arrays["dp_shared_req_big_page_ids"].arr.shape == (3, 2, 4) + else: + assert module.shared_req_big_page_ids is None + assert list(arrays) == ["dp_shared_req_infos"] + + +def test_fill_orders_radix_and_owned_pages_and_clears_reused_rows(monkeypatch: pytest.MonkeyPatch): + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace(is_master_in_dp=True, node_nccl_group=None) + module.dp_rank_in_node = 0 + module.big_page_tokens = 8 + module.shared_req_infos = SimpleNamespace(arr=np.full((2, 2, 2), 99, dtype=np.int64)) + module.shared_req_big_page_ids = SimpleNamespace(arr=np.full((2, 2, 4), 99, dtype=np.int64)) + node = object() + context = SimpleNamespace(radix_cache=SimpleNamespace(get_big_page_ids_by_node=lambda _: [40])) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + req = SimpleNamespace( + cur_kv_len=24, req_idx=7, shared_kv_node=node, + hybrid_len_to_big_page_id={24: 3, 16: 12, 32: 88}, + ) + second = SimpleNamespace(cur_kv_len=8, req_idx=5, shared_kv_node=None, hybrid_len_to_big_page_id={8: 6}) + module.fill_reqs_info([req, second]) + assert module.shared_req_infos.arr[:, 0].tolist() == [[24, 7], [8, 5]] + assert module.shared_req_big_page_ids.arr[:, 0].tolist() == [[40, 12, 3, -1], [6, -1, -1, -1]] + assert (module.shared_req_big_page_ids.arr[:, 1] == 99).all() + + req.cur_kv_len = 8 + module.fill_reqs_info([req]) + assert module.shared_req_big_page_ids.arr[0, 0].tolist() == [40, -1, -1, -1] + module.backend.is_master_in_dp = False + req.cur_kv_len = 0 + module.fill_reqs_info([req]) + assert module.shared_req_infos.arr[0, 0, 0] == 8 + + +@pytest.mark.parametrize("capacity,src_kv_len", [(2, 26), (1, 26), (0, 14)]) +def test_build_reserves_only_new_big_pages_before_kv( + monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: int, +): + module, req, allocator, allocations = _build_fixture(monkeypatch, capacity, src_kv_len) + if capacity == 1: + with pytest.raises(AssertionError, match="DP fetch big page allocation must succeed"): + module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + assert not allocations + assert req.hold_kv_len == 8 and allocator.free_count == 1 + return + tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + task, = tasks + assert task.dst_req is req and task.src_req_idx == 1 + assert task.src_dp_rank == 1 + assert task.src_token_indexes.tolist() == list(range(40, 32 + src_kv_len)) + assert len(task.dst_token_indexes) == src_kv_len - 8 + assert allocations == [((src_kv_len + 3) // 4) * 4 - 8] + assert req.cur_kv_len == 8 and req.hybrid_len_to_big_page_id == {8: 77} + if src_kv_len == 26: + assert task.src_big_page_ids == [20, 21] + assert task.dst_big_page_ids == [0, 1] + assert allocator.free_count == 0 + else: + assert task.src_big_page_ids == task.dst_big_page_ids == [] + + +def test_missing_source_page_fails_before_allocating(monkeypatch: pytest.MonkeyPatch): + module, req, allocator, allocations = _build_fixture(monkeypatch, 2, 26) + module.shared_req_big_page_ids.arr[0, 1, 2] = -1 + with pytest.raises(AssertionError, match="missing prefix big pages"): + module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + assert not allocations and allocator.free_count == 2 + + +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize("local_dp_rank", [0, 1]) +@pytest.mark.parametrize("local_len,remote_len", [(24, 16), (24, 24), (31, 24), (0, 0)]) +def test_local_longest_prefix_skips_copy_but_keeps_barriers( + monkeypatch: pytest.MonkeyPatch, hybrid: bool, local_dp_rank: int, local_len: int, remote_len: int, +): + # 并列最长且本地为 DP 1 时，argmax 会优先选 DP 0；仍应在选择来源前跳过复制。 + allocate_pages = Mock(side_effect=AssertionError("unexpected big page allocation")) + allocate_tokens = Mock(side_effect=AssertionError("unexpected KV allocation")) + copy_resources = Mock(side_effect=AssertionError("unexpected DP copy")) + allocation_need = Mock(side_effect=AssertionError("unexpected allocation calculation")) + available_tokens = Mock(side_effect=AssertionError("unexpected capacity query")) + barriers = [] + manager = SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace(alloc=allocate_pages), + copy_from_other_dp=copy_resources, + ) + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace( + is_hybrid_att_model=hybrid, is_master_in_dp=True, node_nccl_group="node", + model=SimpleNamespace(att_resource_manager=manager), + dp_att_resource_managers={}, dp_req_to_token_indices={}, _alloc_req_kv_mem=allocate_tokens, + ) + module.dp_rank_in_node = local_dp_rank + module.big_page_tokens = 8 if hybrid else 0 + module.shared_req_infos = SimpleNamespace(arr=np.zeros((1, 2, 2), dtype=np.int64)) + module.shared_req_infos.arr[0, 1 - local_dp_rank] = [remote_len, 9] + module.shared_req_big_page_ids = ( + SimpleNamespace(arr=np.full((1, 2, 4), -1, dtype=np.int64)) if hybrid else None + ) + owned_pages = {length: 70 + length // 8 for length in range(8, local_len + 1, 8)} if hybrid else {} + req = SimpleNamespace( + req_idx=7, cur_kv_len=local_len, hold_kv_len=local_len, + shared_kv_node=None, hybrid_len_to_big_page_id=owned_pages.copy(), + get_cur_total_len=lambda: 32, _kv_cache_alloc_need=allocation_need, + shm_req=SimpleNamespace(shm_cur_kv_len=local_len), + ) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=available_tokens)) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: barriers.append(group)) + + module.fill_reqs_info([req]) + assert module.shared_req_infos.arr[0, local_dp_rank].tolist() == [local_len, req.req_idx] + tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[local_dp_rank]) + assert tasks == [] + module.kv_trans(tasks) + + for operation in (allocate_pages, allocate_tokens, copy_resources, allocation_need, available_tokens): + operation.assert_not_called() + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == local_len + assert req.hybrid_len_to_big_page_id == owned_pages + # 即使不接收数据，也必须参加所有屏障，保证其他 DP 能完成读取后再释放来源资源。 + assert barriers == ["node", "node", "node"] + + +def _build_fixture(monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: int) -> tuple: + allocator = StatePageAllocator(capacity) + allocations = [] + req = SimpleNamespace( + req_idx=0, cur_kv_len=8, hold_kv_len=8, hybrid_len_to_big_page_id={8: 77}, + get_cur_total_len=lambda: 32, + ) + req._kv_cache_alloc_need = lambda length: ((length + 3) // 4) * 4 - req.hold_kv_len + + def allocate(cur_req: SimpleNamespace, count: int) -> torch.Tensor: + allocations.append(count) + cur_req.hold_kv_len += count + return torch.arange(100, 100 + count, dtype=torch.int32) + + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace( + is_hybrid_att_model=True, dp_world_size=2, rank_in_dp=1, node_nccl_group=None, + model=SimpleNamespace(att_resource_manager=SimpleNamespace(bounded_att_big_page_allocator=allocator)), + dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, + _alloc_req_kv_mem=allocate, + ) + module.dp_rank_in_node = 0 + module.big_page_tokens = 8 + module.shared_req_infos = SimpleNamespace(arr=np.array([[[8, 0], [src_kv_len, 1]]], dtype=np.int64)) + module.shared_req_big_page_ids = SimpleNamespace(arr=np.array([[[77, -1, -1, -1], [10, 20, 21, -1]]])) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=lambda: 100)) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + return module, req, allocator, allocations diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index 82422adf8b..c97a54391d 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -97,7 +97,7 @@ def start_offload(req, cpu_kv_cache_stream): offload_calls.append((req.shm_req.request_id, CacheTier.DISK in req.cache_tiers)) return SimpleNamespace(req=req) - module._start_kv_cache_offload_task = start_offload + module.write_req_to_cpu = start_offload monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_hybrid_att_model", False) monkeypatch.setattr( multi_level_kv_cache_impl.g_infer_context, @@ -138,8 +138,7 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): context.is_hybrid_att_model = True context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) context.radix_cache = SimpleNamespace( - small_page_buffers=SimpleNamespace(free_state_cache=freed_small_pages.extend), - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_pages.extend), + _release_pages=lambda kind, slots: (freed_small_pages if kind == "small" else freed_big_pages).extend(slots), ) req = SimpleNamespace( req_idx=0, @@ -181,14 +180,16 @@ def alloc_req_kv_mem(req, alloc_token_num): req.hold_kv_len = new_hold_kv_len return mem_indexes - operator = SimpleNamespace( - load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) + resources = SimpleNamespace( + get_cpu_transfer_layouts=lambda: {}, + read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: loaded_indexes.extend(tokens.tolist()), ) backend = SimpleNamespace( is_master_in_dp=False, + is_hybrid_att_model=False, radix_cache=None, model=SimpleNamespace( - mem_manager=SimpleNamespace(operator=operator), + att_resource_manager=resources, req_manager=SimpleNamespace(req_to_token_indexs=table), ), _alloc_req_kv_mem=alloc_req_kv_mem, @@ -198,7 +199,7 @@ def alloc_req_kv_mem(req, alloc_token_num): module.backend = backend module.init_sync_group = object() module.need_sync_compute_stream = lambda: False - module.cpu_cache_client = SimpleNamespace() + module.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=None) context = SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table, mem_manager=SimpleNamespace(alloc=alloc)), get_can_alloc_token_num=lambda: 200, diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py new file mode 100644 index 0000000000..1bacdfda1d --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import concurrent.futures +import socket +import threading + +import rpyc + +from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as module + +if TYPE_CHECKING: + import pytest + + +def test_concurrent_control_channels_retry_port_before_publishing_metadata(monkeypatch: pytest.MonkeyPatch): + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + first_port = probe.getsockname()[1] + + original = module.ThreadedServer + both_bound = threading.Barrier(2) + + def bind_before_listening(*args: object, **kwargs: object): + server = original(*args, **kwargs) + if kwargs["port"] == first_port: + # Force both SO_REUSEADDR sockets to bind before either listens. + both_bound.wait(timeout=10) + return server + + monkeypatch.setattr(module, "ThreadedServer", bind_before_listening) + channels = [] + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(module._NcclControlChannel, "127.0.0.1", first_port, first_port + 4) for _ in range(2) + ] + for future in futures: + channels.append(future.result(timeout=15)) + assert len({channel.port for channel in channels}) == 2 + for index, channel in enumerate(channels): + payload = f"channel-{index}".encode() + connection = rpyc.connect("127.0.0.1", channel.port) + try: + connection.root.push_notif(payload) + finally: + connection.close() + assert channel.get_notifs() == [payload] + finally: + for channel in channels: + channel.close() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py index 22ccc9e3cd..bf5bd040a3 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -13,24 +13,6 @@ from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl -def _bind_alloc_need(req, page_size): - req.args = SimpleNamespace(page_size=page_size) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req - - -def _reserve_into_table(table, req, alloc_sizes, next_index, alloc_token_num): - if alloc_token_num == 0: - return None - alloc_sizes.append(alloc_token_num) - new_hold_kv_len = req.hold_kv_len + alloc_token_num - mem_indexes = torch.arange(next_index[0], next_index[0] + alloc_token_num, dtype=torch.int32) - table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes - next_index[0] += alloc_token_num - req.hold_kv_len = new_hold_kv_len - return mem_indexes - - @pytest.mark.parametrize("is_hybrid_att_model", [False, True]) def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatch, is_hybrid_att_model): table = torch.full((1, 16), -1, dtype=torch.int32) @@ -101,7 +83,8 @@ def test_pd_prefill_rejects_unaligned_transfer_page_size(): backend._prefill_chuncked_handle_func(req, next_token_id=1, next_token_prob=1.0, output_len=0) -def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypatch): +@pytest.mark.parametrize("tp_size,tp_rank", [(1, 0), (2, 0), (2, 1)]) +def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypatch, tp_size, tp_rank): table = torch.full((1, 16), -1, dtype=torch.int32) table[0, :4] = torch.arange(4, dtype=torch.int32) alloc_sizes = [] @@ -119,10 +102,11 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat backend = SimpleNamespace( is_hybrid_att_model=False, model=SimpleNamespace(mem_manager=SimpleNamespace()), - dp_world_size=1, - rank_in_dp=0, + dp_world_size=tp_size, + rank_in_dp=tp_rank, node_nccl_group=None, - mem_managers=[SimpleNamespace(req_to_token_indexs=source_table)] * 2, + dp_req_to_token_indices={0: source_table, 1: source_table}, + is_hybrid_att_model=False, ) backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table(table, cur_req, alloc_sizes, next_index, size) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) @@ -142,5 +126,82 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat assert req.cur_kv_len == 4 assert req.hold_kv_len == 12 assert len(tasks) == 1 - assert tasks[0].mem_indexes.tolist() == list(range(4, 10)) - assert len(tasks[0].max_kv_len_mem_indexes) == 6 + assert tasks[0].src_dp_rank == 1 + assert tasks[0].dst_token_indexes.tolist() == list(range(4, 10)) + assert len(tasks[0].src_token_indexes) == 6 + + +@pytest.mark.parametrize( + "node_size,tp_size,rank_in_node,expected_ranks", + [ + (4, 2, 0, [0, 2]), + (4, 2, 1, [1, 3]), + (4, 2, 2, [0, 2]), + (4, 2, 3, [1, 3]), + (6, 2, 5, [1, 3, 5]), + (2, 4, 1, [1]), + ], +) +def test_dp_imports_only_matching_tp_shards(monkeypatch, node_size, tp_size, rank_in_node, expected_ranks): + from lightllm.server.router.model_infer.mode_backend import base_backend + from lightllm.server.router.model_infer import shared_att_resources_utils + + imported_ranks = [] + registered_ranks = [] + managers = {} + tables = {} + for rank in expected_ranks: + managers[rank] = SimpleNamespace( + register_page_pools_pin_memory=lambda rank=rank: registered_ranks.append(rank), + ) + tables[rank] = object() + + def load_storage(rank, *, kind): + assert kind == "dp" + imported_ranks.append(rank) + return managers[rank], tables[rank] + + local_manager = managers[rank_in_node] + backend = SimpleNamespace( + args=SimpleNamespace(running_max_req_size=8), + dp_size_in_node=max(1, node_size // tp_size), + dp_world_size=tp_size, + node_world_size=node_size, + rank_in_node=rank_in_node, + model=SimpleNamespace( + att_resource_manager=local_manager, + req_manager=SimpleNamespace(req_to_token_indexs=tables[rank_in_node]), + ), + ) + monkeypatch.setattr(base_backend, "get_current_device_id", lambda: 0) + monkeypatch.setattr(base_backend.torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(dp_shared_kv_trans, "DPKVSharedMoudle", lambda **kwargs: object()) + monkeypatch.setattr(shared_att_resources_utils, "load_att_resources", load_storage) + + base_backend.ModeBackend.init_dp_kv_shared(backend) + + remote_ranks = [rank for rank in expected_ranks if rank != rank_in_node] + assert imported_ranks == registered_ranks == remote_ranks + assert list(backend.dp_att_resource_managers) == list(range(len(expected_ranks))) + assert list(backend.dp_req_to_token_indices) == list(range(len(expected_ranks))) + for dp_rank, rank in enumerate(expected_ranks): + assert backend.dp_att_resource_managers[dp_rank] is managers[rank] + assert backend.dp_req_to_token_indices[dp_rank] is tables[rank] + + +def _bind_alloc_need(req, page_size): + req.args = SimpleNamespace(page_size=page_size) + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + return req + + +def _reserve_into_table(table, req, alloc_sizes, next_index, alloc_token_num): + if alloc_token_num == 0: + return None + alloc_sizes.append(alloc_token_num) + new_hold_kv_len = req.hold_kv_len + alloc_token_num + mem_indexes = torch.arange(next_index[0], next_index[0] + alloc_token_num, dtype=torch.int32) + table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes + next_index[0] += alloc_token_num + req.hold_kv_len = new_hold_kv_len + return mem_indexes diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py index 7bd0940c89..6af1ebcc8a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py @@ -124,8 +124,7 @@ def test_notification_payload(backend, page_kind, method, stage, destination): if backend is NcclKVTransporter: transporter.host_ip = "127.0.0.1" transporter.control_channel = SimpleNamespace(port=1234, send_notif=Mock()) - transporter.kv_move_buffer = SimpleNamespace(dtype="float16") - transporter.page_size = transporter.num_layers = transporter.kv_head_num = transporter.head_dims = 1 + transporter.kv_move_buffer = SimpleNamespace(dtype="uint8", shape=(12, 128)) transporter._get_remote_metadata = Mock(return_value=SimpleNamespace(host_ip="127.0.0.2", control_port=4321)) transporter._get_peer = Mock(return_value=SimpleNamespace(start_recv=Mock())) else: @@ -190,3 +189,33 @@ def test_receiver_rejects_payload_with_decode_local_mark(): ) assert result.returncode != 0 assert b"decode modules are unavailable to the receiver" in result.stderr + + +def test_nixl_rejects_different_peer_page_capacity_before_registering_handles(): + transporter = NixlKVTransporter.__new__(NixlKVTransporter) + transporter.page_len = 128 + transporter.nixl_agent = SimpleNamespace(get_xfer_descs=Mock(), prep_xfer_dlist=Mock()) + with pytest.raises(ValueError, match="same page byte capacity"): + transporter._create_paged_xfer_handles([(4096, 512, 0, None)], 2, "peer") + transporter.nixl_agent.get_xfer_descs.assert_not_called() + transporter._create_paged_xfer_handles([(4096, 256, 0, None)], 2, "peer") + transporter.nixl_agent.get_xfer_descs.assert_called_once_with([(4096, 128, 0), (4224, 128, 0)], "VRAM") + + +def test_nccl_checks_peer_page_shape_before_accepting_agent(): + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclAgentMetadata + + transporter = NcclKVTransporter.__new__(NcclKVTransporter) + transporter.remote_agents = {} + transporter.kv_move_buffer = SimpleNamespace(shape=(2, 128), dtype="torch.uint8") + peer = SimpleNamespace( + agent_name="peer", + agent_metadata=pickle.dumps(NcclAgentMetadata("peer", "127.0.0.1", 1234, 0)), + page_reg_desc=pickle.dumps({"page_shape": (256,), "dtype": "torch.uint8"}), + ) + with pytest.raises(ValueError, match="same page shape"): + transporter.connect_add_remote_agent(peer) + assert not transporter.remote_agents + peer.page_reg_desc = pickle.dumps({"page_shape": (128,), "dtype": "torch.uint8"}) + transporter.connect_add_remote_agent(peer) + assert transporter.remote_agents["peer"] is peer diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py deleted file mode 100644 index 9e44ed774b..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py +++ /dev/null @@ -1,52 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend - - -def test_preload_prompt_cache_keeps_only_complete_kv_pages(monkeypatch): - loaded = {} - - class FakeMemManager: - def alloc(self, size): - loaded["alloc_size"] = size - return torch.arange(size, dtype=torch.int32) - - def load_index_kv_buffer(self, indexes, buffers): - loaded["indexes"] = indexes.clone() - loaded["buffers"] = buffers - - class FakeRadixCache: - def __init__(self): - self.mem_manager = FakeMemManager() - - def insert(self, token_ids, indexes): - loaded["insert"] = (token_ids.clone(), indexes.clone()) - - def match_prefix(self, token_ids, update_refs): - loaded["match"] = (token_ids.clone(), update_refs) - - backend = ModeBackend.__new__(ModeBackend) - backend.logger = SimpleNamespace(info=lambda _: None) - backend.weight_dir = "/model" - backend.args = SimpleNamespace(page_size=4) - backend.radix_cache = FakeRadixCache() - kv_buffer = torch.arange(12).reshape(1, 6, 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) - monkeypatch.setattr(torch, "load", lambda *args, **kwargs: {"kv_buffer": kv_buffer}) - - backend.preload_prompt_cache_kv_buffer( - { - "prompt_cache_kv_buffer": {"rank_0": "prompt_cache.pt"}, - "prompt_cache_token_ids": [10, 11, 12, 13, 14, 15], - } - ) - - assert loaded["alloc_size"] == 4 - assert loaded["indexes"].tolist() == [0, 1, 2, 3] - assert loaded["buffers"]["kv_buffer"].shape == (1, 4, 2) - assert loaded["insert"][0].tolist() == [10, 11, 12, 13] - assert loaded["insert"][1].tolist() == [0, 1, 2, 3] - assert loaded["match"][0].tolist() == [10, 11, 12, 13] - assert loaded["match"][1] is True diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py new file mode 100644 index 0000000000..de2fb835cb --- /dev/null +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -0,0 +1,457 @@ +"""CPU service metadata must match the actual model storage and data path.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +import dataclasses +import json +from types import SimpleNamespace + +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + +from lightllm.models import Deepseek2TpPartModel, Deepseek3_2TpPartModel, LlamaTpPartModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.att_resource_config import cpu_cache_namespace +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_unique_server_name +from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta + +if TYPE_CHECKING: + from collections.abc import Callable + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient + from lightllm.utils.kv_cache_utils import CpuKVCacheMeta + from pathlib import Path + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +@pytest.fixture +def runtime(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + for key, value in { + "LIGHTLLM_UNIQUE_SERVICE_NAME_ID": tmp_path.name, + "LIGHTLLM_CURRENT_RANK_IN_NODE": "0", + "LIGHTLLM_CURRENT_DEVICE_ID": "0", + }.items(): + monkeypatch.setenv(key, value) + yield + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + get_unique_server_name.cache_clear() + calcu_cpu_cache_meta.cache_clear() + + +CASES = [ + ("dense", "None", LlamaTpPartModel), + ("int4", "int4kv", LlamaTpPartModel), + ("int8", "int8kv", LlamaTpPartModel), + ("fp8_head", "fp8kv_sph", LlamaTpPartModel), + ("fp8_tensor", "fp8kv_spt", LlamaTpPartModel), + ("mla", "None", Deepseek2TpPartModel), + ("packed", "None", Deepseek3_2TpPartModel), + ("packed_fp8", "fp8kv_dsa", Deepseek3_2TpPartModel), + ("qwen", "None", Qwen3NextTpPartModel), + ("qwen_head", "fp8kv_sph", Qwen3NextTpPartModel), + ("qwen_tensor", "fp8kv_spt", Qwen3NextTpPartModel), + ("glm", "None", Glm5NextTpPartModel), +] + + +def make_runtime( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, + *, + mtp_layers: int = 0, +): + hybrid = kind.startswith("qwen") or kind == "glm" + config = { + "model_type": "qwen3_5_text" if kind.startswith("qwen") else "glm5_next" if kind == "glm" else "llama", + "architectures": ["LlamaForCausalLM"], + "num_attention_heads": 8, + "num_key_value_heads": 2, + "num_hidden_layers": 4 if hybrid else 2, + "hidden_size": 128, + "head_dim": 16, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "full_attention_interval": 2, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_conv_kernel_dim": 3, + "index_kpool": 4, + "index_head_dim": 128, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "linear_attn_config": {"num_heads": 4, "head_dim": 8, "short_conv_kernel_size": 3, "kda_layers": [0, 1, 2]}, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + args = StartArgs( + model_dir=str(tmp_path), + tp=tp, + dp=1, + data_type="bfloat16", + llm_kv_type=encoding, + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + cpu_cache_token_page_size=8, + enable_cpu_cache=True, + cpu_cache_storage_size=0.001, + linear_att_cache_size=2, + mtp_mode="vanilla_with_att" if mtp_layers else None, + mtp_step=mtp_layers, + ) + if mtp_layers: + draft_config = dict(config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + (draft_dir / "config.json").write_text(json.dumps(draft_config)) + args.mtp_draft_model_dir = [str(draft_dir)] * mtp_layers + draft_cls = Glm5NextMTPModel if kind == "glm" else model_cls + monkeypatch.setattr("lightllm.models.get_draft_model_class", lambda *_: draft_cls) + if encoding in ("fp8kv_sph", "fp8kv_spt"): + per_head = encoding == "fp8kv_sph" + calibration_layers = 2 + mtp_layers + shape = [calibration_layers, 4 if per_head else 2] + calibration = { + "num_head": 2, + "num_layers": calibration_layers, + "scales_shape": shape, + "scales": torch.full(shape, 0.125).tolist(), + "q_calibration": { + "num_head": 2, + "scales_shape": [calibration_layers, 2], + "scales": [[0.25, 0.5]] * calibration_layers, + }, + "qmin": -448, + "qmax": 448, + "quant_type": "per_head" if per_head else "per_tensor", + "architectures": config["architectures"][0], + } + path = tmp_path / "calibration.json" + path.write_text(json.dumps(calibration)) + args.kv_quant_calibration_config_path = str(path) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + calcu_cpu_cache_meta.cache_clear() + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: model_cls) + meta = calcu_cpu_cache_meta() + managers = [] + for rank in range(tp): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + model = model_cls.__new__(model_cls) + model.args, model.config = args, config + model.max_total_token_num, model.max_req_num = 16, 1 + model.tp_world_size_, model.tp_rank_ = tp, rank + prepare_model_states(model) + drafts = [] + for step in range(mtp_layers): + draft_cls = Glm5NextMTPModel if kind == "glm" else model_cls + draft = draft_cls.__new__(draft_cls) + draft.is_mtp_draft_model = True + draft.main_model, draft.mtp_previous_draft_models = model, list(drafts) + draft.args = args + draft.config = dict(config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft.max_total_token_num, draft.max_req_num = 16, 1 + draft.tp_world_size_, draft.tp_rank_ = tp, rank + prepare_model_states(draft, layer_ids=(config["num_hidden_layers"] + step,)) + drafts.append(draft) + model.att_resource_manager.initialize_resources( + args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction + ) + assert not hasattr(model, "mem_manager") + managers.append(model.att_resource_manager) + return meta, managers + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +@pytest.mark.parametrize("tp", [1, 4]) +def test_cpu_service_describes_main_and_draft_as_distinct_resources( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, +): + meta, managers = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, tp, mtp_layers=2) + for resources in managers: + (main,) = resources.get_full_attention_states("main") + (draft,) = resources.get_full_attention_states("draft0") + assert main is not draft + assert set(main.config.layer_ids).isdisjoint(draft.config.layer_ids) + assert main.capacity == draft.capacity == resources.full_att_allocator.size + assert {main.resource_id, draft.resource_id} <= resources.get_cpu_transfer_layouts().keys() + layout_end = max(layout.byte_end for layout in resources.get_cpu_transfer_layouts().values()) + assert (layout_end + 15) // 16 * 16 == meta.page_bytes + assert resources.get_cpu_page_bytes() == meta.page_bytes + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +@pytest.mark.parametrize("tp", [1, 4]) +def test_cpu_service_layout_and_roundtrip_all_encodings( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, +): + meta, managers = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, tp) + # No byte-size formulas duplicated in the test: check the service's actual + # meta-device descriptors against every real model's layout. + for resources in managers: + layout = next(iter(resources.get_cpu_transfer_layouts().values())) + assert 0 <= layout.byte_start < layout.byte_end <= meta.page_bytes + assert meta.page_bytes == resources.get_cpu_page_bytes() + meta.page_num = 3 + client = cpu_cache_factory(meta) + buffer = client.cpu_kv_cache_tensor + expected = [] + for resources in managers: + (full,) = resources.get_full_attention_states("main") + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + field.tensor.fill_(7) + kwargs = {} + if resources.bounded_att_big_page_allocator is not None: + linear = main_linear_state(resources) + linear.conv_state.fill_(3) + linear.ssm_state.fill_(5) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + kwargs = dict(big_page_slot_ids=[slot]) + torch.cuda.synchronize() + expected.append( + { + name: f.tensor[:, :8].contiguous().view(torch.uint8).clone() + for name, f in full_state_test_fields(full, full_state_test_buffers(full)).items() + if "token" in f.axes + } + ) + resources.write_req_to_cpu(torch.arange(8, dtype=torch.long, device="cuda"), buffer, [2], layouts=resources.get_cpu_transfer_layouts(), **kwargs) + torch.cuda.synchronize() + for resources, fields in zip(managers, expected): + (full,) = resources.get_full_attention_states("main") + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, 8:16].zero_() + kwargs = {} + if resources.bounded_att_big_page_allocator is not None: + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + kwargs = dict( + big_page_slot_ids=[slot], + req_idx=1, + ) + resources.read_req_from_cpu(buffer, [2], torch.arange(8, 16, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts(), **kwargs) + if resources.bounded_att_big_page_allocator is not None: + assert (main_linear_state(resources).conv_state[:, 1] == 3).all() + assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() + for name, field in full_state_test_fields(full, full_state_test_buffers(full)).items(): + if "token" in field.axes: + assert torch.equal(field.tensor[:, 8:16].contiguous().view(torch.uint8), fields[name]) + assert not buffer[:2].any() + + +def test_cpu_namespace_isolates_model_revision(runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + args = StartArgs(model_dir=str(tmp_path)) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + (tmp_path / "config.json").write_text("{}") + initial = cpu_cache_namespace() + (tmp_path / "model.safetensors").write_bytes(b"first revision") + updated = cpu_cache_namespace() + assert updated != initial + (tmp_path / "model.safetensors").write_bytes(b"second revision") + assert cpu_cache_namespace() != updated + + +def test_cpu_namespace_includes_draft_revision(runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + args = StartArgs(model_dir=str(tmp_path), mtp_draft_model_dir=[str(draft_dir)]) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + before = cpu_cache_namespace() + (draft_dir / "model.safetensors").write_bytes(b"new draft weights") + assert cpu_cache_namespace() != before + + +@pytest.mark.parametrize("hybrid", [False, True]) +def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], + hybrid: bool, +): + from lightllm.server.multi_level_kv_cache import CacheTier + from lightllm.server.router.model_infer.infer_batch import InferReq + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + kind, encoding, cls = ("qwen", "None", Qwen3NextTpPartModel) if hybrid else ("int4", "int4kv", LlamaTpPartModel) + meta, (resources,) = make_runtime(monkeypatch, tmp_path, kind, encoding, cls, 1) + meta.page_num = 2 + client = cpu_cache_factory(meta) + pages, _ = client.allocate_pages([11, 22], disk_offload_enable=False) + client.cpu_kv_cache_tensor[pages[0]].fill_(123) + client.update_pages_status_to_ready([pages[0]], deref=False) + slots = resources.bounded_att_big_page_allocator.alloc(2) if hybrid else None + write_batches = [] + write_to_cpu = resources.write_req_to_cpu + + def capture_write(*args: object, **kwargs: object) -> None: + write_batches.append((kwargs["cpu_page_indices"], kwargs["cpu_page_readies"])) + write_to_cpu(*args, **kwargs) + + monkeypatch.setattr(resources, "write_req_to_cpu", capture_write) + if hybrid: + # 当前 bounded 页仍由主机复制，测试先准备好源快照。 + # TODO：CPU→CPU kernel 落地后，再覆盖尚未完成 D2H 的快照依赖。 + pool = main_linear_state(resources).big_pages_pool + for slot in slots: + pool.conv_state[slot].fill_(3) + pool.ssm_state[slot].fill_(5) + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + # 延迟当前生产流的 KV 写入，验证搬运流通过 wait_stream 等待。 + torch.cuda._sleep(10_000_000) + (full,) = resources.get_full_attention_states("main") + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + field.tensor.fill_(7) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=16, + cache_tiers=(CacheTier.CPU,), + tail_small_page_buffer_id=None, + hybrid_len_to_big_page_id=dict(zip([8, 16], slots)) if hybrid else {}, + shared_kv_node=None, + cpu_cache_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, + shm_req=SimpleNamespace( + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [8, 16]), + ), + ) + module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) + module.args = get_env_start_args() + module.cpu_cache_client = client + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, + is_hybrid_att_model=hybrid, + dp_world_size=1, + model=SimpleNamespace( + att_resource_manager=resources, + req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(16, device="cuda").view(1, -1)), + ), + ) + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) + monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) + monkeypatch.setattr(common.g_infer_context, "radix_cache", None) + def reject_host_sync(*args: object, **kwargs: object) -> None: + pytest.fail("CPU offload must use stream dependencies instead of host synchronization") + + copy_stream = torch.cuda.Stream() + waited_streams = [] + wait_stream = torch.cuda.Stream.wait_stream + + def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: + if stream == copy_stream: + waited_streams.append(other) + wait_stream(stream, other) + + with monkeypatch.context() as patch: + patch.setattr(torch.cuda.Stream, "synchronize", reject_host_sync) + patch.setattr(torch.cuda.Stream, "wait_stream", record_wait) + task = module.write_req_to_cpu(req, copy_stream) + assert waited_streams == [producer] + task.sync_event.synchronize() + assert write_batches == [(pages, [True, False])] + assert req.cpu_cache_task_status.is_running() + assert (client.cpu_kv_cache_tensor[pages[0]] == 123).all() + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, :8].zero_() + kwargs = {} + if hybrid: + # 写出已完成，释放测试源快照，给读回目标腾出普通大页槽位。 + req.hybrid_len_to_big_page_id.clear() + resources.bounded_att_big_page_allocator.free(slots) + destination = resources.bounded_att_big_page_allocator.alloc(1)[0] + kwargs = dict(big_page_slot_ids=[destination], req_idx=1) + resources.read_req_from_cpu( + client.cpu_kv_cache_tensor, [pages[1]], torch.arange(8, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts(), **kwargs + ) + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + assert (field.tensor[:, :8] == 7).all() + if hybrid: + assert (main_linear_state(resources).conv_state[:, 1] == 3).all() + assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() + + +def test_disk_preserves_resource_pages_and_isolates_new_model_revision( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], +): + import time + + pytest.importorskip("light_mem") + from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker + + meta, (resources,) = make_runtime(monkeypatch, tmp_path, "packed_fp8", "fp8kv_dsa", Deepseek3_2TpPartModel, 1) + meta.page_num = 2 + client = cpu_cache_factory(meta) + pages, _ = client.allocate_pages([11, 22], disk_offload_enable=True) + (full,) = resources.get_full_attention_states("main") + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.random_(0, 256) + resources.write_req_to_cpu(torch.arange(16, dtype=torch.long, device="cuda"), client.cpu_kv_cache_tensor, pages, layouts=resources.get_cpu_transfer_layouts()) + torch.cuda.synchronize() + expected = client.cpu_kv_cache_tensor.clone() + worker = DiskCacheWorker(0.125, client, str(tmp_path / "disk")) + task = worker.service.create(hash_128s=[11, 22], kv_page_indexer=torch.tensor(pages, dtype=torch.int32), mode="w") + deadline = time.monotonic() + 10 + while not task.ready(): + assert time.monotonic() < deadline, "resource page disk write did not finish" + time.sleep(0.01) + assert all(worker.service.query([11, 22])) + client.cpu_kv_cache_tensor.zero_() + assert worker.load_pages([11, 22], pages) + assert torch.equal(client.cpu_kv_cache_tensor, expected) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.zero_() + resources.read_req_from_cpu(client.cpu_kv_cache_tensor, pages, torch.arange(16, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts()) + # 模型版本改变后使用独立磁盘文件，即使页大小及 token 哈希相同也不命中旧数据。 + (tmp_path / "model.safetensors").write_bytes(b"new model revision") + next_meta = dataclasses.replace(meta, cache_namespace=cpu_cache_namespace()) + other_client = cpu_cache_factory(next_meta) + other = DiskCacheWorker(0.125, other_client, str(tmp_path / "disk")) + assert not any(other.service.query([11, 22])) + assert any((tmp_path / "disk").glob(f"cache_file_{meta.cache_namespace}*")) + assert any((tmp_path / "disk").glob(f"cache_file_{next_meta.cache_namespace}*")) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 3111abce81..6e5f7bab37 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -5,16 +5,28 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba -from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.common.att_resources.allocators import FullAttKVAllocator + +from lightllm.common.att_resources.full_att import ( + GQAFullAttentionState, + FullAttentionConfig, +) +from lightllm.common.req_manager import ReqManager +from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule +from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name from lightllm.utils.kv_cache_utils import CpuKVCacheMeta -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] @pytest.fixture(params=[False, True], ids=["default_chunks", "whole_mm_blocks"]) @@ -37,7 +49,7 @@ def linear_cache(monkeypatch, tmp_path, request): monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) get_env_start_args.cache_clear() get_unique_server_name.cache_clear() - config = LinearAttCacheConfig( + config = Qwen3NextResourceConfig( tp_world_size=1, full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, @@ -54,9 +66,20 @@ def linear_cache(monkeypatch, tmp_path, request): ssm_state_dtype=torch.float32, layer_types=("linear_attention", "full_attention"), ) - mem = Qwen3NextMemManager(384, torch.bfloat16, 1, 8, 1, config) - requests = ReqManagerForMamba(2, 512, mem, config) - requests.create_small_page_cache_manager(2) + full = GQAFullAttentionState( + FullAttentionConfig(config.get_full_att_layer_ids(), 1, 8, torch.bfloat16), + page_size=args.page_size, + ) + full.initialize_storage(384) + requests = ReqManager(2, 512) + resources = Qwen3NextAttResourceManager(config, 3, 0) + resources.bind_storage( + full, + FullAttKVAllocator(full.capacity, full.page_size), + big_page_capacity=5, + small_page_capacity=2, + ) + monkeypatch.setattr(g_infer_context, "att_resource_manager", resources) monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(g_infer_context, "req_manager", requests) monkeypatch.setattr( @@ -65,52 +88,57 @@ def linear_cache(monkeypatch, tmp_path, request): SimpleNamespace( get_big_page_ids_by_node=lambda node: [], free_one_small_page_buffer=lambda: None, - small_page_buffers=requests.small_page_buffers, ), ) - yield mem, requests, config + yield full, resources, requests, config get_env_start_args.cache_clear() get_unique_server_name.cache_clear() @pytest.mark.parametrize("token_num", [320, 384]) -def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( +def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( monkeypatch, linear_cache, cpu_cache_factory, token_num ): from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule - mem, requests, config = linear_cache + full, resources, requests, config = linear_cache disable_chunked_mm_input = get_env_start_args().disable_chunked_mm_input first_page_has_state = not disable_chunked_mm_input - big = mem.big_page_buffers + linear = resources.main_gdn_state + big = linear.big_pages_pool + small = linear.small_pages_pool big_states = {} for end in range(128, token_num + 1, 128): - slot = big.alloc_one_state_cache() + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] big_states[end] = slot - for tensor in big.get_state_cache(slot): + for tensor in (big.conv_state[slot], big.ssm_state[slot]): tensor.normal_() tail_slot = None if token_num % 128: - tail_slot = requests.small_page_buffers.alloc_one_state_cache() - for tensor in requests.small_page_buffers.get_state_cache(tail_slot): + tail_slot = resources.bounded_att_small_page_allocator.alloc(1)[0] + for tensor in (small.conv_state[tail_slot], small.ssm_state[tail_slot]): tensor.normal_() - final_states = requests.small_page_buffers.get_state_cache(tail_slot) + final_states = (small.conv_state[tail_slot], small.ssm_state[tail_slot]) else: - final_states = big.get_state_cache(big_states[token_num]) + final_states = (big.conv_state[big_states[token_num]], big.ssm_state[big_states[token_num]]) expected_states = [tensor.clone() for tensor in final_states] - expected_big_states = {end: [t.clone() for t in big.get_state_cache(slot)] for end, slot in big_states.items()} - mem.kv_buffer.normal_() - expected_kv = mem.kv_buffer[:, :token_num].clone() + expected_big_states = { + end: [t.clone() for t in (big.conv_state[slot], big.ssm_state[slot])] for end, slot in big_states.items() + } + full.k_storage.normal_() + full.v_storage.normal_() + expected_k = full.k_storage[:, :token_num].clone() + expected_v = full.v_storage[:, :token_num].clone() client = cpu_cache_factory( - CpuKVCacheMeta(3, 1, 1, 1, config.get_cpu_cache_big_page_bytes(), torch.uint8, 0, torch.uint8), + CpuKVCacheMeta(3, resources.get_cpu_page_bytes()), ) - pages, _ = client.allocate_pages([11, 22, 33], disk_offload_enable=False) indexes = torch.arange(token_num, dtype=torch.int32, device="cuda") - page_ids = torch.arange(3, dtype=torch.int32, device="cuda") req = SimpleNamespace( req_idx=0, + cache_tiers=(CacheTier.CPU,), + cpu_cache_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, cur_kv_len=token_num, hold_kv_len=token_num, hybrid_cache_len=token_num, @@ -123,27 +151,40 @@ def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( input_len=token_num + 1, disk_prompt_cache_len=0, cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: pages), + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22, 33]), token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256, token_num]), ), ) - mem.operator.offload_gpu_kv_to_cpu_cache(indexes, page_ids, torch.zeros_like(page_ids), client, req) - torch.cuda.synchronize() + # 通过接收 InferReq 的生产入口完成页来源查询、CPU 分配与快照写出。 + requests.req_to_token_indexs[0, :token_num] = indexes module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = get_env_start_args() module.cpu_cache_client = client + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, + model=SimpleNamespace(req_manager=requests, att_resource_manager=resources), + ) + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) + task = module.write_req_to_cpu(req, torch.cuda.Stream()) + assert task is not None + task.sync_event.synchronize() + pages = task.page_indexes.tolist() client.update_pages_status_to_ready(pages, deref=False) - big.free_state_cache(list(big_states.values())) + resources.bounded_att_big_page_allocator.free(list(big_states.values())) if tail_slot is not None: - requests.small_page_buffers.free_state_cache([tail_slot]) + resources.bounded_att_small_page_allocator.free([tail_slot]) req.hybrid_len_to_big_page_id = {} req.tail_small_page_buffer_id = None req.cur_kv_len = req.hold_kv_len = 0 - mem.kv_buffer.zero_() - big.conv_state_cache.buffer.zero_() - big.ssm_state_cache.buffer.zero_() + full.k_storage.zero_() + full.v_storage.zero_() + big.conv_state.zero_() + big.ssm_state.zero_() def allocate(req, count): - allocated = mem.alloc(count) + allocated = resources.full_att_allocator.alloc(count) requests.req_to_token_indexs[req.req_idx, :count] = allocated.cuda() req.hold_kv_len = count return allocated @@ -152,7 +193,8 @@ def allocate(req, count): module.backend = SimpleNamespace( is_master_in_dp=True, _alloc_req_kv_mem=allocate, - model=SimpleNamespace(mem_manager=mem, req_manager=requests), + model=SimpleNamespace(req_manager=requests, att_resource_manager=resources), + is_hybrid_att_model=True, ) module.init_sync_group = None module.need_sync_compute_stream = lambda: False @@ -161,20 +203,23 @@ def allocate(req, count): module.load_cpu_cache_to_reqs([req]) assert req.cur_kv_len == req.hold_kv_len == token_num - assert torch.equal(mem.kv_buffer[:, :token_num], expected_kv) + restored_indices = requests.req_to_token_indexs[req.req_idx, :token_num] + assert torch.equal(full.k_storage[:, restored_indices], expected_k) + assert torch.equal(full.v_storage[:, restored_indices], expected_v) assert sorted(req.hybrid_len_to_big_page_id) == list(range(128, token_num + 1, 128)) - # Even the invalid page is copied normally; only restoring from it is forbidden. + # Every published CPU page has a complete, length-checked linear state. for end, slot in req.hybrid_len_to_big_page_id.items(): - for actual, expected in zip(big.get_state_cache(slot), expected_big_states[end]): + for actual, expected in zip((big.conv_state[slot], big.ssm_state[slot]), expected_big_states[end]): torch.testing.assert_close(actual, expected) - torch.testing.assert_close(requests.req_to_conv_state.buffer[:, 0].cpu(), expected_states[0]) - torch.testing.assert_close(requests.req_to_ssm_state.buffer[:, 0].cpu(), expected_states[1]) + torch.testing.assert_close(linear.conv_state[:, 0].cpu(), expected_states[0]) + torch.testing.assert_close(linear.ssm_state[:, 0].cpu(), expected_states[1]) - # Loaded intermediate pages must remain unusable as radix restore points. - tree = HybridAttPagedRadixCache(384, 103, 64, 2, mem, requests.small_page_buffers) + # Model data validity and multimodal resume boundaries are separate checks. + tree = HybridAttPagedRadixCache(384, 103, 64, 2, resources) cache_len = max(req.hybrid_len_to_big_page_id) tree.insert( torch.arange(cache_len), + value=restored_indices[:cache_len].cpu(), block_hashs=list(range(1, cache_len // 64 + 1)), block_state_idxs=[None] * (cache_len // 64), len_to_big_page_id=req.hybrid_len_to_big_page_id, @@ -195,3 +240,78 @@ def allocate(req, count): mm_input_spans=req.mm_input_spans, ) tree.dec_node_ref_counter(req.shared_kv_node) + + +@pytest.mark.parametrize("shared_prefix", [False, True]) +@pytest.mark.parametrize("allocated_pages", [0, 1, 2]) +def test_offload_resolves_only_allocated_pages_and_requires_big_slots( + monkeypatch, linear_cache, shared_prefix, allocated_pages, +): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + full, resources, requests, config = linear_cache + slots = resources.bounded_att_big_page_allocator.alloc(allocated_pages) if allocated_pages else [] + shared_slots = [slots[0]] if shared_prefix and slots else [] + monkeypatch.setattr(g_infer_context, "radix_cache", SimpleNamespace(get_big_page_ids_by_node=lambda node: shared_slots)) + big_slots = {128 * (i + 1): slot for i, slot in enumerate(slots) if not shared_prefix or i > 0} + req = SimpleNamespace( + req_idx=0, shared_kv_node=None, hybrid_len_to_big_page_id=big_slots, + tail_small_page_buffer_id=None, cur_kv_len=256, cache_tiers=(CacheTier.CPU,), + shm_req=SimpleNamespace( + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), + ), + ) + requested_hashes = [] + writes = [] + + def allocate_pages(hashes, *, disk_offload_enable): + requested_hashes.append(list(hashes)) + return list(range(allocated_pages)) + [-1] * (len(hashes) - allocated_pages), [False] * len(hashes) + + def write_to_cpu(token_indices, cpu_page_buffers, cpu_page_indices, **kwargs): + writes.append((token_indices.numel(), cpu_page_indices, kwargs["big_page_slot_ids"], kwargs["small_page_slot_ids"])) + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = get_env_start_args() + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, + model=SimpleNamespace(att_resource_manager=resources, req_manager=requests), + ) + module.cpu_cache_client = SimpleNamespace( + lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), + allocate_pages=allocate_pages, cpu_kv_cache_tensor=torch.empty((2, 1), dtype=torch.uint8), + ) + monkeypatch.setattr(resources, "write_req_to_cpu", write_to_cpu) + broadcasts = [] + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda data, **kwargs: broadcasts.append(data)) + task = module.write_req_to_cpu(req, torch.cuda.Stream()) + if allocated_pages == 0: + assert task is None + assert req.cpu_cache_task_status == InferReq._CpuCacheTaskStatus.FINISHED + assert requested_hashes == [[11, 22]] + assert broadcasts == [[0]] + assert writes == [] + return + task.sync_event.synchronize() + assert broadcasts == [[{ + "item_size": allocated_pages, + "page_list": list(range(allocated_pages)), + "ready_list": [False] * allocated_pages, + }]] + # 即使 CPU 只分配一页，也申请完整可缓存范围；未选中的 256 位置不要求提供槽位。 + assert requested_hashes == [[11, 22]] + assert writes == [(allocated_pages * 128, list(range(allocated_pages)), slots, [])] + + # 实际选中页缺失时直接暴露错误，不缩短写出范围，也不提交部分搬运。 + missing_length = allocated_pages * 128 + if shared_prefix and allocated_pages == 1: + shared_slots.clear() + else: + del req.hybrid_len_to_big_page_id[missing_length] + with pytest.raises(KeyError) as error: + module.write_req_to_cpu(req, torch.cuda.Stream()) + assert error.value.args == (missing_length,) + assert requested_hashes == [[11, 22], [11, 22]] + assert len(writes) == 1 diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index f40dd1b98c..c7ff8758b8 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -32,6 +32,50 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): req.cur_kv_len = end +@pytest.mark.usefixtures("kv_allocator_environment") +def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): + from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context + from unit_tests.common.att_resources.test_manager_helpers import composition, init_source + + resources = composition() + slots = init_source(resources) + tail = resources.main_indexer_kpool_c4_tail_state + resources.reset_bounded_att_state([1]) + assert (tail.tail_buffer[:, 1] == 9).all() + resources.restore_bounded_att_pages([1], page_kind="big", slot_ids=[slots[-1]]) + for state in resources.get_bounded_attention_states(): + state.mtp_state_index.fill_(2) + + monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) + monkeypatch.setattr(g_infer_context, "radix_cache", None) + monkeypatch.setattr(g_infer_context, "att_resource_manager", resources) + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace( + linear_att_hash_page_size=4, linear_att_page_block_num=2, max_req_total_len=16, page_size=4 + ) + req.req_idx = 1 + req.cur_kv_len = req.cur_output_len = 0 + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.shm_req = SimpleNamespace( + input_len=1, hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [], size=0) + ) + req._hybrid_match_radix_cache() + + for state in resources.get_bounded_attention_states(): + assert not state.conv_state[:, 1].any() + assert not state.ssm_state[:, 3:6].any() + assert state.mtp_state_index.tolist() == [2, 0, 2] + assert state.conv_state[:, 0].any() + pool = state.big_pages_pool + expected_conv, expected_ssm = (23, 47) if state is resources.main_kda_state else (24, 48) + assert (pool.conv_state[slots[-1]] == expected_conv).all() + assert (pool.ssm_state[slots[-1]] == expected_ssm).all() + assert not tail.tail_buffer[:, 1].any() + assert (tail.tail_buffer[:, 0] == 3).all() + assert resources.bounded_att_big_page_allocator.free_count == 1 + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 + + @pytest.mark.parametrize( "input_len,big_page_num,spans,expected_big,expected_small", [ @@ -74,23 +118,15 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( context = InferenceContext() context.args = req.args context.is_hybrid_att_model = True - context.radix_cache = SimpleNamespace( - big_page_buffers=SimpleNamespace(alloc_one_state_cache=lambda: next(slots)), - small_page_buffers=SimpleNamespace(alloc_one_state_cache=lambda: 7), - free_one_small_page_buffer=lambda: None, + context.radix_cache = SimpleNamespace(free_one_small_page_buffer=lambda: None) + context.att_resource_manager = SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: [next(slots)]), + bounded_att_small_page_allocator=SimpleNamespace(alloc=lambda count: [7]), + save_bounded_att_pages=lambda reqs, page_kind, slot_ids: ( + big_snapshots if page_kind == "big" else small_snapshots + ).append(req.get_chuncked_input_token_len()), ) - def save_big_page_states(b_req_idx, reqs, buffer_indexes): - assert b_req_idx.tolist() == [req.req_idx for req in reqs] - for req, buffer_idx in zip(reqs, buffer_indexes): - if buffer_idx != -1: - big_snapshots.append(req.get_chuncked_input_token_len()) - - def save_state(req, buffer_idx, state_cache_manager): - small_snapshots.append(req.get_chuncked_input_token_len()) - - context.req_manager = SimpleNamespace(save_big_page_states=save_big_page_states, save_state=save_state) - while req.cur_kv_len < req.get_cur_total_len(): context.save_hybrid_state_to_cache(torch.tensor([0]), [req]) req.cur_kv_len = req.get_chuncked_input_token_len() @@ -144,7 +180,7 @@ def test_full_attention_gpu_restore_stays_outside_image( assert radix_cache.get_refed_tokens_num() == gpu_end hashes = compute_token_list_hash(source, 256) - client = cpu_cache_factory(CpuKVCacheMeta(len(hashes), 256, 1, 1, 8, torch.bfloat16, 0, torch.bfloat16)) + client = cpu_cache_factory(CpuKVCacheMeta(len(hashes), 4096)) pages, _ = client.allocate_pages(hashes, False) client.update_pages_status_to_ready(pages) manager = MultiLevelKVCacheManager.__new__(MultiLevelKVCacheManager) @@ -200,8 +236,8 @@ def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disabl req.multimodal_params = SimpleNamespace( to_dict=lambda: { "images": [ - {"start_idx": 450, "token_id": 450, "token_num": 150}, - {"start_idx": 700, "token_id": 700, "token_num": 200}, + {"start_idx": 450, "token_num": 150}, + {"start_idx": 700, "token_num": 200}, ], "audios": [], } @@ -243,7 +279,7 @@ def test_cpu_cache_match_and_load_use_a_complete_restore_point( from lightllm.utils.kv_cache_utils import CpuKVCacheMeta ends = list(range(2048, load_end + 1, 2048)) - client = cpu_cache_factory(CpuKVCacheMeta(max(1, len(ends)), 1, 1, 1, 16, torch.uint8, 0, torch.uint8)) + client = cpu_cache_factory(CpuKVCacheMeta(max(1, len(ends)), 16)) page_ids, _ = client.allocate_pages(list(range(10, 10 + len(ends))), disk_offload_enable=False) client.update_pages_status_to_ready(page_ids, deref=False) loaded = [] @@ -267,16 +303,17 @@ def allocate(req, count): return torch.arange(req.cur_kv_len, req.cur_kv_len + count, dtype=torch.int32) module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace(page_size=256) + module.args = SimpleNamespace(page_size=256, cpu_cache_token_page_size=256) module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=hybrid, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), - mem_manager=SimpleNamespace( - operator=SimpleNamespace( - load_cpu_cache_to_gpu=lambda **kwargs: loaded.append(req.cur_kv_len), - ) + att_resource_manager=SimpleNamespace( + get_cpu_transfer_layouts=lambda: {}, + bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), + read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: loaded.append(req.hold_kv_len), ), ), ) @@ -288,6 +325,7 @@ def allocate(req, count): monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) + req.hybrid_len_to_big_page_id = {} module.load_cpu_cache_to_reqs([req]) assert req.cur_kv_len == req.hold_kv_len == expected assert loaded == ([expected] if expected > loaded_start else []) diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py new file mode 100644 index 0000000000..3ffa28da55 --- /dev/null +++ b/unit_tests/server/router/test_kv_capacity.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.att_resources.allocators import full_att_kv_allocator + + +pytestmark = pytest.mark.usefixtures("kv_allocator_environment") + + +@pytest.mark.parametrize( + "tp,dp,nnodes,selected_ranks", + [(4, 1, 1, [0]), (8, 4, 1, [0, 2, 4, 6]), (8, 4, 2, [0, 2]), (8, 8, 2, [0, 1, 2, 3]), (8, 1, 2, [0])], +) +@pytest.mark.parametrize("reader_first", [False, True]) +def test_router_reads_each_local_dp_groups_allocator( + monkeypatch: pytest.MonkeyPatch, + tp: int, + dp: int, + nnodes: int, + selected_ranks: list[int], + reader_first: bool, +): + monkeypatch.setattr(full_att_kv_allocator, "get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=dp, nnodes=nnodes)) + reader = full_att_kv_allocator.KVCapacityReader() if reader_first else None + allocators = [] + for rank in range(tp // nnodes): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", str(rank)) + allocator = FullAttKVAllocator(16 + 4 * rank, 4, pin_memory=False) + assert allocator.shared_can_use_token_num.get_value() == allocator.size + allocator.alloc(4) + allocators.append(allocator) + reader = reader or full_att_kv_allocator.KVCapacityReader() + assert len(reader.counters) == len(selected_ranks) + for dp_rank, rank in enumerate(selected_ranks): + allocator = allocators[rank] + initial = allocator.size - 4 + assert reader.get_available_tokens(dp_rank) == initial + tokens = allocator.alloc(4) + assert reader.get_available_tokens(dp_rank) == initial - 4 + allocator.free(tokens) + assert reader.get_available_tokens(dp_rank) == initial + allocator.free_all() + assert reader.get_available_tokens(dp_rank) == allocator.size + + +def test_failed_page_operations_leave_shared_capacity_unchanged(): + allocator = FullAttKVAllocator(16, 4, pin_memory=False) + counter = allocator.shared_can_use_token_num + tokens = allocator.alloc(16) + with pytest.raises(AssertionError): + allocator.alloc(4) + with pytest.raises(AssertionError): + allocator.free(tokens[:3]) + assert counter.get_value() == allocator.can_use_mem_size == 0 + allocator.free(tokens) + assert counter.get_value() == allocator.can_use_mem_size == 16 diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py index 1e63550f45..800f0cdda3 100644 --- a/unit_tests/server/test_dp_request_capacity.py +++ b/unit_tests/server/test_dp_request_capacity.py @@ -45,7 +45,6 @@ def test_reject_invalid_global_capacity(capacity): def _make_dp_queue(monkeypatch, balancer): monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) - monkeypatch.setattr("lightllm.server.router.req_queue.base_queue.get_fixed_kv_len", lambda: 0) args = StartArgs( running_max_req_size=32, dp=8, diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 6b57270125..56433d1b3b 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -20,7 +20,7 @@ ) -class FakeMemManager: +class FakeKVAllocator: def __init__(self, page_size: int): self.page_size = page_size self.next_index = 0 @@ -35,15 +35,14 @@ def alloc(self, size: int) -> torch.Tensor: def make_executor(page_size: int = 4): - mem_manager = FakeMemManager(page_size=page_size) + allocator = FakeKVAllocator(page_size=page_size) req_manager = SimpleNamespace( - mem_manager=mem_manager, req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), ) - model = SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager) + model = SimpleNamespace(att_resource_manager=SimpleNamespace(full_att_allocator=allocator), req_manager=req_manager) args = SimpleNamespace(mtp_mode=None, mtp_step=0) executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) - return executor, mem_manager, req_manager + return executor, allocator, req_manager def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): @@ -53,7 +52,7 @@ def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): def test_prefill_allocates_and_reuses_complete_pages_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) + executor, allocator, req_manager = make_executor(page_size=4) req_idx = torch.tensor([0, 1], dtype=torch.int32) model_input = executor._make_prefill_input( @@ -62,7 +61,7 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): ready_cache_len=0, ) - assert mem_manager.alloc_sizes == [4, 4] + assert allocator.alloc_sizes == [4, 4] assert model_input.b_is_decode_req.tolist() == [False, False] assert executor._hold_kv_len_by_req == {0: 4, 1: 4} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :4], page_size=4) @@ -77,7 +76,7 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): input_ids=torch.tensor([1, 2], dtype=torch.int64), max_kv_seq_len=4, ) - assert mem_manager.alloc_sizes == [4, 4] + assert allocator.alloc_sizes == [4, 4] # 跨过页边界后，每个请求分别追加一整页。 executor._make_decode_input( @@ -88,14 +87,14 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): input_ids=torch.tensor([3, 4], dtype=torch.int64), max_kv_seq_len=5, ) - assert mem_manager.alloc_sizes == [4, 4, 4, 4] + assert allocator.alloc_sizes == [4, 4, 4, 4] assert executor._hold_kv_len_by_req == {0: 8, 1: 8} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) def test_mtp_expanded_rows_allocate_once_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) + executor, allocator, req_manager = make_executor(page_size=4) executor.args.mtp_mode = "eagle_with_att" executor.args.mtp_step = 3 @@ -110,7 +109,7 @@ def test_mtp_expanded_rows_allocate_once_per_request(): ) assert model_input.batch_size == 4 - assert mem_manager.alloc_sizes == [8, 8] + assert allocator.alloc_sizes == [8, 8] assert executor._hold_kv_len_by_req == {0: 8, 1: 8} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) diff --git a/unit_tests/utils/test_config_utils_fixed_kv.py b/unit_tests/utils/test_config_utils_fixed_kv.py deleted file mode 100644 index 116e299e24..0000000000 --- a/unit_tests/utils/test_config_utils_fixed_kv.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.utils import config_utils - - -@pytest.fixture(autouse=True) -def clear_fixed_kv_len_cache(): - config_utils.get_fixed_kv_len.cache_clear() - yield - config_utils.get_fixed_kv_len.cache_clear() - - -def test_fixed_kv_len_supports_page_size_one(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=1, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 3 - - -def test_fixed_kv_len_uses_only_complete_pages(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=2, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 2 - - -def test_no_fixed_kv_supports_paged_kv(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=4, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {}) - - assert config_utils.get_fixed_kv_len() == 0 From fa53ba4f6ebc637d59172e40606a850f242c0974 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 01:51:33 +0000 Subject: [PATCH 02/56] refactor(att_resources): distinguish logical layers from shared storage Group logical full-attention layers by physical storage for allocation, calibration and CPU/PD layouts. Bind Gemma4 shared KV layers through the State mapping. Clarify logical layer IDs and physical storage indices across State interfaces, model setup, tests and design documentation. --- docs/cache_manager_design.md | 40 ++++-- .../att_resources/att_tail/base_state.py | 2 +- .../att_resources/att_tail/kpool_state.py | 21 +-- .../att_resources/bounded_att/base_state.py | 2 +- .../bounded_att/conv_ssm_types.py | 4 +- .../bounded_att/gdn_pages_pool.py | 4 +- .../att_resources/bounded_att/gdn_state.py | 25 ++-- .../bounded_att/kda_pages_pool.py | 6 +- .../att_resources/bounded_att/kda_state.py | 25 ++-- .../att_resources/full_att/base_state.py | 33 ++++- .../full_att/create_gqa_state.py | 4 +- .../full_att/create_mla_state.py | 2 +- .../att_resources/full_att/gqa_state.py | 27 ++-- .../att_resources/full_att/gqa_transfer.py | 6 +- .../full_att/group_quantized_state.py | 33 +++-- .../full_att/group_quantized_transfer.py | 6 +- .../att_resources/full_att/mla_state.py | 23 +-- .../att_resources/full_att/mla_transfer.py | 6 +- .../full_att/packed_mla_state.py | 25 ++-- .../full_att/packed_mla_transfer.py | 6 +- .../full_att/static_fp8_state.py | 33 +++-- .../full_att/static_fp8_transfer.py | 6 +- .../common/basemodel/attention/base_att.py | 9 +- .../common/basemodel/attention/fa3/fp8.py | 18 +-- .../basemodel/attention/flashinfer/fp8.py | 12 +- .../common/basemodel/attention/linear/gdn.py | 4 +- .../common/basemodel/attention/linear/kda.py | 4 +- lightllm/common/basemodel/basemodel_utils.py | 59 +++++--- .../transformer_layer_infer_template.py | 2 +- .../layer_infer/transformer_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 8 +- .../layer_infer/transformer_layer_infer.py | 16 +-- lightllm/models/gemma4/model.py | 15 ++ .../models/glm5_next/att_resource_manager.py | 45 +++--- lightllm/models/glm5_next/indexer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 4 +- lightllm/models/glm5_next/model.py | 12 +- lightllm/models/glm5_next_mtp/model.py | 2 +- .../layer_infer/transformer_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 2 +- .../models/qwen3next/att_resource_config.py | 2 +- .../models/qwen3next/att_resource_manager.py | 18 +-- lightllm/models/qwen3next/model.py | 12 +- lightllm/utils/att_resource_config.py | 75 +++++----- test/cpu_cache_kernel/test_speed.py | 2 +- .../common/att_resources/model_helpers.py | 20 +-- .../att_resources/test_backend_startup.py | 4 +- .../att_resources/test_bounded_state.py | 12 +- .../common/att_resources/test_capacity.py | 40 +++--- .../att_resources/test_compute_binding.py | 4 +- .../att_resources/test_fp8_compute_binding.py | 6 +- .../att_resources/test_full_att_cpu_kernel.py | 2 +- .../common/att_resources/test_full_state.py | 6 +- .../test_full_storage_initialization.py | 4 +- .../att_resources/test_gqa_pd_kernel.py | 6 +- .../common/att_resources/test_gqa_transfer.py | 2 +- .../common/att_resources/test_layer_groups.py | 133 ++++++++++++++++++ .../att_resources/test_model_calibration.py | 2 +- .../test_model_initialization.py | 10 +- .../att_resources/test_model_resources.py | 4 +- .../test_model_storage_definition.py | 19 ++- .../att_resources/test_no_attention_draft.py | 2 +- .../att_resources/test_page_pool_storage.py | 4 +- .../test_quantized_full_state.py | 6 +- .../test_quantized_paged_storage.py | 2 +- .../att_resources/test_state_creation.py | 12 +- .../att_resources/test_state_identity.py | 4 +- .../common/att_resources/test_tail_state.py | 2 +- .../common/att_resources/test_tp_config.py | 12 +- unit_tests/models/glm5_next/test_cache.py | 4 +- unit_tests/models/glm5_next/test_kernels.py | 5 +- unit_tests/models/glm5_next/test_layers.py | 4 +- unit_tests/models/glm5_next/test_mtp.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 25 ++-- .../model_infer/test_cpu_resource_layout.py | 6 +- .../test_hybrid_cpu_state_cache.py | 2 +- 78 files changed, 653 insertions(+), 393 deletions(-) create mode 100644 unit_tests/common/att_resources/test_layer_groups.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 9ec820c1b6..9df001f10a 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -332,7 +332,7 @@ lightllm/models// `KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_pages`、`restore_bounded_att_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 -`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；CPU meta 定义入口为 `create_bounded_attention_definition(...)`。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 +`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；CPU meta 定义入口为 `create_bounded_attention_definition(...)`。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 `FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；其余复制视图函数返回源、目标视图列表，由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 @@ -430,6 +430,24 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttResourceManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 +#### 逻辑层与物理存储层 + +`FullAttentionConfig.logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 + +```python +logical_layer_groups = ( + (0, 4), # 逻辑层 0、4 共用物理层 0。 + (8,), # 逻辑层 8 使用物理层 1。 +) +# 不共享时使用 ((0,), (4,), (8,))，对应三份物理层存储。 +``` + +层命名统一使用 `logical_` 和 `storage_` 区分含义：`logical_layer_groups` 描述共享分组，`logical_layer_id` 是推理层传入的逻辑层号，`storage_layer_index` 是 Tensor layer 轴上的物理下标。State 的 `_logical_to_storage_layer` 保存两者的对应关系。KDA/GDN 和 K-pool 当前仍是一对一存储，配置使用 `logical_layer_ids`；三类 State 的 `get_layer_view(logical_layer_id=...)` 使用相同参数语义。FP8 标定的 `storage_layer_count`、`storage_layer_start` 表示物理行数和偏移，不能传逻辑层号。Tensor 轴标签 `"layer"` 保留，表示存储层轴。 + +State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get_layer_view(logical_layer_id=0)` 与 `get_layer_view(logical_layer_id=4)` 返回同一份存储视图；`logical_layer_id=8` 返回另一份。分配、单 token 字节数、CPU/PD 布局和静态量化标定行数均按 `len(logical_layer_groups)` 计算，不按组内逻辑层总数计算，不重复保存物理层数或另一份配置映射。 + +同组逻辑层共享该 State 的全部配套字段，必须使用兼容的布局和参数。分组只表达数据共享关系，不指定写入者或执行顺序；模型负责跳过共享消费者的 KV 写入，避免覆盖生产层数据。普通模型每个逻辑层单独一组；Gemma4 将后部共享层加入同类型的最后一个非共享层所在组。推理层按自己的逻辑层号访问，CPU 服务与模型使用同一套配置构建物理分组。DP/CPU/PD 传输按相同的物理层顺序搬运，main/draft 仍各自拥有独立 State。 + #### 继承与内部 storage 的关系 建议以 `FullAttentionState` 作为公共接口，由具体布局子类实现差异。子类内部的 storage 持有 buffer 和布局信息，operator 可以复用兼容的写入、量化及复制 kernel。布局选择集中在具体子类及其配套实现，公共调用方使用稳定接口。 @@ -476,10 +494,10 @@ class FullAttentionState: # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> ViewT: + def get_layer_view(self, logical_layer_id: int) -> ViewT: ... - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: WriteT) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: WriteT) -> None: ... def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -541,8 +559,8 @@ class FullAttentionState: | 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | | `initialize_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | | `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | -| `get_layer_view(layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | -| `write_full_att_kv(layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | +| `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | +| `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | | `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | 容量、device、层映射及布局标识应从同一份稳定配置和描述中取得，避免多处维护互相矛盾的副本。配置相同且只改变容量时，容量计算和真实 buffer 创建应覆盖相同的数据组成。共享底层内存按实际分配去重记账。 @@ -607,7 +625,7 @@ manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字 State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 显式接收必需的 `scales` 和可选的 `q_scales`；meta 阶段保留真实标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 -`layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 +`logical_layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 混合模型将实际 full-attention 逻辑层号及追加的 draft 层号写入 State 配置；例如逻辑层 3、7、8 可以分别对应物理行 0、1、2。计算侧直接以逻辑层号调用 `write_full_att_kv` 和 `get_layer_view`，由 State 完成映射，CPU/PD 描述复用同一份层标识。KV/Q 标定按物理存储行组织，backend 必须使用对应行的参数。 @@ -1316,7 +1334,7 @@ class BoundedAttentionState(ABC): # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== @abstractmethod - def get_layer_view(self, layer_id: int) -> Any: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... @@ -1394,7 +1412,7 @@ CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_pag | --- | --- | | 初始化 | 构造函数按 `request_capacity` 和 `mtp_step` 创建运行态及层映射，大小页池先为 `None`；不创建 allocator，不分配请求号或页槽位 | | `initialize_page_pools(...)` | 容量规划完成后创建大小页池，覆盖共享 allocator 的槽位域；容量为零的页池保持 `None`，不重建运行态或重置 MTP 游标 | -| `get_layer_view(layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | +| `get_layer_view(logical_layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | | `reset_bounded_att_state(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | | `commit_bounded_att_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | @@ -1499,7 +1517,7 @@ class AttentionTailState(ABC): # 二、Attention tail：计算视图与无快照重置 # ======================================================================== @abstractmethod - def get_layer_view(self, layer_id: int) -> Any: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... @@ -1528,7 +1546,7 @@ class AttentionTailState(ABC): | 接口 | 语义 | | --- | --- | | 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | -| `get_layer_view(layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | +| `get_layer_view(logical_layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | | `reset_attention_tail_state(req_indices)` | 清零指定请求的 tail，适用于新请求或可重置的缓存前缀；K-pool 清零完整 ring，其他实现清零对应数据与控制信息。模型配置和缓存流程负责保证边界，接口不接收长度 | `get_layer_view` 供计算 backend 直接使用；池化与写 tail 仍在模型 kernel 中完成，不增加通用的 `append_token` 或 attention forward 包装。执行顺序必须满足“先池化、再更新 tail”。 @@ -1579,7 +1597,7 @@ BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。分配完成后的 backend 绑定、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方可按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 +单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_storage_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 `AttResourceManager.initialize_resources()` 由基类提供通用实现，保留三个阶段：标记 State 创建完成、估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/att_resources/att_tail/base_state.py index 1e65380fb6..33add9cc6f 100644 --- a/lightllm/common/att_resources/att_tail/base_state.py +++ b/lightllm/common/att_resources/att_tail/base_state.py @@ -29,7 +29,7 @@ def get_state_type(self) -> str: # ======================================================================== @abstractmethod - def get_layer_view(self, layer_id: int) -> Any: + def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回尾部状态计算所需的存储视图。""" ... diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py index 4106623a67..10410bdb76 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -15,8 +15,8 @@ @dataclass(frozen=True) class KPoolStateConfig: - # 本 State 管理的逻辑层号；顺序对应 tail_buffer 的 layer 轴，不要求编号连续。 - layer_ids: tuple[int, ...] + # 本 State 管理的逻辑层号；当前每层独占一份 tail，列表顺序对应物理层轴，不要求编号连续。 + logical_layer_ids: tuple[int, ...] # K-pool 的池化步长/压缩比例，单位为原始 token；例如 4 表示每 4 个 token 形成一个池化组。 # 无快照前缀恢复要求长度对齐该步长；ring 的实际槽位数为 pool_size + State.mtp_step。 pool_size: int @@ -34,7 +34,7 @@ class KPoolStateConfig: def __post_init__(self) -> None: """构造配置时检查维度、层号和 TP 范围。""" assert min(self.pool_size, self.head_dim) > 0, "invalid K-pool dimensions" - assert self.layer_ids and len(set(self.layer_ids)) == len(self.layer_ids), ( + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids), ( "tail layers must be nonempty and unique" ) assert 0 <= self.tp_rank < self.tp_world_size, "invalid TP rank" @@ -61,10 +61,13 @@ def __init__( self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step # ring 额外保留的 MTP 历史槽位数。 self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } self.tail_buffer = torch.zeros( ( - len(config.layer_ids), + len(config.logical_layer_ids), request_capacity, config.pool_size + mtp_step, 2 * config.head_dim, @@ -81,9 +84,9 @@ def get_state_type(self) -> str: # 二、Attention tail：计算视图与无快照重置 # ======================================================================== - def get_layer_view(self, layer_id: int) -> torch.Tensor: + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """返回该层所有请求的原生 ring 视图，末轴包含 key 和 gate。""" - return self.tail_buffer[self._layer_mapping[layer_id]] + return self.tail_buffer[self._logical_to_storage_layer[logical_layer_id]] def reset_attention_tail_state( self, @@ -106,10 +109,10 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: tail_region = TransferRegion( field="tail", byte_start=0, - byte_end=len(config.layer_ids) * window_size * channels * config.dtype.itemsize, + byte_end=len(config.logical_layer_ids) * window_size * channels * config.dtype.itemsize, dtype=config.dtype, axes=("layer", "window", "channel"), - shape=(len(config.layer_ids), window_size, channels), + shape=(len(config.logical_layer_ids), window_size, channels), axis_slices=( AxisSlice(axis="window", src_start=0, dst_start=0, length=window_size), AxisSlice(axis="channel", src_start=0, dst_start=0, length=channels), diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/att_resources/bounded_att/base_state.py index 461f53b9ab..3756fe2e82 100644 --- a/lightllm/common/att_resources/bounded_att/base_state.py +++ b/lightllm/common/att_resources/bounded_att/base_state.py @@ -44,7 +44,7 @@ def initialize_page_pools( # ======================================================================== @abstractmethod - def get_layer_view(self, layer_id: int) -> Any: + def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回运行态和算法控制信息的计算视图。""" ... diff --git a/lightllm/common/att_resources/bounded_att/conv_ssm_types.py b/lightllm/common/att_resources/bounded_att/conv_ssm_types.py index 4858ed11eb..aa0666eb3b 100644 --- a/lightllm/common/att_resources/bounded_att/conv_ssm_types.py +++ b/lightllm/common/att_resources/bounded_att/conv_ssm_types.py @@ -9,8 +9,8 @@ class ConvSSMStateConfig: """KDA/GDN 的层布局、维度、dtype 和 TP 配置；不包含模型角色或资源身份。""" - # 本 State 管理的逻辑层编号；顺序对应 Conv/SSM Tensor 的物理层轴，编号不能重复。 - layer_ids: tuple[int, ...] + # 本 State 管理的逻辑层编号；当前每个逻辑层独占一份运行态，列表顺序对应物理层轴，编号不能重复。 + logical_layer_ids: tuple[int, ...] # TP 切分前的全局 K head 数，不计 TP 副本；Q 使用相同的 head 数和维度。 global_k_heads: int # TP 切分前的全局 V head 数；GDN 可与 K 不同，KDA 要求两者相等。 diff --git a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py index 268859db57..1b1c0745c8 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py +++ b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py @@ -23,13 +23,13 @@ def __init__( device = torch.device(device) pin_memory = pin_memory and device.type == "cpu" self.conv_state = create_page_pool_tensor( - size=(size, len(config.layer_ids), config.conv_channels, config.conv_width), + size=(size, len(config.logical_layer_ids), config.conv_channels, config.conv_width), dtype=config.conv_dtype, device=device, pin_memory=pin_memory, ) self.ssm_state = create_page_pool_tensor( - size=(size, len(config.layer_ids), config.num_v_heads, config.key_dim, config.value_dim), + size=(size, len(config.logical_layer_ids), config.num_v_heads, config.key_dim, config.value_dim), dtype=config.ssm_dtype, device=device, pin_memory=pin_memory, diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py index 5c684b794d..624d961134 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -40,15 +40,18 @@ def __init__( self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step - self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} - num_layers = len(config.layer_ids) + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) self.conv_state = torch.zeros( - (num_layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), dtype=config.conv_dtype, device=device, ) self.ssm_state = torch.zeros( - (num_layers, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + (storage_layer_count, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), dtype=config.ssm_dtype, device=device, ) @@ -86,10 +89,10 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== - def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """返回该层 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - layer = self._layer_mapping[layer_id] - return self.conv_state[layer], self.ssm_state[layer] + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] def reset_bounded_att_state(self, req_indices: List[int]) -> None: """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" @@ -305,9 +308,9 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 config = self.config - layer_count = len(config.layer_ids) + storage_layer_count = len(config.logical_layer_ids) conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (layer_count, conv_channels, config.conv_width) + conv_shape = (storage_layer_count, conv_channels, config.conv_width) conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize regions = [] logger.debug( @@ -360,7 +363,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads head_start = config.tp_rank // replica_count * config.num_v_heads write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (layer_count, config.global_v_heads, config.key_dim, config.value_dim) + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) element_bytes = config.ssm_dtype.itemsize ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes ssm_region = TransferRegion( diff --git a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py index b06bfb1cf6..e70d52e376 100644 --- a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py +++ b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py @@ -25,15 +25,15 @@ def __init__( self.size = size device = torch.device(device) pin_memory = pin_memory and device.type == "cpu" - layers, heads, dim = len(config.layer_ids), config.num_k_heads, config.key_dim + storage_layer_count, heads, dim = len(config.logical_layer_ids), config.num_k_heads, config.key_dim self.conv_state = create_page_pool_tensor( - size=(size, layers, 3 * heads * dim, config.conv_width), + size=(size, storage_layer_count, 3 * heads * dim, config.conv_width), dtype=config.conv_dtype, device=device, pin_memory=pin_memory, ) self.ssm_state = create_page_pool_tensor( - size=(size, layers, heads, dim, dim), + size=(size, storage_layer_count, heads, dim, dim), dtype=config.ssm_dtype, device=device, pin_memory=pin_memory, diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py index 6e8a5dbbe7..a800eed33e 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -40,15 +40,18 @@ def __init__( self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step - self._layer_mapping = {layer: i for i, layer in enumerate(config.layer_ids)} - num_layers = len(config.layer_ids) + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) self.conv_state = torch.zeros( - (num_layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), dtype=config.conv_dtype, device=device, ) self.ssm_state = torch.zeros( - (num_layers, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + (storage_layer_count, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), dtype=config.ssm_dtype, device=device, ) @@ -86,10 +89,10 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== - def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """返回该层 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - layer = self._layer_mapping[layer_id] - return self.conv_state[layer], self.ssm_state[layer] + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] def reset_bounded_att_state(self, req_indices: List[int]) -> None: """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" @@ -305,9 +308,9 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 config = self.config - layer_count = len(config.layer_ids) + storage_layer_count = len(config.logical_layer_ids) conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (layer_count, conv_channels, config.conv_width) + conv_shape = (storage_layer_count, conv_channels, config.conv_width) conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize regions = [] logger.debug( @@ -360,7 +363,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads head_start = config.tp_rank // replica_count * config.num_v_heads write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (layer_count, config.global_v_heads, config.key_dim, config.value_dim) + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) element_bytes = config.ssm_dtype.itemsize ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes ssm_region = TransferRegion( diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index b584e03c14..39d49766e7 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -15,8 +15,13 @@ @dataclass(frozen=True) class FullAttentionConfig: - # 当前 State 管理的逻辑层号；元组顺序对应 storage 的 layer 轴顺序，不要求层号连续。 - layer_ids: tuple[int, ...] + # 物理存储层与逻辑推理层的对应关系，外层顺序就是所有 Tensor 的 layer 轴顺序。 + # 每个非空内层元组列出共享同一份存储的逻辑层号；层号可以不连续，但不能重复出现。 + # 例如 ((0, 4), (8,)) 只分配两层：逻辑层 0/4 使用 storage[0]，逻辑层 8 使用 storage[1]。 + # 不共享时写成 ((0,), (4,), (8,))，此时分配三层；容量和传输按 len(logical_layer_groups) 计算。 + # 组内共享该方案的全部字段（KV、scale、indexer 等），因此布局及配套参数必须兼容。 + # 分组不决定执行顺序或写入者：模型负责让共享消费者跳过 KV 写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] # TP 切分前的 KV head 数：GQA 的 K/V 各有这么多个，不把二者相加，也不计 TP 副本。 # 不是 query head 数；例如全局 8 个 KV head，实际有 8 个 K head 和 8 个 V head。 # MLA 的压缩 latent 作为一个 KV 字段存储，通常设为 1，并在各 TP rank 复制。 @@ -43,6 +48,20 @@ def head_num(self) -> int: return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) def __post_init__(self) -> None: + # 只保留这一份分组配置；层号映射在 State 构造时展开，不另存逻辑/物理层数。 + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(logical_layer_group, tuple) and logical_layer_group + for logical_layer_group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + logical_layer_ids = [ + logical_layer_id + for logical_layer_group in self.logical_layer_groups + for logical_layer_id in logical_layer_group + ] + assert all(isinstance(logical_layer_id, int) and logical_layer_id >= 0 for logical_layer_id in logical_layer_ids), ( + "logical layer IDs must be nonnegative integers" + ) + assert len(logical_layer_ids) == len(set(logical_layer_ids)), "a logical layer must belong to exactly one storage group" assert self.global_head_num > 0 and 0 <= self.tp_rank < self.tp_world_size assert ( self.replicated @@ -82,12 +101,12 @@ def __init__( @abstractmethod def initialize_storage(self, token_capacity: int) -> None: - """按 token 容量创建独立存储，沿用已设置的页大小、设备及配套参数。""" + """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" raise NotImplementedError() @abstractmethod def get_token_bytes(self) -> int: - """返回当前 TP rank 的单 token 存储字节数，包含动态量化参数和行 padding。""" + """累计当前 TP rank 各物理层的单 token 字节数，包含动态量化参数和行 padding；共享逻辑层不重复计数。""" raise NotImplementedError() # ======================================================================== @@ -95,12 +114,12 @@ def get_token_bytes(self) -> int: # ======================================================================== @abstractmethod - def get_layer_view(self, layer_id: int) -> Any: - """按逻辑层号返回计算所需的存储视图，具体字段由该布局定义。""" + def get_layer_view(self, logical_layer_id: int) -> Any: + """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" raise NotImplementedError() @abstractmethod - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。""" raise NotImplementedError() diff --git a/lightllm/common/att_resources/full_att/create_gqa_state.py b/lightllm/common/att_resources/full_att/create_gqa_state.py index 527e1bcf7e..c31a95c63b 100644 --- a/lightllm/common/att_resources/full_att/create_gqa_state.py +++ b/lightllm/common/att_resources/full_att/create_gqa_state.py @@ -33,7 +33,7 @@ def create_gqa_state( return GQAFullAttentionState(layout_config=config, page_size=page_size, device=device) if kv_type in ("int4kv", "int8kv"): quant_config = GroupQuantizedAttentionConfig( - layer_ids=config.layer_ids, + logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, head_dim=config.head_dim, dtype=config.dtype, @@ -48,7 +48,7 @@ def create_gqa_state( assert scales is not None, "static FP8 requires calibration scales" assert kv_type != "fp8kv_sph" or q_scales is not None, "per-head FP8 requires Q calibration scales" quant_config = StaticFP8AttentionConfig( - layer_ids=config.layer_ids, + logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, head_dim=config.head_dim, dtype=torch.uint8, diff --git a/lightllm/common/att_resources/full_att/create_mla_state.py b/lightllm/common/att_resources/full_att/create_mla_state.py index e87332cbe1..e412b12e20 100644 --- a/lightllm/common/att_resources/full_att/create_mla_state.py +++ b/lightllm/common/att_resources/full_att/create_mla_state.py @@ -41,7 +41,7 @@ def create_mla_state( kv_bytes = 656 if fp8 else config.head_dim * config.dtype.itemsize row_bytes = (kv_bytes + indexer_dim + 4 + 15) // 16 * 16 packed_config = PackedMLAAttentionConfig( - layer_ids=config.layer_ids, + logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, head_dim=row_bytes // storage_dtype.itemsize, dtype=storage_dtype, diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index 4a077b6cc1..a87a844982 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -49,7 +49,12 @@ def __init__( self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化方案显式设置，区别于 scale 精度。 self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 @@ -71,7 +76,7 @@ def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" config = self.config return ( - len(config.layer_ids) * 2 * config.head_num * config.head_dim + len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize ) @@ -79,23 +84,23 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" - layer = self._layer_mapping[layer_id] - return self.k_storage[layer], self.v_storage[layer] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - layer = self._layer_mapping[layer_id] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv(K=kv[:, :self.config.head_num], DestLoc=token_indices, Out=self.k_storage[layer]) - destindex_copy_kv(K=kv[:, self.config.head_num:], DestLoc=token_indices, Out=self.v_storage[layer]) + destindex_copy_kv(K=kv[:, :self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index]) + destindex_copy_kv(K=kv[:, self.config.head_num:], DestLoc=token_indices, Out=self.v_storage[storage_layer_index]) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" @@ -276,7 +281,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) global_shape = list(local_shape) global_shape[axes.index("token")] = page_tokens head_axis = axes.index("head") @@ -332,7 +337,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.layer_ids), + len(self.config.logical_layer_groups), self.capacity + self.page_size, self.config.head_num, self.config.head_dim, diff --git a/lightllm/common/att_resources/full_att/gqa_transfer.py b/lightllm/common/att_resources/full_att/gqa_transfer.py index 55cd40f05e..3f4537d151 100644 --- a/lightllm/common/att_resources/full_att/gqa_transfer.py +++ b/lightllm/common/att_resources/full_att/gqa_transfer.py @@ -207,7 +207,7 @@ def _copy_gqa_pages( if copy_block: dim = offsets % WIDTH head = offsets // WIDTH % HEADS - layer = offsets // (WIDTH * HEADS * token_count) + storage_layer_index = offsets // (WIDTH * HEADS * token_count) slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) @@ -215,14 +215,14 @@ def _copy_gqa_pages( # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 K/V 字段的大小。 local_offset = ( - layer * LOCAL_STRIDES[0] + storage_layer_index * LOCAL_STRIDES[0] + slot * LOCAL_STRIDES[1] + head * LOCAL_STRIDES[2] + dim * LOCAL_STRIDES[3] ) external_offset = ( page * EXTERNAL_STRIDES[0] - + layer * EXTERNAL_STRIDES[1] + + storage_layer_index * EXTERNAL_STRIDES[1] + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + head * EXTERNAL_STRIDES[3] + dim * EXTERNAL_STRIDES[4] diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index ed857731df..180618e96e 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -58,7 +58,12 @@ def __init__( """按分组量化配置准备 KV 与动态 scale 存储。""" self.config = layout_config self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 @@ -79,7 +84,7 @@ def initialize_storage(self, token_capacity: int) -> None: self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) shape = ( - len(self.config.layer_ids), token_capacity + self.page_size, + len(self.config.logical_layer_groups), token_capacity + self.page_size, self.config.head_num, self.config.head_dim // self.config.group_size, ) self.k_scales = torch.empty(size=shape, dtype=self.config.dtype, device=self.device) @@ -90,21 +95,21 @@ def get_token_bytes(self) -> int: config = self.config head_bytes = config.head_dim * config.bits // 8 head_bytes += config.head_dim // config.group_size * config.dtype.itemsize - return len(config.layer_ids) * 2 * config.head_num * head_bytes + return len(config.logical_layer_groups) * 2 * config.head_num * head_bytes # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + def get_layer_view(self, logical_layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """返回该层 K/V 编码数据及各自的分组 scale 视图。""" - layer = self._layer_mapping[layer_id] - return (self.k_storage[layer], self.k_scales[layer]), (self.v_storage[layer], self.v_scales[layer]) + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return (self.k_storage[storage_layer_index], self.k_scales[storage_layer_index]), (self.v_storage[storage_layer_index], self.v_scales[storage_layer_index]) - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """量化输入浮点 K/V，并将编码数据和分组 scale 写入相同 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - layer = self._layer_mapping[layer_id] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" @@ -121,8 +126,8 @@ def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torc destindex_copy_int4kv( KV=values, DestLoc=token_indices, - KV_buffer=storage[layer], - KV_scale_buffer=scales_buffer[layer], + KV_buffer=storage[storage_layer_index], + KV_scale_buffer=scales_buffer[storage_layer_index], quant_group_size=self.config.group_size, ) else: @@ -131,8 +136,8 @@ def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torc destindex_copy_quantize_kv( K=values, DestLoc=token_indices, - Out=storage[layer], - Out_scale=scales_buffer[layer], + Out=storage[storage_layer_index], + Out_scale=scales_buffer[storage_layer_index], quant_group_dim=self.config.group_size, ) @@ -347,7 +352,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) global_shape = list(local_shape) global_shape[axes.index("token")] = page_tokens head_axis = axes.index("head") @@ -403,6 +408,6 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.layer_ids), self.capacity + self.page_size, + len(self.config.logical_layer_groups), self.capacity + self.page_size, self.config.head_num, self.config.head_dim * self.config.bits // 8, ) diff --git a/lightllm/common/att_resources/full_att/group_quantized_transfer.py b/lightllm/common/att_resources/full_att/group_quantized_transfer.py index 9b47b269fd..b4376c7235 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_transfer.py +++ b/lightllm/common/att_resources/full_att/group_quantized_transfer.py @@ -217,7 +217,7 @@ def _copy_group_quantized_pages( if copy_block: dim = offsets % WIDTH head = offsets // WIDTH % HEADS - layer = offsets // (WIDTH * HEADS * token_count) + storage_layer_index = offsets // (WIDTH * HEADS * token_count) slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) @@ -225,14 +225,14 @@ def _copy_group_quantized_pages( # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 local_offset = ( - layer * LOCAL_STRIDES[0] + storage_layer_index * LOCAL_STRIDES[0] + slot * LOCAL_STRIDES[1] + head * LOCAL_STRIDES[2] + dim * LOCAL_STRIDES[3] ) external_offset = ( page * EXTERNAL_STRIDES[0] - + layer * EXTERNAL_STRIDES[1] + + storage_layer_index * EXTERNAL_STRIDES[1] + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + head * EXTERNAL_STRIDES[3] + dim * EXTERNAL_STRIDES[4] diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index 62e27b058f..9a8b776008 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -39,7 +39,12 @@ def __init__( self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 @@ -59,7 +64,7 @@ def get_token_bytes(self) -> int: """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" config = self.config return ( - len(config.layer_ids) * config.head_num * config.head_dim + len(config.logical_layer_groups) * config.head_num * config.head_dim * self.storage_dtype.itemsize ) @@ -67,20 +72,20 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> torch.Tensor: + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" - return self.storage[self._layer_mapping[layer_id]] + return self.storage[self._logical_to_storage_layer[logical_layer_id]] - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """通过 CUDA 算子将形状匹配的 MLA latent 数据写入指定物理 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - layer = self._layer_mapping[layer_id] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[layer]) + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index]) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" @@ -244,7 +249,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: field_specs = (("kv", self.storage_dtype, config.head_dim, "dim"),) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) global_shape = list(local_shape) global_shape[axes.index("token")] = page_tokens head_axis = axes.index("head") @@ -300,7 +305,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.layer_ids), + len(self.config.logical_layer_groups), self.capacity + self.page_size, self.config.head_num, self.config.head_dim, diff --git a/lightllm/common/att_resources/full_att/mla_transfer.py b/lightllm/common/att_resources/full_att/mla_transfer.py index cc763e570a..42145c3825 100644 --- a/lightllm/common/att_resources/full_att/mla_transfer.py +++ b/lightllm/common/att_resources/full_att/mla_transfer.py @@ -202,7 +202,7 @@ def _copy_mla_pages( if copy_block: dim = offsets % WIDTH head = offsets // WIDTH % HEADS - layer = offsets // (WIDTH * HEADS * token_count) + storage_layer_index = offsets // (WIDTH * HEADS * token_count) slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) @@ -210,14 +210,14 @@ def _copy_mla_pages( # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 local_offset = ( - layer * LOCAL_STRIDES[0] + storage_layer_index * LOCAL_STRIDES[0] + slot * LOCAL_STRIDES[1] + head * LOCAL_STRIDES[2] + dim * LOCAL_STRIDES[3] ) external_offset = ( page * EXTERNAL_STRIDES[0] - + layer * EXTERNAL_STRIDES[1] + + storage_layer_index * EXTERNAL_STRIDES[1] + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + head * EXTERNAL_STRIDES[3] + dim * EXTERNAL_STRIDES[4] diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index adf5767e41..8f8142285f 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -80,7 +80,12 @@ def __init__( self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 @@ -101,7 +106,7 @@ def get_token_bytes(self) -> int: """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" config = self.config return ( - len(config.layer_ids) * config.head_num * config.head_dim + len(config.logical_layer_groups) * config.head_num * config.head_dim * self.storage_dtype.itemsize ) @@ -121,18 +126,18 @@ def indexer_offset(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> PackedMLALayerView: + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" - row = self.storage[self._layer_mapping[layer_id]] + row = self.storage[self._logical_to_storage_layer[logical_layer_id]] raw = row.view(torch.uint8) kv = raw[..., : self.kv_nbytes] if self.config.fp8_nope_dim else row[..., : self.config.kv_dim] indexer = raw[..., self.indexer_offset : self.indexer_offset + self.config.indexer_dim + 4] return PackedMLALayerView(kv=kv, indexer_k=indexer) - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """写入 MLA KV 及其量化参数，保留同一 token 行中的 indexer 区域。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - layer = self._layer_mapping[layer_id] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" @@ -141,7 +146,7 @@ def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torc destindex_copy_kv_flashmla_fp8, ) - raw = self.storage[layer] + raw = self.storage[storage_layer_index] destindex_copy_kv_flashmla_fp8( KV_nope=kv[..., :512], KV_rope=kv[..., 512:], @@ -153,7 +158,7 @@ def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torc else: from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[layer, ..., : self.config.kv_dim]) + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.kv_dim]) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" @@ -350,7 +355,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: )) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) global_shape = list(local_shape) global_shape[axes.index("token")] = page_tokens head_axis = axes.index("head") @@ -406,7 +411,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.layer_ids), + len(self.config.logical_layer_groups), self.capacity + self.page_size, self.config.head_num, self.config.head_dim, diff --git a/lightllm/common/att_resources/full_att/packed_mla_transfer.py b/lightllm/common/att_resources/full_att/packed_mla_transfer.py index b7ef3640ee..b0310eaa37 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_transfer.py +++ b/lightllm/common/att_resources/full_att/packed_mla_transfer.py @@ -226,7 +226,7 @@ def _copy_packed_mla_pages( if copy_block: dim = offsets % WIDTH head = offsets // WIDTH % HEADS - layer = offsets // (WIDTH * HEADS * token_count) + storage_layer_index = offsets // (WIDTH * HEADS * token_count) slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) @@ -234,14 +234,14 @@ def _copy_packed_mla_pages( # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 local_offset = ( - layer * LOCAL_STRIDES[0] + storage_layer_index * LOCAL_STRIDES[0] + slot * LOCAL_STRIDES[1] + head * LOCAL_STRIDES[2] + dim * LOCAL_STRIDES[3] ) external_offset = ( page * EXTERNAL_STRIDES[0] - + layer * EXTERNAL_STRIDES[1] + + storage_layer_index * EXTERNAL_STRIDES[1] + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + head * EXTERNAL_STRIDES[3] + dim * EXTERNAL_STRIDES[4] diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index dc36d675ab..6097f433fb 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -56,14 +56,14 @@ def __init__( """按静态 FP8 配置准备 KV 存储，并保留传入的标定参数。""" # 标定 Tensor 由 State 持有，在接收构造参数时校验，不重复放入 config。 heads = 1 if layout_config.scale_mode == "per_tensor" else layout_config.head_num - assert scales is not None and scales.shape == (len(layout_config.layer_ids), 2 * heads), ( + assert scales is not None and scales.shape == (len(layout_config.logical_layer_groups), 2 * heads), ( "calibration scales do not match KV layers and heads" ) assert scales.dtype == torch.float32 and bool((torch.isfinite(input=scales) & (scales > 0)).all()), ( "calibration scales must be finite positive float32 values" ) if q_scales is not None: - assert q_scales.shape == (len(layout_config.layer_ids), layout_config.head_num) and ( + assert q_scales.shape == (len(layout_config.logical_layer_groups), layout_config.head_num) and ( q_scales.dtype == torch.float32 ), "Q calibration does not match layers and heads" self.scales = scales @@ -71,7 +71,12 @@ def __init__( self.cpu_scales = None # 实际存储初始化时准备计算 backend 使用的 CPU 标定值。 self.config = layout_config self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 - self._layer_mapping = {layer: i for i, layer in enumerate(layout_config.layer_ids)} + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 @@ -101,7 +106,7 @@ def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" config = self.config return ( - len(config.layer_ids) * 2 * config.head_num * config.head_dim + len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize ) @@ -109,15 +114,15 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def get_layer_view(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" - layer = self._layer_mapping[layer_id] - return self.k_storage[layer], self.v_storage[layer] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] - def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """按固定 per-tensor 或 per-head scale 量化输入 K/V，写入 FP8 编码。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - layer = self._layer_mapping[layer_id] + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" @@ -126,10 +131,10 @@ def write_full_att_kv(self, layer_id: int, token_indices: torch.Tensor, kv: torc destindex_copy_kv_fp8( KV=kv, DestLoc=token_indices, - scales=self.scales[layer], - Out=self.k_storage[layer].view(torch.float8_e4m3fn), + scales=self.scales[storage_layer_index], + Out=self.k_storage[storage_layer_index].view(torch.float8_e4m3fn), is_per_tensor_quant=self.config.scale_mode == "per_tensor", - Out_v=self.v_storage[layer].view(torch.float8_e4m3fn), + Out_v=self.v_storage[storage_layer_index].view(torch.float8_e4m3fn), ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -318,7 +323,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.layer_ids), page_tokens, config.head_num, field_width) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) global_shape = list(local_shape) global_shape[axes.index("token")] = page_tokens head_axis = axes.index("head") @@ -374,7 +379,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.layer_ids), + len(self.config.logical_layer_groups), self.capacity + self.page_size, self.config.head_num, self.config.head_dim, diff --git a/lightllm/common/att_resources/full_att/static_fp8_transfer.py b/lightllm/common/att_resources/full_att/static_fp8_transfer.py index 107ec9bda7..1674f7d5e0 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_transfer.py +++ b/lightllm/common/att_resources/full_att/static_fp8_transfer.py @@ -209,7 +209,7 @@ def _copy_static_fp8_pages( if copy_block: dim = offsets % WIDTH head = offsets // WIDTH % HEADS - layer = offsets // (WIDTH * HEADS * token_count) + storage_layer_index = offsets // (WIDTH * HEADS * token_count) slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) @@ -217,14 +217,14 @@ def _copy_static_fp8_pages( # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 local_offset = ( - layer * LOCAL_STRIDES[0] + storage_layer_index * LOCAL_STRIDES[0] + slot * LOCAL_STRIDES[1] + head * LOCAL_STRIDES[2] + dim * LOCAL_STRIDES[3] ) external_offset = ( page * EXTERNAL_STRIDES[0] - + layer * EXTERNAL_STRIDES[1] + + storage_layer_index * EXTERNAL_STRIDES[1] + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + head * EXTERNAL_STRIDES[3] + dim * EXTERNAL_STRIDES[4] diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 938a294f6b..095ef834e2 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -76,13 +76,14 @@ def uses_causal_attention(self) -> bool: is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") return not is_parallel_block_draft - def _find_layer_index( + def _find_storage_layer_index( self, k: torch.Tensor, v: torch.Tensor, state: "FullAttentionState" ) -> int: - for layer_index, layer_id in enumerate(state.config.layer_ids): - layer_k, layer_v = state.get_layer_view(layer_id) + # 标定参数按物理层排列；每组只检查一个代表逻辑层，避免共享层重复计数。 + for storage_layer_index, logical_layer_group in enumerate(state.config.logical_layer_groups): + layer_k, layer_v = state.get_layer_view(logical_layer_id=logical_layer_group[0]) if k.data_ptr() == layer_k.data_ptr() and v.data_ptr() == layer_v.data_ptr(): - return layer_index + return storage_layer_index raise AssertionError("K/V views do not belong to the same State layer") diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index c82de810b8..0629e2215e 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -89,8 +89,8 @@ def _fp8_prefill_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", ) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, state=state) - static_q_scales = state.q_scales[layer_index] + storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) + static_q_scales = state.q_scales[storage_layer_index] q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales) q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) o = flash_attn_with_kvcache( @@ -106,8 +106,8 @@ def _fp8_prefill_att( window_size=(-1, -1), softcap=0.0, q_descale=q_scale, - k_descale=self.k_descale[layer_index], - v_descale=self.v_descale[layer_index], + k_descale=self.k_descale[storage_layer_index], + v_descale=self.v_descale[storage_layer_index], return_softmax_lse=False, ) return o @@ -181,13 +181,13 @@ def _fp8_decode_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", ) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, state=state) + storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) q_head_num = q.shape[1] att_batch_size = self.b_att_seq_len.shape[0] static_q_scales = state.q_scales - q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales[layer_index]) - q_scale = static_q_scales[layer_index].view(1, k_head_num).expand(att_batch_size, k_head_num) + q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales[storage_layer_index]) + q_scale = static_q_scales[storage_layer_index].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, @@ -201,8 +201,8 @@ def _fp8_decode_att( window_size=(-1, -1), softcap=0.0, q_descale=q_scale.view(att_batch_size, k_head_num), - k_descale=self.k_descale[layer_index], - v_descale=self.v_descale[layer_index], + k_descale=self.k_descale[storage_layer_index], + v_descale=self.v_descale[storage_layer_index], return_softmax_lse=False, ) return o diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 28a22733cd..022eb16cb6 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -70,9 +70,9 @@ def _fp8_prefill_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", ) - layer_index = self.backend._find_layer_index(k=k, v=v, state=state) - k_descale = state.cpu_scales[layer_index][0] - v_descale = state.cpu_scales[layer_index][1] + storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) + k_descale = state.cpu_scales[storage_layer_index][0] + v_descale = state.cpu_scales[storage_layer_index][1] self.prefill_wrapper.run( q, (k, v), @@ -126,10 +126,10 @@ def _fp8_decode_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", ) - layer_index = self.backend._find_layer_index(k=k, v=v, state=state) + storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) - k_descale = state.cpu_scales[layer_index][0] - v_descale = state.cpu_scales[layer_index][1] + k_descale = state.cpu_scales[storage_layer_index][0] + v_descale = state.cpu_scales[storage_layer_index][1] self.decode_wrapper.run( q, (k, v), diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 1d487e05b9..fde9c74c7e 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -141,7 +141,7 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: LinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -283,7 +283,7 @@ def decode_att( backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 5d359d9f8f..3807107666 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -85,7 +85,7 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -206,7 +206,7 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(layer_num) + conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) conv_update = causal_conv1d_update if backend.mtp_step > 0: diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index c7af388b98..b063aaef00 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -29,7 +29,7 @@ def get_full_attention_options( args: StartArgs | EasyDict, *, model_class: type[TpPartBaseModel] | None = None, - model_layer_ids: Sequence[int] | None = None, + model_logical_layer_ids: Sequence[int] | None = None, ): """Return default KV geometry or the model's specialized storage options.""" build_options = getattr(model_class, "_full_attention_options", None) @@ -38,10 +38,22 @@ def get_full_attention_options( else: heads = config.get("num_key_value_heads", config["num_attention_heads"]) dim = config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]) - layers = config["num_hidden_layers"] - options = dict(heads=heads, dim=dim, layer_ids=tuple(range(layers))) - if model_layer_ids is not None: - options["layer_ids"] = tuple(model_layer_ids[i] for i in options["layer_ids"] if i < len(model_layer_ids)) + logical_layer_count = config["num_hidden_layers"] + options = dict( + heads=heads, dim=dim, + logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in range(logical_layer_count)), + ) + if model_logical_layer_ids is not None: + # 模型选项使用本模型内的层号；转换为实际逻辑层号时保留共享分组及物理层顺序。 + options["logical_layer_groups"] = tuple( + tuple( + model_logical_layer_ids[local_logical_layer_id] + for local_logical_layer_id in logical_layer_group + if local_logical_layer_id < len(model_logical_layer_ids) + ) + for logical_layer_group in options["logical_layer_groups"] + if any(local_logical_layer_id < len(model_logical_layer_ids) for local_logical_layer_id in logical_layer_group) + ) return options @@ -52,22 +64,22 @@ def create_full_attention_definition( *, tp_rank: int = 0, tp_world_size: int = 1, - model_layer_ids: Sequence[int] | None = None, - calibration_layer_start: int | None = None, + model_logical_layer_ids: Sequence[int] | None = None, + calibration_storage_layer_start: int | None = None, ): """Build a meta State using the model's configuration and storage hooks.""" from lightllm.utils.att_resource_config import build_full_attention_state config = get_resource_config(model_config, model_class=model_class) - options = get_full_attention_options(config, args, model_class=model_class, model_layer_ids=model_layer_ids) - if not options["layer_ids"]: + options = get_full_attention_options(config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids) + if not options["logical_layer_groups"]: return None state = build_full_attention_state( args, **options, tp_rank=tp_rank, tp_world_size=tp_world_size, - calibration_layer_start=calibration_layer_start, + calibration_storage_layer_start=calibration_storage_layer_start, ) return state @@ -77,7 +89,7 @@ def create_bounded_attention_definition( model_config: dict[str, object], args: StartArgs | EasyDict, *, - model_layer_ids: Sequence[int], + model_logical_layer_ids: Sequence[int], tp_rank: int = 0, tp_world_size: int = 1, ): @@ -86,15 +98,15 @@ def create_bounded_attention_definition( if create_definition is None: return None return create_definition( - model_config, args, model_layer_ids=model_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size + model_config, args, model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size ) -def get_model_calibration_layer_start(model: TpPartBaseModel) -> int: - """Count preceding models' full-attention layers without initializing their States.""" +def get_model_calibration_storage_layer_start(model: TpPartBaseModel) -> int: + """按前序模型的物理层数累计标定表行偏移，不创建 State，也不使用逻辑层号作为行号。""" if not model.is_mtp_draft_model: return 0 - layer_start = 0 + storage_layer_start = 0 for previous in (model.main_model, *model.mtp_previous_draft_models): if previous.is_mtp_draft_model and previous.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): continue @@ -103,10 +115,10 @@ def get_model_calibration_layer_start(model: TpPartBaseModel) -> int: config, previous.args, model_class=type(previous), - model_layer_ids=tuple(layer.layer_num_ for layer in previous.layers_infer), + model_logical_layer_ids=tuple(layer.layer_num_ for layer in previous.layers_infer), ) - layer_start += len(options["layer_ids"]) - return layer_start + storage_layer_start += len(options["logical_layer_groups"]) + return storage_layer_start def init_model_full_attention(model: TpPartBaseModel) -> None: @@ -119,17 +131,18 @@ def init_model_full_attention(model: TpPartBaseModel) -> None: model.args, tp_rank=model.tp_rank_, tp_world_size=model.tp_world_size_, - model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - calibration_layer_start=get_model_calibration_layer_start(model), + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + calibration_storage_layer_start=get_model_calibration_storage_layer_start(model), ) if definition is None: return model.att_resource_manager.set_att_state(model.att_resource_role, definition) # 层和 indexer 只记录规范成员名，推理时直接访问管理器中已分配的 State。 name = definition.resource_id + # 绑定所有逻辑层，包括共享消费者；State 自己将逻辑层号映射到物理层。 + logical_layer_ids = {logical_layer_id for logical_layer_group in definition.config.logical_layer_groups for logical_layer_id in logical_layer_group} for layer in model.layers_infer: - layer_id = layer.kv_share_target_layer_ if getattr(layer, "is_kv_shared_", False) else layer.layer_num_ - if layer_id in definition.config.layer_ids: + if layer.layer_num_ in logical_layer_ids: layer.full_att_state_name = name if hasattr(layer, "indexer"): layer.indexer.full_att_state_name = name @@ -139,7 +152,7 @@ def init_model_bounded_attention(model: TpPartBaseModel) -> None: """按当前模型的配置创建请求运行态，并直接挂到共享管理器；大小页稍后分配。""" definition = create_bounded_attention_definition( type(model), model.config, model.args, - model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + 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_, ) if definition is None: diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index cf3df4b0aa..9a2183326b 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -35,7 +35,7 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[to def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight): state = getattr(infer_state.att_resource_manager, self.full_att_state_name) state.write_full_att_kv( - layer_id=self.layer_num_, + logical_layer_id=self.layer_num_, token_indices=infer_state.mem_index, kv=cache_kv, ) diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index 1b1bada67c..a018399bdb 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -29,7 +29,7 @@ def _context_attention_kernel( out=None, ) -> torch.Tensor: state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -45,7 +45,7 @@ def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 58fd1ba930..a5a061cbd2 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -104,7 +104,7 @@ def _token_attention_kernel( q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - kv = state.get_layer_view(layer_id=self.layer_num_) + kv = state.get_layer_view(logical_layer_id=self.layer_num_) out = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), @@ -121,7 +121,7 @@ def _decompress_kv( layer_weight: Deepseek2TransformerLayerWeight, ): state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - compressed_kv = state.get_layer_view(self.layer_num_) + compressed_kv = state.get_layer_view(logical_layer_id=self.layer_num_) total_token_num = infer_state.total_token_num sampled_compressed_kv_nope = self.alloc_tensor( diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index a812f8750f..43f9998a52 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -102,7 +102,7 @@ def _context_attention_kernel( state = getattr(infer_state.att_resource_manager, self.full_att_state_name) mla_out = infer_state.prefill_att_state.prefill_att( q=q_all, - k=state.get_layer_view(layer_id=self.layer_num_).kv, + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -145,7 +145,7 @@ def _token_attention_kernel( state = getattr(infer_state.att_resource_manager, self.full_att_state_name) o_tensor = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=state.get_layer_view(layer_id=self.layer_num_).kv, + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -203,7 +203,7 @@ def _get_indices( K_fp8=k_fp8, K_scale=k_scale, DestLoc=infer_state.mem_index, - O_buffer=state.get_layer_view(self.layer_idx_).indexer_k, + O_buffer=state.get_layer_view(logical_layer_id=self.layer_idx_).indexer_k, ) weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale @@ -219,7 +219,7 @@ def _get_indices( mtp_step = get_env_start_args().mtp_step # Use efficient Triton kernel to extract FP8 keys and scales from buffer k_fp8_, k_scale_ = extract_indexer_ks( - I_buffer=state.get_layer_view(self.layer_idx_).indexer_k, + I_buffer=state.get_layer_view(logical_layer_id=self.layer_idx_).indexer_k, b_seq_len=infer_state.b_seq_len, b_req_idx=infer_state.b_req_idx, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 742b2fc996..2010a674ee 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -84,17 +84,6 @@ def __init__(self, layer_num, network_config): kv_shared_count = network_config.get("num_kv_shared_layers") or 0 total_layers = network_config["num_hidden_layers"] self.is_kv_shared_ = kv_shared_count > 0 and layer_num >= total_layers - kv_shared_count - self.kv_share_target_layer_ = None - if self.is_kv_shared_: - cutoff = total_layers - kv_shared_count - for j in range(cutoff - 1, -1, -1): - if network_config["layer_types"][j] == layer_type: - self.kv_share_target_layer_ = j - break - assert self.kv_share_target_layer_ is not None, ( - f"layer {layer_num} ({layer_type}) is KV-shared but no earlier non-shared " - f"layer of the same type found below cutoff={cutoff}" - ) # Always 1.0: NoPE dims for full-attn layers are zero-padded into # cos/sin (cos=1, sin=0 → identity), so the kernel walks the whole @@ -193,10 +182,9 @@ def _att_control(self): return AttControl(use_sliding_window=False, sliding_window=(-1, -1)) def _get_layer_kv(self, infer_state: InferStateInfo): - # KV-shared layers read from the target layer's cache slot. - layer_idx = self.kv_share_target_layer_ if self.is_kv_shared_ else self.layer_num_ + # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k_raw, _v_raw = state.get_layer_view(layer_id=layer_idx) + _k_raw, _v_raw = state.get_layer_view(logical_layer_id=self.layer_num_) # _k_raw / _v_raw shape (S, cache_slot_num, cache_slot_dim). Use .view # (not .reshape) so any non-contiguous layout from a future mem_manager # backend fails loudly instead of silently copying — slice + view is diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 61e933ef5f..003840ab77 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -101,6 +101,21 @@ def _full_attention_options(cls, config, args): raise ValueError("Gemma4 KV width must contain complete storage heads") options = get_full_attention_options(config, args) options.update(heads=per_token_k_width // head_dim, dim=head_dim) + shared_logical_layer_count = config.get("num_kv_shared_layers") or 0 + if shared_logical_layer_count: + # 后部共享层读取同类型的最后一个非共享层；仅生产层占一份物理 KV。 + # 由模型配置生成分组，GPU State 和 CPU cache 布局使用完全相同的层顺序。 + storage_layer_count = config["num_hidden_layers"] - shared_logical_layer_count + assert 0 < storage_layer_count < config["num_hidden_layers"], "invalid KV-shared layer count" + logical_layer_groups = [[logical_layer_id] for logical_layer_id in range(storage_layer_count)] + for logical_layer_id in range(storage_layer_count, config["num_hidden_layers"]): + for source_logical_layer_id in range(storage_layer_count - 1, -1, -1): + if config["layer_types"][source_logical_layer_id] == config["layer_types"][logical_layer_id]: + logical_layer_groups[source_logical_layer_id].append(logical_layer_id) + break + else: + raise AssertionError("KV-shared layer requires an earlier producer of the same attention type") + options["logical_layer_groups"] = tuple(tuple(logical_layer_group) for logical_layer_group in logical_layer_groups) return options def _init_att_backend(self): diff --git a/lightllm/models/glm5_next/att_resource_manager.py b/lightllm/models/glm5_next/att_resource_manager.py index e9b57b6dd0..5c43191b44 100644 --- a/lightllm/models/glm5_next/att_resource_manager.py +++ b/lightllm/models/glm5_next/att_resource_manager.py @@ -24,14 +24,17 @@ def init_model_att_tail(model: Glm5NextTpPartModel) -> None: config = get_resource_config(model.config, model_class=type(model)) options = get_full_attention_options( config, model.args, model_class=type(model), - model_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), ) - if not options["layer_ids"]: + if not options["logical_layer_groups"]: return cfg = model.linear_config state = create_kpool_state( config=KPoolStateConfig( - layer_ids=options["layer_ids"], pool_size=cfg.index_kpool, + logical_layer_ids=tuple( + logical_layer_id for logical_layer_group in options["logical_layer_groups"] + for logical_layer_id in logical_layer_group + ), pool_size=cfg.index_kpool, head_dim=cfg.index_head_dim, dtype=cfg.full_att_dtype, tp_rank=model.tp_rank_, tp_world_size=model.tp_world_size_, ), @@ -54,19 +57,19 @@ def __init__( tp_rank: int = 0, device: str | torch.device | int = "cuda", role: str = "main", - linear_layer_ids: Sequence[int] | None = None, - full_layer_ids: Sequence[int] | None = None, + bounded_logical_layer_ids: Sequence[int] | None = None, + full_att_logical_layer_ids: Sequence[int] | None = None, ): super().__init__(request_capacity=request_capacity) self.cache_config = cache_config self.cpu_tail_restore = True - layers = ( + logical_layer_ids = ( tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) - if linear_layer_ids is None - else tuple(linear_layer_ids) + if bounded_logical_layer_ids is None + else tuple(bounded_logical_layer_ids) ) config = ConvSSMStateConfig( - layer_ids=layers, + logical_layer_ids=logical_layer_ids, global_k_heads=cache_config.global_linear_k_heads, global_v_heads=cache_config.global_linear_v_heads, key_dim=cache_config.head_linear_k_dim, @@ -77,11 +80,11 @@ def __init__( tp_rank=tp_rank, tp_world_size=cache_config.tp_world_size, ) - if layers: + if logical_layer_ids: state = create_kda_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) self.set_att_state(role, state) self.att_tail_type = f"indexer_kpool_c{cache_config.index_kpool}_tail" - own_full_layers = ( + own_full_att_logical_layer_ids = ( tuple( range( cache_config.full_attention_interval - 1, @@ -89,12 +92,12 @@ def __init__( cache_config.full_attention_interval, ) ) - if full_layer_ids is None - else tuple(full_layer_ids) + if full_att_logical_layer_ids is None + else tuple(full_att_logical_layer_ids) ) - tail_layers = [(role, own_full_layers)] + tail_logical_layers_by_role = [(role, own_full_att_logical_layer_ids)] if role == "main" and cache_config.draft_full_att_kv_layer_num: - tail_layers.append( + tail_logical_layers_by_role.append( ( "draft0", tuple( @@ -105,11 +108,11 @@ def __init__( ), ) ) - for role, layers in tail_layers: - if layers: + for role, logical_layer_ids in tail_logical_layers_by_role: + if logical_layer_ids: state = create_kpool_state( config=KPoolStateConfig( - layer_ids=layers, + logical_layer_ids=logical_layer_ids, pool_size=cache_config.index_kpool, head_dim=cache_config.index_head_dim, dtype=cache_config.full_att_dtype, @@ -132,14 +135,16 @@ def bind_storage( """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" assert self.full_att_allocator is None, "attention storage is already bound" assert all( - layer < self.cache_config.all_layer_num for layer in full_state.config.layer_ids + logical_layer_id < self.cache_config.all_layer_num + for logical_layer_group in full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "main State must only contain main layers" assert (draft_full_state is not None) == bool( self.cache_config.draft_full_att_kv_layer_num ), "draft State must match the model configuration" if draft_full_state is not None: assert all( - layer >= self.cache_config.all_layer_num for layer in draft_full_state.config.layer_ids + logical_layer_id >= self.cache_config.all_layer_num + for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "draft State must only contain draft layers" state = self.main_kda_state diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 8f4cddb2d4..a0afccb1c0 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -40,9 +40,9 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig role = infer_state.att_resource_role tail = getattr( infer_state.att_resource_manager, f"{role}_indexer_kpool_c{self.pool_size}_tail_state" - ).get_layer_view(self.layer_idx) + ).get_layer_view(logical_layer_id=self.layer_idx) state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - packed_buffer = state.get_layer_view(self.layer_idx).indexer_k + packed_buffer = state.get_layer_view(logical_layer_id=self.layer_idx).indexer_k compress_pools( raw=raw, tail=tail, diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index b741d0ffc2..0a02b295e6 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -162,7 +162,7 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): state = getattr(infer_state.att_resource_manager, self.full_att_state_name) return infer_state.prefill_att_state.prefill_att( q=q, - k=state.get_layer_view(layer_id=self.layer_num_).kv, + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_prefill=True, @@ -244,7 +244,7 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): state = getattr(infer_state.att_resource_manager, self.full_att_state_name) return infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=state.get_layer_view(layer_id=self.layer_num_).kv, + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_decode=True, diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 2fc9559e65..c4f9af4c22 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -26,18 +26,18 @@ class Glm5NextTpPartModel(TpPartBaseModel): transformer_layer_infer_class = Glm5NextTransformerLayerInfer @classmethod - def create_bounded_attention_definition(cls, model_config, args, *, model_layer_ids, tp_rank=0, tp_world_size=1): + def create_bounded_attention_definition(cls, model_config, args, *, model_logical_layer_ids, tp_rank=0, tp_world_size=1): from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig from lightllm.common.basemodel.basemodel_utils import get_resource_config config = get_resource_config(model_config, model_class=cls) cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) - layers = tuple(layer for i, layer in enumerate(model_layer_ids) if (i + 1) % cfg.full_attention_interval) - if not layers: + logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) + if not logical_layer_ids: return None state = create_kda_state( config=ConvSSMStateConfig( - layer_ids=layers, + logical_layer_ids=logical_layer_ids, global_k_heads=cfg.global_linear_k_heads, global_v_heads=cfg.global_linear_v_heads, key_dim=cfg.head_linear_k_dim, @@ -96,7 +96,7 @@ def _init_att_resource_manager(self): if not self.is_mtp_draft_model: self.att_resource_manager = Glm5NextAttResourceManager( self.linear_config, self.max_req_num + 1, self.args.mtp_step, - tp_rank=self.tp_rank_, linear_layer_ids=(), full_layer_ids=(), + tp_rank=self.tp_rank_, bounded_logical_layer_ids=(), full_att_logical_layer_ids=(), ) init_model_bounded_attention(self) init_model_att_tail(self) @@ -107,7 +107,7 @@ def _full_attention_options(cls, config, args): if args.llm_kv_type != "None": raise ValueError(f"GLM K-pool does not support llm_kv_type={args.llm_kv_type}") cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) - return dict(heads=1, dim=config["kv_lora_rank"], layer_ids=cfg.get_full_att_layer_ids(), packed=True) + return dict(heads=1, dim=config["kv_lora_rank"], logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in cfg.get_full_att_logical_layer_ids()), packed=True) def _init_att_backend(self): self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 43b24ab354..294c350e91 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -80,7 +80,7 @@ def _resource_config(cls, model_config): @classmethod def _full_attention_options(cls, config, args): - return dict(heads=1, dim=config["kv_lora_rank"], layer_ids=(0,), packed=True) + return dict(heads=1, dim=config["kv_lora_rank"], logical_layer_groups=((0,),), packed=True) def _init_att_backend(self): backend_class = ( diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index 1cd117f336..f698747191 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -73,7 +73,7 @@ def _context_attention_kernel( use_sliding_window = False state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -101,7 +101,7 @@ def _token_attention_kernel( use_sliding_window = False state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 12610ede89..36ea15e0b6 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -45,7 +45,7 @@ def _context_attention_kernel( layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -63,7 +63,7 @@ def _token_attention_kernel( layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - _k, _v = state.get_layer_view(layer_id=self.layer_num_) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor) return o_tensor.view(q.shape) diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b80ebb564f..20afacf6f8 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -14,7 +14,7 @@ def __init__(self, layer_num, network_config): def _token_attention_kernel(self, q, infer_state, layer_weight): state = getattr(infer_state.att_resource_manager, self.full_att_state_name) - k, v = state.get_layer_view(layer_id=self.layer_num_) + k, v = state.get_layer_view(logical_layer_id=self.layer_num_) k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) q = q.view(-1, self.tp_q_head_num_, self.head_dim_) diff --git a/lightllm/models/qwen3next/att_resource_config.py b/lightllm/models/qwen3next/att_resource_config.py index 3c6f45cc40..5c78a7cc7c 100644 --- a/lightllm/models/qwen3next/att_resource_config.py +++ b/lightllm/models/qwen3next/att_resource_config.py @@ -74,7 +74,7 @@ def get_main_model_full_att_layer_num(self): def get_full_att_kv_layer_num_with_draft_model(self): return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num - def get_full_att_layer_ids(self): + def get_full_att_logical_layer_ids(self): """Logical model layers stored in each physical full-attention row.""" return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) + tuple( range(self.all_layer_num, self.all_layer_num + self.draft_full_att_kv_layer_num) diff --git a/lightllm/models/qwen3next/att_resource_manager.py b/lightllm/models/qwen3next/att_resource_manager.py index d9fdefe593..63438488ab 100644 --- a/lightllm/models/qwen3next/att_resource_manager.py +++ b/lightllm/models/qwen3next/att_resource_manager.py @@ -29,17 +29,17 @@ def __init__( tp_rank: int = 0, device: str | torch.device | int = "cuda", role: str = "main", - layer_ids: Sequence[int] | None = None, + logical_layer_ids: Sequence[int] | None = None, ): super().__init__(request_capacity=request_capacity) self.cache_config = cache_config - layers = ( + logical_layer_ids = ( tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) - if layer_ids is None - else tuple(layer_ids) + if logical_layer_ids is None + else tuple(logical_layer_ids) ) config = ConvSSMStateConfig( - layer_ids=layers, + logical_layer_ids=logical_layer_ids, global_k_heads=cache_config.global_linear_k_heads, global_v_heads=cache_config.global_linear_v_heads, key_dim=cache_config.head_linear_k_dim, @@ -50,7 +50,7 @@ def __init__( tp_rank=tp_rank, tp_world_size=cache_config.tp_world_size, ) - if layers: + if logical_layer_ids: state = create_gdn_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) self.set_att_state(role, state) @@ -66,14 +66,16 @@ def bind_storage( """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" assert self.full_att_allocator is None, "attention storage is already bound" assert all( - layer < self.cache_config.all_layer_num for layer in full_state.config.layer_ids + logical_layer_id < self.cache_config.all_layer_num + for logical_layer_group in full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "main State must only contain main layers" assert (draft_full_state is not None) == bool( self.cache_config.draft_full_att_kv_layer_num ), "draft State must match the model configuration" if draft_full_state is not None: assert all( - layer >= self.cache_config.all_layer_num for layer in draft_full_state.config.layer_ids + logical_layer_id >= self.cache_config.all_layer_num + for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "draft State must only contain draft layers" state = self.main_gdn_state diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 5bbc5ffc6f..13263d56f4 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -36,18 +36,18 @@ def __init__(self, kvargs) -> None: super().__init__(kvargs) @classmethod - def create_bounded_attention_definition(cls, model_config, args, *, model_layer_ids, tp_rank=0, tp_world_size=1): + def create_bounded_attention_definition(cls, model_config, args, *, model_logical_layer_ids, tp_rank=0, tp_world_size=1): from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig from lightllm.common.basemodel.basemodel_utils import get_resource_config config = get_resource_config(model_config, model_class=cls) cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) - layers = tuple(layer for i, layer in enumerate(model_layer_ids) if (i + 1) % cfg.full_attention_interval) - if not layers: + logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) + if not logical_layer_ids: return None state = create_gdn_state( config=ConvSSMStateConfig( - layer_ids=layers, + logical_layer_ids=logical_layer_ids, global_k_heads=cfg.global_linear_k_heads, global_v_heads=cfg.global_linear_v_heads, key_dim=cfg.head_linear_k_dim, @@ -91,7 +91,7 @@ def _init_att_resource_manager(self): raise ValueError("MTP linear attention does not support mixed prefill/decode") if not self.is_mtp_draft_model: self.att_resource_manager = Qwen3NextAttResourceManager( - self.linear_config, self.max_req_num + 1, start_args.mtp_step, tp_rank=self.tp_rank_, layer_ids=() + self.linear_config, self.max_req_num + 1, start_args.mtp_step, tp_rank=self.tp_rank_, logical_layer_ids=() ) init_model_bounded_attention(self) init_model_full_attention(self) @@ -102,7 +102,7 @@ def _full_attention_options(cls, config, args): raise ValueError(f"Qwen3Next does not support llm_kv_type={args.llm_kv_type}") cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) return dict( - heads=cfg.full_att_all_num_kv_heads, dim=cfg.full_att_head_dim, layer_ids=cfg.get_full_att_layer_ids() + heads=cfg.full_att_all_num_kv_heads, dim=cfg.full_att_head_dim, logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in cfg.get_full_att_logical_layer_ids()) ) def _init_att_backend1(self): diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/att_resource_config.py index be0b66ab44..d8bcd9c5ea 100644 --- a/lightllm/utils/att_resource_config.py +++ b/lightllm/utils/att_resource_config.py @@ -17,7 +17,7 @@ from lightllm.utils.torch_dtype_utils import get_torch_dtype if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Mapping from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs @@ -26,40 +26,47 @@ def load_static_fp8_calibration( args: StartArgs | EasyDict, *, per_head: bool, - layers: int, + storage_layer_count: int, heads: int, tp_rank: int = 0, tp_world_size: int = 1, - layer_start: int | None = None, + storage_layer_start: int | None = None, ): - """Load and validate the same fixed parameters for model and cache metadata. + """读取模型和 CPU cache 共用的静态标定参数。 - Returned tensors remain on CPU so validation and capacity planning precede - GPU allocation. TP replicas select the same calibrated head values. + storage_layer_count 是当前 State 的物理层数；storage_layer_start 是其在标定表中的 + 物理行起点，两者都不是逻辑层号。共享逻辑层使用同一行标定参数，不重复占用行。 + 返回 Tensor 保留在 CPU；当前 TP 分片在每行中选取对应的 head 参数。 """ path = args.kv_quant_calibration_config_path if path is None: if per_head: raise ValueError("fp8kv_sph requires kv_quant_calibration_config_path with q_calibration") - return {"scale": "unit"}, torch.ones((layers, 2), dtype=torch.float32), None + return {"scale": "unit"}, torch.ones((storage_layer_count, 2), dtype=torch.float32), None with open(path) as f: cfg = json.load(f) - stored_layers = cfg["num_layers"] - draft_layers = get_added_mtp_kv_layer_num() if layer_start is None else 0 - if layer_start is not None: - if layer_start < 0 or layer_start + layers > stored_layers: + calibrated_storage_layer_count = cfg["num_layers"] + expected_draft_storage_layer_count = get_added_mtp_kv_layer_num() if storage_layer_start is None else 0 + if storage_layer_start is not None: + if storage_layer_start < 0 or storage_layer_start + storage_layer_count > calibrated_storage_layer_count: raise ValueError("calibration does not cover this model's attention layers") elif per_head and ("num_target_layers" in cfg or "num_draft_layers" in cfg): - target, draft = cfg["num_target_layers"], cfg["num_draft_layers"] + calibrated_target_storage_layer_count = cfg["num_target_layers"] + calibrated_draft_storage_layer_count = cfg["num_draft_layers"] if ( - draft < 0 - or target + draft != stored_layers - or target != layers - draft_layers - or (draft_layers > 0 and draft != draft_layers) + calibrated_draft_storage_layer_count < 0 + or calibrated_target_storage_layer_count + calibrated_draft_storage_layer_count != calibrated_storage_layer_count + or calibrated_target_storage_layer_count != storage_layer_count - expected_draft_storage_layer_count + or ( + expected_draft_storage_layer_count > 0 + and calibrated_draft_storage_layer_count != expected_draft_storage_layer_count + ) ): raise ValueError("calibration target/draft layer layout does not match the runtime") - elif stored_layers != layers: - raise ValueError(f"calibration num_layers {stored_layers} does not match runtime {layers}") + elif calibrated_storage_layer_count != storage_layer_count: + raise ValueError( + f"calibration num_layers {calibrated_storage_layer_count} does not match runtime {storage_layer_count}" + ) if not per_head: from lightllm.utils.config_utils import get_model_architectures @@ -79,8 +86,8 @@ def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]): raise ValueError(f"calibration scales and scales_shape must match {shape}") if not bool((torch.isfinite(tensor) & (tensor > 0)).all()): raise ValueError("calibration scales must be finite and positive") - start = layer_start or 0 - return tensor[start : start + layers] + first_storage_layer_index = storage_layer_start or 0 + return tensor[first_storage_layer_index : first_storage_layer_index + storage_layer_count] def shard(tensor: torch.Tensor, calibrated_heads: int): if calibrated_heads <= 0 or (heads * tp_world_size) % calibrated_heads: @@ -90,15 +97,15 @@ def shard(tensor: torch.Tensor, calibrated_heads: int): if per_head: calibrated_heads = cfg["num_head"] - all_scales = scale_tensor(cfg, (stored_layers, 2 * calibrated_heads)) + all_scales = scale_tensor(cfg, (calibrated_storage_layer_count, 2 * calibrated_heads)) k, v = all_scales.split(calibrated_heads, dim=-1) scales = torch.cat((shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) q_cfg = cfg.get("q_calibration") if not q_cfg: raise ValueError("fp8kv_sph calibration config requires q_calibration") - q_scales = shard(scale_tensor(q_cfg, (stored_layers, q_cfg["num_head"])), q_cfg["num_head"]) + q_scales = shard(scale_tensor(q_cfg, (calibrated_storage_layer_count, q_cfg["num_head"])), q_cfg["num_head"]) else: - scales, q_scales = scale_tensor(cfg, (stored_layers, 2)), None + scales, q_scales = scale_tensor(cfg, (calibrated_storage_layer_count, 2)), None return cfg, scales, q_scales @@ -107,19 +114,19 @@ def build_full_attention_state( *, heads: int, dim: int, - layer_ids: Sequence[int], + logical_layer_groups: tuple[tuple[int, ...], ...], mla: bool = False, packed: bool = False, dtype: torch.dtype | None = None, replicated: bool = False, tp_rank: int = 0, tp_world_size: int = 1, - calibration_layer_start: int | None = None, + calibration_storage_layer_start: int | None = None, ) -> FullAttentionState: """将模型和启动配置转换为算法 create 参数，返回容量规划所需的 meta State。""" dtype = get_torch_dtype(args.data_type) if dtype is None else dtype config = FullAttentionConfig( - layer_ids=tuple(layer_ids), + logical_layer_groups=logical_layer_groups, global_head_num=heads, head_dim=dim, dtype=dtype, @@ -137,11 +144,11 @@ def build_full_attention_state( calibration, scales, q_scales = load_static_fp8_calibration( args, per_head=args.llm_kv_type == "fp8kv_sph", - layers=len(layer_ids), + storage_layer_count=len(logical_layer_groups), heads=config.head_num, tp_rank=tp_rank, tp_world_size=tp_world_size, - layer_start=calibration_layer_start, + storage_layer_start=calibration_storage_layer_start, ) calibration_id = hashlib.sha256(json.dumps(calibration, sort_keys=True).encode()).hexdigest() return create_gqa_state( @@ -174,22 +181,22 @@ def build_cpu_cache_layouts(): models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 - layer_start = calibration_start = 0 + logical_layer_start = calibration_storage_layer_start = 0 for role, model_class, config in models: count = get_resource_config(config, model_class=model_class)["num_hidden_layers"] - layer_ids = tuple(range(layer_start, layer_start + count)) + logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) full = create_full_attention_definition( - model_class, config, args, model_layer_ids=layer_ids, calibration_layer_start=calibration_start + model_class, config, args, model_logical_layer_ids=logical_layer_ids, calibration_storage_layer_start=calibration_storage_layer_start ) bounded = create_bounded_attention_definition( - model_class, config, args, model_layer_ids=layer_ids + model_class, config, args, model_logical_layer_ids=logical_layer_ids ) for state in (full, bounded): if state is None: continue manager.set_att_state(role, state) - layer_start += count - calibration_start += len(full.config.layer_ids) if full is not None else 0 + logical_layer_start += count + calibration_storage_layer_start += len(full.config.logical_layer_groups) if full is not None else 0 manager.finish_state_creation() if manager.get_bounded_attention_states(): diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index f28f549022..e7fc6f238b 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -53,7 +53,7 @@ def main(): ) tokens = args.tokens_per_page * args.pages full = GQAFullAttentionState( - FullAttentionConfig(tuple(range(12)), 1, 256, torch.bfloat16, 0, 8), + FullAttentionConfig(tuple((layer,) for layer in tuple(range(12))), 1, 256, torch.bfloat16, 0, 8), ) full.initialize_storage(tokens) resources = Qwen3NextAttResourceManager(cfg, request_capacity=1, mtp_step=0) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 38ddc83fbe..49b27d72e8 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -67,11 +67,11 @@ def main_linear_state(resources: AttResourceManager): return resources.main_gdn_state -def prepare_model_states(model: TpPartBaseModel, *, layer_ids: Sequence[int] | None = None): +def prepare_model_states(model: TpPartBaseModel, *, logical_layer_ids: Sequence[int] | None = None): config = model.config.get("text_config", model.config) - if layer_ids is None: - layer_ids = range(config.get("num_hidden_layers", config.get("n_layer"))) - model.layers_infer = [SimpleNamespace(layer_num_=i) for i in layer_ids] + if logical_layer_ids is None: + logical_layer_ids = range(config.get("num_hidden_layers", config.get("n_layer"))) + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in logical_layer_ids] if not hasattr(model, "torch_memory_saver"): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) model.prepare_att_resources() @@ -81,16 +81,18 @@ def create_full_attention_roles( template: FullAttentionState, main_layer_count: int, ) -> tuple[FullAttentionState, FullAttentionState | None]: """仅用于测试：按配置层号创建各自拥有独立存储的 main 和单个 draft。""" - layers = template.config.layer_ids - boundary = sum(layer < main_layer_count for layer in layers) - assert boundary and all(layer < main_layer_count for layer in layers[:boundary]) - assert all(layer >= main_layer_count for layer in layers[boundary:]), "draft layers must follow main layers" + layers = template.config.logical_layer_groups + boundary = sum(all(layer < main_layer_count for layer in group) for group in layers) + assert boundary and all(layer < main_layer_count for group in layers[:boundary] for layer in group) + assert all(layer >= main_layer_count for group in layers[boundary:] for layer in group), ( + "draft layers must follow main layers and cannot share main storage" + ) states = [] for start, end in ((0, boundary), (boundary, len(layers))): if start == end: states.append(None) continue - config = replace(template.config, layer_ids=layers[start:end]) + config = replace(template.config, logical_layer_groups=layers[start:end]) if isinstance(template, StaticFP8FullAttentionState): state = rebuild_full_state( template, config=config, scales=template.scales[start:end].clone(), diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index bf932d6a17..98122c4b9d 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -250,7 +250,7 @@ def start_runtime(): if main_linear: assert manager.main_gdn_state.mtp_state_index[0] == draft_count for i, state in enumerate(draft_states): - assert len(state.config.layer_ids) == 1 + assert len(state.config.logical_layer_ids) == 1 manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([i + 2], dtype=torch.int32, device="cuda"), role=f"draft{i}", @@ -265,7 +265,7 @@ def start_runtime(): assert {buffer.data_ptr() for buffer in full_state_test_buffers(main_state).values()}.isdisjoint( buffer.data_ptr() for buffer in full_state_test_buffers(draft_state).values() ) - assert len(draft_state.config.layer_ids) == 1 + assert len(draft_state.config.logical_layer_groups) == 1 else: assert not manager.get_full_attention_states(f"draft{i}") diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index a710d3de6b..72e6316ffc 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -30,7 +30,7 @@ def make_state( mtp_step: int = 2, ): config = ConvSSMStateConfig( - layer_ids=(1, 3), global_k_heads=heads, global_v_heads=heads, key_dim=2, value_dim=2, conv_width=2, + logical_layer_ids=(1, 3), global_k_heads=heads, global_v_heads=heads, key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, ) state = cls(config, request_capacity=3, mtp_step=mtp_step, device=device) @@ -42,10 +42,10 @@ def make_state( def page_layout(state: KDAState | GDNState): cfg = state.config - conv_shape = (cfg.conv_channels, len(cfg.layer_ids), cfg.conv_width) - conv_bytes = cfg.conv_channels * len(cfg.layer_ids) * cfg.conv_width * 4 - ssm_shape = (len(cfg.layer_ids), cfg.num_v_heads, cfg.key_dim, cfg.value_dim) - ssm_bytes = len(cfg.layer_ids) * cfg.num_v_heads * cfg.key_dim * cfg.value_dim * 4 + conv_shape = (cfg.conv_channels, len(cfg.logical_layer_ids), cfg.conv_width) + conv_bytes = cfg.conv_channels * len(cfg.logical_layer_ids) * cfg.conv_width * 4 + ssm_shape = (len(cfg.logical_layer_ids), cfg.num_v_heads, cfg.key_dim, cfg.value_dim) + ssm_bytes = len(cfg.logical_layer_ids) * cfg.num_v_heads * cfg.key_dim * cfg.value_dim * 4 end = 16 + conv_bytes + ssm_bytes return StateTransferLayout( resource_id=state.resource_id, @@ -94,7 +94,7 @@ def test_prefill_state_is_saved_and_restored( if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") config = ConvSSMStateConfig( - layer_ids=(1, 3), global_k_heads=1, global_v_heads=1, key_dim=2, value_dim=2, conv_width=2, + logical_layer_ids=(1, 3), global_k_heads=1, global_v_heads=1, key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, ) state = cls(config, request_capacity=3, mtp_step=2, device=device) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index e4dc14b509..4515812be5 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -39,7 +39,7 @@ def definition(kind: str): - config = FullAttentionConfig((0, 2), 2, 16, torch.bfloat16) + config = FullAttentionConfig(((0,), (2,)), 2, 16, torch.bfloat16) if kind in ("int4", "int8"): return GroupQuantizedFullAttentionState( GroupQuantizedAttentionConfig(**vars(config), bits=4 if kind == "int4" else 8), @@ -158,19 +158,19 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( pd_args: EasyDict, ): config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=2, global_v_heads=2 if state_cls is KDAState else 4, key_dim=3, + logical_layer_ids=(0,), global_k_heads=2, global_v_heads=2 if state_cls is KDAState else 4, key_dim=3, value_dim=3 if state_cls is KDAState else 5, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) tail = KPoolAttentionTailState( - KPoolStateConfig(layer_ids=(1,), pool_size=4, head_dim=3, dtype=torch.bfloat16), + KPoolStateConfig(logical_layer_ids=(1,), pool_size=4, head_dim=3, dtype=torch.bfloat16), request_capacity=3, mtp_step=3, device="cpu", ) reference = state_cls(config, request_capacity=3, mtp_step=3, device="cpu") conv = reference.conv_state[:, 0, :, : config.conv_width] ssm = reference.ssm_state[:, 0] bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() - tail_bytes = len(tail.config.layer_ids) * tail.config.pool_size * 2 * tail.config.head_dim * tail.config.dtype.itemsize + tail_bytes = len(tail.config.logical_layer_ids) * tail.config.pool_size * 2 * tail.config.head_dim * tail.config.dtype.itemsize manager = AttResourceManager(request_capacity=3) manager.set_att_state("main", reference) manager.set_att_state("main", tail) @@ -214,7 +214,7 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( (GQAFullAttentionState, "main", 3, torch.bfloat16), (GQAFullAttentionState, "main", 1, torch.float32), ): - config = FullAttentionConfig((0,), 1, dim, dtype) + config = FullAttentionConfig(((0,),), 1, dim, dtype) manager.set_att_state(role, state_cls(config, page_size=1, device="meta")) manager.finish_state_creation() # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 @@ -261,7 +261,7 @@ def _get_transfer_regions(self, page_tokens: int): pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 config = GroupQuantizedAttentionConfig( - layer_ids=(0,), global_head_num=1, head_dim=128, + logical_layer_groups=((0,),), global_head_num=1, head_dim=128, dtype=torch.float32, bits=8, group_size=128, ) manager = AttResourceManager(request_capacity=1) @@ -328,11 +328,11 @@ def get_pd_transfer_regions(self): manager = AttResourceManager(request_capacity=1) if full_page_tokens is not None: pd_args.pd_kv_page_size = full_page_tokens - config = FullAttentionConfig((0,), 1, 1, torch.float32) + config = FullAttentionConfig(((0,),), 1, 1, torch.float32) manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) state = PaddedGDNState( ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ), request_capacity=1, device="cpu", @@ -390,7 +390,7 @@ def _get_transfer_regions(self, page_tokens: int): pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig((0,), 1, 3, torch.bfloat16) + config = FullAttentionConfig(((0,),), 1, 3, torch.bfloat16) manager.set_att_state("main", PaddedState(config, page_size=1, device="meta")) manager.finish_state_creation() get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes @@ -412,19 +412,19 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( pd_args: EasyDict, ): manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig((0,), 1, 1, torch.float32) + config = FullAttentionConfig(((0,),), 1, 1, torch.float32) full_state = GQAFullAttentionState(config, page_size=1, device="meta") manager.set_att_state("main", full_state) for cls, role in ((KDAState, "main"), (GDNState, "draft0")): config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, conv_dtype=torch.float32, ssm_dtype=torch.float32, ) state = cls(config, request_capacity=1, device="meta") manager.set_att_state(role, state) if with_tail: state = KPoolAttentionTailState( - KPoolStateConfig(layer_ids=(0,), pool_size=4, head_dim=1, dtype=torch.float32), + KPoolStateConfig(logical_layer_ids=(0,), pool_size=4, head_dim=1, dtype=torch.float32), request_capacity=1, mtp_step=0, device="meta", ) manager.set_att_state("main", state) @@ -450,15 +450,15 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( cpu_args.cpu_cache_token_page_size = page_tokens template = definition(kind) config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) - scales = torch.ones((len(config.layer_ids), 2 * config.head_num)) if kind == "static_head" else None - q_scales = torch.ones((len(config.layer_ids), config.head_num)) if kind == "static_head" else None + scales = torch.ones((len(config.logical_layer_groups), 2 * config.head_num)) if kind == "static_head" else None + q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None manager = AttResourceManager(request_capacity=3) for role in ("main", "draft0"): manager.set_att_state(role, rebuild_full_state(template, config=config, scales=scales, q_scales=q_scales)) for state_cls, v_heads, value_dim in ((KDAState, 1, 1), (GDNState, 2, 3)): # 奇数长度的 BF16 Conv 与 FP32 SSM 之间需要补齐，MTP 候选区不应计入页大小。 bounded_config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=v_heads, key_dim=1, value_dim=value_dim, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=v_heads, key_dim=1, value_dim=value_dim, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, tp_rank=tp_world_size - 1, tp_world_size=tp_world_size, ) @@ -495,7 +495,7 @@ def get_cpu_transfer_regions(self): manager = AttResourceManager(request_capacity=1) state = PaddedGDNState( ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ), request_capacity=1, device="meta", @@ -557,7 +557,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas manager.bounded_att_small_page_allocator = StatePageAllocator(1) for state_cls, role, value_dim in ((KDAState, "draft0", 1), (GDNState, "main", 3)): config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=value_dim, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=value_dim, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) state = state_cls(config, request_capacity=1, mtp_step=2, device="cpu") @@ -600,11 +600,11 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation pd_args.pd_kv_page_size = 3 pd_args.cpu_cache_token_page_size = 3 manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig((0,), 1, 3, torch.bfloat16) + config = FullAttentionConfig(((0,),), 1, 3, torch.bfloat16) manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) manager.set_att_state("main", GDNState( ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, + logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, conv_dtype=torch.bfloat16, ssm_dtype=torch.bfloat16, ), request_capacity=1, device="cpu", @@ -652,7 +652,7 @@ def get_cpu_transfer_regions(self): pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 config = FullAttentionConfig( - layer_ids=(0,), global_head_num=1, head_dim=2, dtype=torch.float32 + logical_layer_groups=((0,),), global_head_num=1, head_dim=2, dtype=torch.float32 ) manager = AttResourceManager(request_capacity=1) prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index 1a5eadd06b..0001baac9d 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -31,7 +31,7 @@ @pytest.mark.parametrize("capture", [False, True]) def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, capture: bool): template = definition(kind) - state = rebuild_full_state(template, config=replace(template.config, layer_ids=(3, 8)), token_capacity=8, device="cuda") + state = rebuild_full_state(template, config=replace(template.config, logical_layer_groups=((3,), (8,))), token_capacity=8, device="cuda") for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() manager = AttResourceManager(request_capacity=2) @@ -70,7 +70,7 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, actual = actual.view(torch.float8_e4m3fn).float() * state.scales[1][None, :, None] # Backends index fixed scales by physical row, even when model layer # IDs are sparse or include appended draft layers. - assert BaseAttBackend._find_layer_index(None, k, v, state) == 1 + assert BaseAttBackend._find_storage_layer_index(None, k, v, state) == 1 assert state.cpu_scales.shape == (2, 4) assert torch.equal(actual, kv) diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index dc8f428d97..88088abcf0 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -43,7 +43,7 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( scales[1] += 10 q_scales = None if kind == "static_tensor" else template.q_scales + value states.append(rebuild_full_state( - template, config=replace(template.config, layer_ids=(3, 8)), token_capacity=8, device="cuda", + template, config=replace(template.config, logical_layer_groups=((3,), (8,))), token_capacity=8, device="cuda", scales=scales, q_scales=q_scales, )) manager = AttResourceManager(request_capacity=2) @@ -120,7 +120,7 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu scales = torch.arange(1, scale_count + 1, dtype=torch.float32, device="cuda").unsqueeze(0) / 4 state = StaticFP8FullAttentionState( layout_config=StaticFP8AttentionConfig( - layer_ids=(5,), global_head_num=head_num, head_dim=128, dtype=torch.bfloat16, + logical_layer_groups=((5,),), global_head_num=head_num, head_dim=128, dtype=torch.bfloat16, scale_mode=scale_mode, calibration_id="split-storage", ), page_size=4, device="cuda", scales=scales, @@ -131,7 +131,7 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu state.v_storage.zero_() indices = torch.tensor([1, 6], dtype=torch.int32, device="cuda") kv = torch.randn((2, 2 * head_num, 128), dtype=torch.bfloat16, device="cuda") - state.write_full_att_kv(layer_id=5, token_indices=indices, kv=kv) + state.write_full_att_kv(logical_layer_id=5, token_indices=indices, kv=kv) k, v = state.get_layer_view(5) assert k.untyped_storage().data_ptr() != v.untyped_storage().data_ptr() actual = torch.cat((k[indices], v[indices]), dim=1) diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index 8ae0feb50e..a991630b1a 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -200,7 +200,7 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: cpu_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) pd_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.write_full_att_kv(state.config.layer_ids[0], cpu_ids, torch.empty(0, device="cuda")) + state.write_full_att_kv(state.config.logical_layer_groups[0][0], cpu_ids, torch.empty(0, device="cuda")) for src, dst in ((cpu_ids, cuda_ids), (cuda_ids, cpu_ids)): with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_full_att_tokens(src, dst) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index f59c6e8341..f5a03d59ef 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -24,7 +24,7 @@ def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): - prepared_full_state = cls(FullAttentionConfig((2, 5), 1, 4, torch.float32), page_size=4, device=device) + prepared_full_state = cls(FullAttentionConfig(((2,), (5,)), 1, 4, torch.float32), page_size=4, device=device) prepared_full_state.initialize_storage(12) return prepared_full_state @@ -158,7 +158,7 @@ def test_invalid_late_region_shape_cannot_partially_write_pd_page(): def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") - config = FullAttentionConfig((1, 3, 5, 7, 9, 11), 1, 4, torch.float32) + config = FullAttentionConfig(((1,), (3,), (5,), (7,), (9,), (11,)), 1, 4, torch.float32) source = GQAFullAttentionState(config, page_size=4, device=device) source.initialize_storage(12) target = GQAFullAttentionState(config, page_size=4, device=device) @@ -185,7 +185,7 @@ def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int) # 拒绝将单层数据广播到多层，失败时目标 token 保持原值。 single_layer_source = GQAFullAttentionState( - replace(config, layer_ids=config.layer_ids[:1]), page_size=4, device=device + replace(config, logical_layer_groups=config.logical_layer_groups[:1]), page_size=4, device=device ) single_layer_source.initialize_storage(12) for storage_buffer in full_state_test_buffers(single_layer_source).values(): diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index e75a38790b..5a816ae05a 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -46,8 +46,8 @@ def test_model_storage_write_and_copy_preserve_all_fields( ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") dim = config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else config.head_dim kv = torch.randn(2, (1 if "kv" in full_state_test_buffers(state) else 2) * config.head_num, dim, device="cuda", dtype=torch.bfloat16) - state.write_full_att_kv(config.layer_ids[-1], ids, kv) - if len(config.layer_ids) > 1: + state.write_full_att_kv(config.logical_layer_groups[-1][0], ids, kv) + if len(config.logical_layer_groups) > 1: assert not read_kv_storage(state)[0].any() resources.copy_full_att_tokens(ids, torch.tensor([1, 5], dtype=torch.int32, device="cuda")) for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py index 76b698b619..934803990f 100644 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -30,7 +30,7 @@ def test_gqa_pd_direct_copy_between_tp_layouts( pd_args.pd_kv_page_size = 4 # P 保存完整 heads，D 只恢复 rank 1 的一半 heads；两侧外部页格式相同。 source = GQAFullAttentionState( - layout_config=FullAttentionConfig(layer_ids=(2, 5), global_head_num=4, head_dim=128, dtype=dtype), + layout_config=FullAttentionConfig(logical_layer_groups=((2,), (5,)), global_head_num=4, head_dim=128, dtype=dtype), page_size=4, device="cuda", ) @@ -115,7 +115,7 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: # MQA 的 KV head 在 TP ranks 间复制，非首 rank 禁止写出但仍能读回。 state = GQAFullAttentionState( layout_config=FullAttentionConfig( - layer_ids=(0, 2), global_head_num=1, head_dim=128, dtype=torch.bfloat16, tp_world_size=2, tp_rank=1, + logical_layer_groups=((0,), (2,)), global_head_num=1, head_dim=128, dtype=torch.bfloat16, tp_world_size=2, tp_rank=1, ), page_size=4, device="cuda", @@ -153,7 +153,7 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt pytest.skip("CUDA peer access required") pd_args.pd_kv_page_size = 4 state = GQAFullAttentionState( - layout_config=FullAttentionConfig(layer_ids=(0, 2), global_head_num=2, head_dim=128, dtype=torch.bfloat16), + layout_config=FullAttentionConfig(logical_layer_groups=((0,), (2,)), global_head_num=2, head_dim=128, dtype=torch.bfloat16), page_size=4, device="cuda:1", ) diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index b8cd1e170f..6f767e34f0 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -18,7 +18,7 @@ @pytest.mark.parametrize("page_device", ["cpu", "cuda"]) def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_args: "EasyDict") -> None: pd_args.pd_kv_page_size = 32 - config = FullAttentionConfig(layer_ids=(0, 2, 4), global_head_num=3, head_dim=128, dtype=torch.bfloat16) + config = FullAttentionConfig(logical_layer_groups=((0,), (2,), (4,)), global_head_num=3, head_dim=128, dtype=torch.bfloat16) source = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") target = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") source.initialize_storage(token_capacity=128) diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py new file mode 100644 index 0000000000..91efbee2c5 --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -0,0 +1,133 @@ +"""逻辑层共享物理 KV：容量、视图、标定索引和外部页传输保持一致。""" + +from dataclasses import is_dataclass, replace +from typing import TYPE_CHECKING + +import pytest +import torch +from torch.utils._pytree import tree_leaves + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.full_att import FullAttentionConfig +from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition +from lightllm.models.gemma4.model import Gemma4TpPartModel +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + rebuild_full_state, +) +from unit_tests.common.att_resources.test_capacity import definition, allocated_bytes + +if TYPE_CHECKING: + from easydict import EasyDict + + +SCHEMES = ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] + + +@pytest.mark.parametrize("kind", SCHEMES) +def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") -> None: + template = definition(kind) + config = replace(template.config, logical_layer_groups=((0, 4), (8, 11))) + state = rebuild_full_state(template, config=config, token_capacity=8, device="cpu") + assert all(buffer.shape[0] == 2 for buffer in full_state_test_buffers(state).values()) + assert state.get_token_bytes() == template.get_token_bytes() + assert allocated_bytes(state, token_only=True) == state.get_token_bytes() * (8 + state.page_size) + + for getter in (state.get_cpu_transfer_regions, state.get_pd_transfer_regions): + assert all(region.shape[region.axes.index("layer")] == 2 for region in getter()) + + # 同组返回相同地址，不同组返回独立数据；包括量化 scale 和 packed indexer 的计算视图。 + views = [] + for layer_id in (0, 4, 8, 11): + view = state.get_layer_view(logical_layer_id=layer_id) + if is_dataclass(view): + view = tuple(vars(view).values()) + views.append([(tensor.data_ptr(), tensor.shape, tensor.stride()) for tensor in tree_leaves(view)]) + assert views[0] == views[1] + assert views[2] == views[3] + assert views[0] != views[2] + + if kind in ("gqa", "static_head", "static_tensor"): + k, v = state.get_layer_view(logical_layer_id=11) + assert BaseAttBackend._find_storage_layer_index(None, k, v, state) == 1 + + +@pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) +def test_invalid_layer_groups_are_rejected(groups: tuple) -> None: + with pytest.raises(AssertionError): + FullAttentionConfig(logical_layer_groups=groups, global_head_num=1, head_dim=128, dtype=torch.bfloat16) + + +def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "EasyDict") -> None: + pd_args.data_type = "bfloat16" + config = dict( + num_hidden_layers=6, num_attention_heads=4, num_key_value_heads=2, hidden_size=512, + head_dim=128, global_head_dim=128, num_kv_shared_layers=2, + layer_types=["sliding_attention", "full_attention"] * 3, + ) + state = create_full_attention_definition( + model_class=Gemma4TpPartModel, model_config=config, args=pd_args, + model_logical_layer_ids=(10, 11, 12, 13, 14, 15), + ) + # 共享消费者 14/15 分别使用最后一个同类型生产层 12/13，保持四份物理层存储。 + assert state.config.logical_layer_groups == ((10,), (11,), (12, 14), (13, 15)) + assert state.get_token_bytes() == 4 * 2 * 2 * 128 * state.config.dtype.itemsize + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", SCHEMES) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_shared_layers_transfer_each_physical_layer_once( + kind: str, transfer: str, pd_args: "EasyDict", +) -> None: + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + template = definition(kind) + source = rebuild_full_state( + template, config=replace(template.config, logical_layer_groups=((0, 4), (8,))), token_capacity=8, device="cuda", + ) + target = rebuild_full_state( + template, config=replace(template.config, logical_layer_groups=((20, 24), (28,))), token_capacity=8, device="cuda", + ) + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + for buffer in source_buffers.values(): + buffer.view(torch.uint8).random_(0, 256) + for buffer in target_buffers.values(): + buffer.zero_() + managers = [] + for state in (source, target): + manager = AttResourceManager(request_capacity=1) + manager.set_att_state(model_name="main", state=state) + manager.finish_state_creation() + managers.append(manager) + src_ids = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([1, 4, 7], dtype=torch.int32, device="cuda") + if transfer == "cpu": + pages = torch.zeros((2, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + managers[0].write_req_to_cpu( + token_indices=src_ids, cpu_page_buffers=pages, cpu_page_indices=[1], + layouts=managers[0].get_cpu_transfer_layouts(), + ) + managers[1].read_req_from_cpu( + cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=dst_ids, + layouts=managers[1].get_cpu_transfer_layouts(), + ) + else: + pages = torch.zeros((2, managers[0].get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + managers[0].write_to_pd_buffer( + pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=src_ids, + layouts=managers[0].get_pd_full_att_transfer_layouts(), + ) + managers[1].read_from_pd_buffer( + pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=dst_ids, + layouts=managers[1].get_pd_full_att_transfer_layouts(), + ) + torch.cuda.synchronize() + source_fields = full_state_test_fields(source, source_buffers) + target_fields = full_state_test_fields(target, target_buffers) + for name, field in source_fields.items(): + expected = field.tensor[:, src_ids].contiguous().view(torch.uint8) + actual = target_fields[name].tensor[:, dst_ids].contiguous().view(torch.uint8) + assert torch.equal(actual, expected) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index e109cbe0aa..0fc9f4dd2a 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -109,7 +109,7 @@ def test_calibration_slices_follow_model_configs_before_other_states_exist(tmp_p if not expected_layers: assert state is None continue - assert state.config.layer_ids == expected_layers + assert state.config.logical_layer_groups == tuple((layer,) for layer in expected_layers) assert read_kv_storage(state) is None assert torch.equal(state.scales[:, 0], scales[start:end]) if encoding == "fp8kv_sph": diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 78a6bbe228..a1f254edf3 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -120,7 +120,7 @@ def check_runtime_before_profile(device: str | torch.device | int): draft.max_req_num, draft.max_total_token_num = model.max_req_num, None draft.config = dict(model.config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) draft._init_req_manager() - prepare_model_states(draft, layer_ids=(4,)) + prepare_model_states(draft, logical_layer_ids=(4,)) assert draft.att_resource_manager is model.att_resource_manager assert read_kv_storage(owned_states(draft.att_resource_manager.__dict__, FullAttentionState)[0]) is None drafts.append(draft) @@ -135,8 +135,8 @@ def check_runtime_before_profile(device: str | torch.device | int): assert resources.full_att_allocator.size == full_state.capacity == 32 assert not hasattr(model, "mem_manager") assert not hasattr(model, "full_att_state") - assert full_state.config.layer_ids == tuple( - layer for layer in model.linear_config.get_full_att_layer_ids() if layer < model.config["num_hidden_layers"] + assert full_state.config.logical_layer_groups == tuple( + (layer,) for layer in model.linear_config.get_full_att_logical_layer_ids() if layer < model.config["num_hidden_layers"] ) assert observed == [main_linear_state(resources).conv_state.data_ptr()] assert resources.bounded_att_big_page_allocator.size == 6 @@ -155,7 +155,9 @@ def check_runtime_before_profile(device: str | torch.device | int): (draft_state,) = resources.get_full_attention_states("draft0") assert draft_state is not full_state assert draft_state.capacity == full_state.capacity - assert set(draft_state.config.layer_ids).isdisjoint(full_state.config.layer_ids) + assert {layer for group in draft_state.config.logical_layer_groups for layer in group}.isdisjoint( + layer for group in full_state.config.logical_layer_groups for layer in group + ) assert not hasattr(draft, "mem_manager") assert not hasattr(draft, "mtp_state_cursors") assert observed == [main_linear_state(resources).conv_state.data_ptr()] diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 4781c32919..e7390ce2f5 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -63,7 +63,7 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int if glm: full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - (3, *range(4, 4 + draft_layers)), 1, 584, torch.bfloat16, rank, tp, kv_dim=512 + tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), 1, 584, torch.bfloat16, rank, tp, kv_dim=512 ), page_size=4, ) @@ -71,7 +71,7 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int else: full = GQAFullAttentionState( FullAttentionConfig( - (1, 3, *range(4, 4 + draft_layers)), 2, 16, torch.bfloat16, rank, tp + tuple((layer,) for layer in (1, 3, *range(4, 4 + draft_layers))), 2, 16, torch.bfloat16, rank, tp ), page_size=4, ) diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index 609acafdb0..f3b6fd7160 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -82,6 +82,14 @@ def args(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): 256, 4, ), + ( + Gemma4TpPartModel, + {"num_key_value_heads": 4, "head_dim": 256, "global_head_dim": 256, + "num_kv_shared_layers": 1, "layer_types": ["full_attention", "full_attention"]}, + 4, + 256, + 1, + ), ] @@ -113,13 +121,18 @@ def test_model_and_cpu_service_use_identical_storage_geometry( resources.initialize_resources(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) (state,) = resources.get_full_attention_states("main") assert state.config.head_num == local_heads and state.config.head_dim == dim - assert state.k_storage.shape == state.v_storage.shape == (2, 20, local_heads, dim) + storage_layers = 2 - (config.get("num_kv_shared_layers") or 0) + assert state.k_storage.shape == state.v_storage.shape == (storage_layers, 20, local_heads, dim) + if storage_layers == 1: + assert state.config.logical_layer_groups == ((0, 1),) + assert model.layers_infer[0].full_att_state_name == model.layers_infer[1].full_att_state_name + assert state.get_layer_view(logical_layer_id=0)[0].data_ptr() == state.get_layer_view(logical_layer_id=1)[0].data_ptr() assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) assert resources.full_att_allocator.size == 16 layout = resources.get_cpu_transfer_layouts()[state.resource_id] assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) - # 两层 GQA，K/V 各 global_heads 个 head，按实际 KV 精度计算全局传输页。 - expected_pd_bytes = 2 * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() + # 按物理层数计算传输页；共享消费者不重复增加 KV 或传输容量。 + expected_pd_bytes = storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() assert resources.get_pd_page_bytes() == expected_pd_bytes fields = layout.regions assert all(region.shape[2] == global_heads for region in fields) diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index f5c50fb4cc..0393431051 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -72,5 +72,5 @@ def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: s (state,) = resources.get_full_attention_states("draft0") assert not hasattr(draft, "full_att_state") assert state.resource_id == "draft0_full_att_bf16_state" - assert state.config.layer_ids == (2,) + assert state.config.logical_layer_groups == ((2,),) assert read_kv_storage(state) is None diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index c95475996a..f813bb5802 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -28,7 +28,7 @@ def test_page_pool_allocation_strategy( pytest.skip("CUDA required for pinned memory") cpu_args.enable_dp_prompt_cache_fetch = dp_fetch config = ConvSSMStateConfig( - layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, + logical_layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, key_dim=4, value_dim=4, conv_width=3, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) @@ -99,7 +99,7 @@ def test_pinned_page_pool_survives_queue_on_another_gpu( cpu_args.enable_dp_prompt_cache_fetch = True torch.cuda.set_device(0) config = ConvSSMStateConfig( - layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, + logical_layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, key_dim=4, value_dim=4, conv_width=3, conv_dtype=torch.float32, ssm_dtype=torch.float32, ) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 40433b2820..3de2789a81 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -37,14 +37,14 @@ def make_quant_state(kind: str, device: str | torch.device | int): if kind in ("int4", "int8"): - config = GroupQuantizedAttentionConfig((1, 3), 2, 16, torch.float16, bits=int(kind[-1])) + config = GroupQuantizedAttentionConfig(((1,), (3,)), 2, 16, torch.float16, bits=int(kind[-1])) prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) prepared_full_state.initialize_storage(8) return prepared_full_state if kind.startswith("fp8_"): per_head = kind == "fp8_head" config = StaticFP8AttentionConfig( - (1, 3), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor", calibration_id="test" + ((1,), (3,)), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor", calibration_id="test" ) scales = torch.arange(1, 9 if per_head else 5, dtype=torch.float32, device=device).reshape(2, -1) / 16 prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device, scales=scales) @@ -52,7 +52,7 @@ def make_quant_state(kind: str, device: str | torch.device | int): return prepared_full_state fp8 = kind == "packed_fp8" config = PackedMLAAttentionConfig( - (1, 3), 1, 800 if fp8 else 648, torch.uint8 if fp8 else torch.bfloat16, fp8_nope_dim=512 if fp8 else 0 + ((1,), (3,)), 1, 800 if fp8 else 648, torch.uint8 if fp8 else torch.bfloat16, fp8_nope_dim=512 if fp8 else 0 ) prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) prepared_full_state.initialize_storage(8) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py index 96c3cc67d2..a7ce8744a7 100644 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -12,7 +12,7 @@ def test_quantized_kv_buffers_include_the_full_hold_page(): allocations = [] for bits in (8, 4): state = GroupQuantizedFullAttentionState( - GroupQuantizedAttentionConfig((0, 1, 2), 2, 16, torch.float16, bits=bits), + GroupQuantizedAttentionConfig(((0,), (1,), (2,)), 2, 16, torch.float16, bits=bits), page_size=4, device="cpu", ) state.initialize_storage(32) diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index b61adeef99..983fa4c6e5 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -32,7 +32,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( kv_type: str, expected_type: str, storage_dtype: torch.dtype, width: int, ): config = FullAttentionConfig( - layer_ids=(3, 7), global_head_num=8, head_dim=128, dtype=torch.bfloat16, tp_rank=2, tp_world_size=4, + logical_layer_groups=((3,), (7,)), global_head_num=8, head_dim=128, dtype=torch.bfloat16, tp_rank=2, tp_world_size=4, ) scales = torch.full((2, 4 if kv_type == "fp8kv_sph" else 2), 0.5, dtype=torch.float32) q_scales = torch.ones((2, 2), dtype=torch.float32) if kv_type == "fp8kv_sph" else None @@ -42,7 +42,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( ) assert state.get_state_type() == expected_type assert state.resource_id is None - assert state.config.layer_ids == (3, 7) + assert state.config.logical_layer_groups == ((3,), (7,)) assert state.config.global_head_num == 8 and state.config.head_num == 2 assert state.config.tp_rank == 2 and state.config.tp_world_size == 4 assert state.config.head_dim == config.head_dim and config.dtype == torch.bfloat16 @@ -73,7 +73,7 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( packed: bool, kv_type: str, kv_dim: int, row_width: int, expected_type: str, ): config = FullAttentionConfig( - layer_ids=(2, 6), global_head_num=1, head_dim=kv_dim, dtype=torch.bfloat16, tp_rank=1, tp_world_size=4, + logical_layer_groups=((2,), (6,)), global_head_num=1, head_dim=kv_dim, dtype=torch.bfloat16, tp_rank=1, tp_world_size=4, ) state = create_mla_state(config=config, packed=packed, kv_type=kv_type, page_size=4, device="meta") assert state.get_state_type() == expected_type @@ -91,7 +91,7 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( @pytest.mark.parametrize("create_state", [create_kda_state, create_gdn_state]) def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Callable[..., BoundedAttentionState]): config = ConvSSMStateConfig( - layer_ids=(0, 2), global_k_heads=2, global_v_heads=2, + logical_layer_ids=(0, 2), global_k_heads=2, global_v_heads=2, key_dim=8, value_dim=8, conv_width=3, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) state = create_state(config=config, request_capacity=3, mtp_step=2, device="cpu") @@ -103,7 +103,7 @@ def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Calla def test_kpool_creation_preserves_pool_and_mtp_dimensions(): - config = KPoolStateConfig(layer_ids=(3, 7), pool_size=4, head_dim=128, dtype=torch.bfloat16) + config = KPoolStateConfig(logical_layer_ids=(3, 7), pool_size=4, head_dim=128, dtype=torch.bfloat16) state = create_kpool_state(config=config, request_capacity=3, mtp_step=2, device="cpu") assert state.config is config and state.resource_id is None assert state.tail_buffer.shape == (2, 3, 6, 256) @@ -111,7 +111,7 @@ def test_kpool_creation_preserves_pool_and_mtp_dimensions(): def test_create_rejects_encoding_from_another_logical_family(): - config = FullAttentionConfig(layer_ids=(0,), global_head_num=1, head_dim=576, dtype=torch.bfloat16) + config = FullAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=576, dtype=torch.bfloat16) with pytest.raises(AssertionError, match="unsupported GQA encoding"): create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") with pytest.raises(AssertionError, match="unsupported packed MLA encoding"): diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 34d0169666..4bc1d07c7c 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -155,7 +155,7 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp cpu_args.cpu_cache_token_page_size = 4 physical = allocated_state(kind, device="cuda") main, draft = create_full_attention_roles(physical, 1) - assert main.config.layer_ids == (0,) and draft.config.layer_ids == (2,) + assert main.config.logical_layer_groups == ((0,),) and draft.config.logical_layer_groups == ((2,),) for role, state, layer in (("main", main, 0), ("draft0", draft, 1)): physical_buffers = full_state_test_buffers(physical) for name, buffer in full_state_test_buffers(state).items(): @@ -203,7 +203,7 @@ def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_la manager.set_att_state("draft0", draft) manager.full_att_allocator = FullAttKVAllocator(main.capacity, main.page_size, pin_memory=False) template = allocated_state("int8") - quant = rebuild_full_state(template, config=replace(template.config, layer_ids=(4, 6))) + quant = rebuild_full_state(template, config=replace(template.config, logical_layer_groups=((4,), (6,)))) manager.set_att_state("main", quant) if overlap: overlapping = allocated_state("static_head") diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py index 20ca398659..9840d4622e 100644 --- a/unit_tests/common/att_resources/test_tail_state.py +++ b/unit_tests/common/att_resources/test_tail_state.py @@ -11,7 +11,7 @@ def make_tail(*, device: str | torch.device | int = "cpu", mtp_step: int = 2): return KPoolAttentionTailState( - KPoolStateConfig(layer_ids=(3, 7), pool_size=4, head_dim=2, dtype=torch.float32), + KPoolStateConfig(logical_layer_ids=(3, 7), pool_size=4, head_dim=2, dtype=torch.float32), request_capacity=3, mtp_step=mtp_step, device=device ) diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index 743cb67c63..ffa41946c8 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -57,7 +57,7 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( if not torch.cuda.is_available(): pytest.skip("CUDA required") cpu_args.cpu_cache_token_page_size = 4 - config = FullAttentionConfig((0,), heads, 4, torch.float32, replicated=replicated) + config = FullAttentionConfig(((0,),), heads, 4, torch.float32, replicated=replicated) reference = state_cls(config, page_size=4, device="cuda") reference.initialize_storage(4) for storage_buffer in full_state_test_buffers(reference).values(): @@ -91,7 +91,7 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( def test_bounded_config_keeps_distinct_global_heads_when_local_heads_are_replicated(cpu_args: EasyDict): cpu_args.cpu_cache_token_page_size = 4 config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=3, conv_width=2, + logical_layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=3, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_rank=7, tp_world_size=8, ) source = GDNState(config, request_capacity=1, device="meta") @@ -120,15 +120,15 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( pd_args.pd_kv_page_size = page_tokens if bounded: config = ConvSSMStateConfig( - layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=2, conv_width=2, + logical_layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_world_size=8, ) source = GDNState(config, request_capacity=1, device="meta") - other_config = replace(config, layer_ids=(1,), global_k_heads=4) + other_config = replace(config, logical_layer_ids=(1,), global_k_heads=4) else: - config = FullAttentionConfig((0,), 2, 4, torch.float32, tp_world_size=8) + config = FullAttentionConfig(((0,),), 2, 4, torch.float32, tp_world_size=8) source = GQAFullAttentionState(config, page_size=4, device="meta") - other_config = replace(config, layer_ids=(1,), global_head_num=4) + other_config = replace(config, logical_layer_groups=((1,),), global_head_num=4) if bounded: other = type(source)( other_config, request_capacity=source.request_capacity, mtp_step=source.mtp_step, device="meta" diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index e1d9131cd4..414a1d766a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -85,7 +85,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t ) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - config.get_full_att_layer_ids(), + tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), 1, 584, torch.bfloat16, @@ -163,7 +163,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t # belong to live requests. Bytewise equality checks FP8 scale integrity. packed_bytes = full.storage.view(torch.uint8) packed_bytes[:, 0].random_(0, 256) - full.copy_full_att_tokens(torch.tensor([0]), torch.tensor([7])) + full.copy_full_att_tokens(torch.tensor([0], device="cuda"), torch.tensor([7], device="cuda")) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index db1186b792..cd3af18217 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -585,11 +585,12 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) ragged = torch.randperm(tokens + 9, device="cuda", dtype=torch.int32)[:tokens] full_att_state = SimpleNamespace( - get_layer_view=lambda _: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), + get_layer_view=lambda logical_layer_id: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), ) infer = SimpleNamespace( + att_resource_role="main", att_resource_manager=SimpleNamespace( - main_indexer_kpool_c4_tail_state=SimpleNamespace(get_layer_view=lambda _: tail), + main_indexer_kpool_c4_tail_state=SimpleNamespace(get_layer_view=lambda logical_layer_id: tail), main_mla_packed_bf16_state=full_att_state, ), is_draft_model=False, diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 33a7ff73a4..455b90b1c9 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -173,7 +173,7 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.full_att_state_name = "main_mla_packed_bf16_state" state.att_resource_manager = SimpleNamespace( main_mla_packed_bf16_state=SimpleNamespace( - get_layer_view=lambda layer_id: SimpleNamespace(kv=kv) + get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv) ) ) state.prefill_att_state = Glm5NextSparsePrefillState() @@ -433,7 +433,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): assert state is not main.att_resource_manager.main_mla_packed_bf16_state assert not hasattr(draft, "full_att_state") assert state.storage.device.type == "meta" - assert state.config.layer_ids == (4 + step,) + assert state.config.logical_layer_groups == ((4 + step,),) assert state.resource_id == f"draft{step}_mla_packed_bf16_state" assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None assert getattr(draft.att_resource_manager, f"draft{step}_indexer_kpool_c4_tail_state").tail_buffer.is_cuda diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 9ccab62b6a..ef3879a991 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -193,7 +193,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): ) backend.linear_state = SimpleNamespace( mtp_state_index=accepted_offsets, - get_layer_view=lambda layer: (conv, ssm), + get_layer_view=lambda logical_layer_id: (conv, ssm), ) manager = SimpleNamespace(HOLD_REQUEST_ID=3) infer = InferStateInfo() diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 067643b514..a6074a4661 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -73,7 +73,7 @@ def create(tp, target_mtp_step=None): monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - config.get_full_att_layer_ids(), + tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), 1, 584, torch.bfloat16, @@ -115,7 +115,7 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) mtp_size = get_env_start_args().mtp_step + 1 dst_mtp_size = decode_mtp_step + 1 - layers = sum(len(state.config.layer_ids) for state in source[0].full_states) + layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() @@ -130,14 +130,14 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) offset = 0 for state in mem.full_states: - end = offset + len(state.config.layer_ids) + end = offset + len(state.config.logical_layer_groups) state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] offset = end mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) offset = 0 for state in mem.tails: - end = offset + len(state.config.layer_ids) + end = offset + len(state.config.logical_layer_ids) state.tail_buffer[:, src_req].copy_(tail[offset:end]) offset = end @@ -198,8 +198,8 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) for state, source_state in zip(mem.full_states, source[0].full_states): expected = full_state_test_fields(source_state, {"kv": source_state.storage}) - for field in full_state_test_fields(state, {"kv": state.storage}).values(): - source_field = expected[field.name] + for name, field in full_state_test_fields(state, {"kv": state.storage}).items(): + source_field = expected[name] assert torch.equal( field.tensor[:, dst_indexes].contiguous().view(torch.uint8), source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), @@ -251,11 +251,16 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem assert torch.equal(*outputs) -def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch): +def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - with pytest.raises(AssertionError, match="must be at least the complete attention runtime"): - mem.resources.get_pd_page_bytes() + full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_full_att_transfer_layouts().values()) + runtime_bytes = max( + layout.byte_end for layout in mem.resources.get_pd_bounded_att_or_tail_transfer_layouts().values() + ) + assert runtime_bytes > full_bytes + assert mem.resources.get_pd_page_bytes() == runtime_bytes + assert "smaller than the complete attention runtime" in caplog.text @pytest.mark.parametrize("source_step,target_step", [(0, 1), (0, 5), (3, 2), (5, 1)]) @@ -266,7 +271,7 @@ def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, from lightllm.common.att_resources.types import StateTransferLayout torch.manual_seed(53) - config = KPoolStateConfig(layer_ids=(0,), pool_size=4, head_dim=128, dtype=torch.bfloat16) + config = KPoolStateConfig(logical_layer_ids=(0,), pool_size=4, head_dim=128, dtype=torch.bfloat16) source = create_kpool_state(config=config, request_capacity=2, mtp_step=source_step, device="cuda") target = create_kpool_state(config=config, request_capacity=2, mtp_step=target_step, device="cuda") raw = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index de2fb835cb..837296705c 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -174,7 +174,7 @@ def make_runtime( draft.config = dict(config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) draft.max_total_token_num, draft.max_req_num = 16, 1 draft.tp_world_size_, draft.tp_rank_ = tp, rank - prepare_model_states(draft, layer_ids=(config["num_hidden_layers"] + step,)) + prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) drafts.append(draft) model.att_resource_manager.initialize_resources( args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction @@ -200,7 +200,9 @@ def test_cpu_service_describes_main_and_draft_as_distinct_resources( (main,) = resources.get_full_attention_states("main") (draft,) = resources.get_full_attention_states("draft0") assert main is not draft - assert set(main.config.layer_ids).isdisjoint(draft.config.layer_ids) + assert {layer for group in main.config.logical_layer_groups for layer in group}.isdisjoint( + layer for group in draft.config.logical_layer_groups for layer in group + ) assert main.capacity == draft.capacity == resources.full_att_allocator.size assert {main.resource_id, draft.resource_id} <= resources.get_cpu_transfer_layouts().keys() layout_end = max(layout.byte_end for layout in resources.get_cpu_transfer_layouts().values()) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 6e5f7bab37..54c859875c 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -67,7 +67,7 @@ def linear_cache(monkeypatch, tmp_path, request): layer_types=("linear_attention", "full_attention"), ) full = GQAFullAttentionState( - FullAttentionConfig(config.get_full_att_layer_ids(), 1, 8, torch.bfloat16), + FullAttentionConfig(tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), 1, 8, torch.bfloat16), page_size=args.page_size, ) full.initialize_storage(384) From 2b15f7e9548b8b682a0e27ee469eafb9d8eed1c6 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 02:26:42 +0000 Subject: [PATCH 03/56] refactor(att_resources): unify state identities and cached lookup Name resources by model role, attention algorithm and storage scheme. Resolve States by role, type substring and logical layer after creation, with cached queries and unique-match checks. Remove inference-layer name binding and migrate compute, cache and transfer callers. Update design documentation and regression coverage. --- docs/cache_manager_design.md | 78 +++++--- .../att_resources/att_resource_manager.py | 44 +++- .../att_resources/att_tail/base_state.py | 19 +- .../att_resources/att_tail/kpool_state.py | 9 +- .../att_resources/bounded_att/base_state.py | 20 +- .../att_resources/bounded_att/gdn_state.py | 12 +- .../att_resources/bounded_att/kda_state.py | 12 +- .../att_resources/full_att/base_state.py | 23 ++- .../att_resources/full_att/gqa_state.py | 6 +- .../full_att/group_quantized_state.py | 6 +- .../att_resources/full_att/mla_state.py | 4 + .../full_att/packed_mla_state.py | 6 +- .../full_att/static_fp8_state.py | 6 +- lightllm/common/att_resources/types.py | 2 +- .../common/basemodel/attention/fa3/fp8.py | 8 +- .../basemodel/attention/flashinfer/fp8.py | 4 +- .../common/basemodel/attention/linear/gdn.py | 6 +- .../common/basemodel/attention/linear/kda.py | 6 +- lightllm/common/basemodel/basemodel_utils.py | 13 +- .../transformer_layer_infer_template.py | 8 +- .../layer_infer/transformer_layer_infer.py | 15 +- .../layer_infer/transformer_layer_infer.py | 15 +- .../layer_infer/transformer_layer_infer.py | 22 +- .../layer_infer/transformer_layer_infer.py | 8 +- .../models/glm5_next/att_resource_manager.py | 9 +- lightllm/models/glm5_next/indexer.py | 14 +- .../layer_infer/transformer_layer_infer.py | 16 +- .../layer_infer/transformer_layer_infer.py | 15 +- .../layer_infer/transformer_layer_infer.py | 15 +- .../layer_infer/transformer_layer_infer.py | 15 +- .../models/qwen3next/att_resource_manager.py | 7 +- lightllm/models/qwen3next/model.py | 2 +- test/att_resources/validate_model_cache.py | 2 +- test/cpu_cache_kernel/test_speed.py | 2 +- .../common/att_resources/manager_fixtures.py | 12 +- .../common/att_resources/model_helpers.py | 5 +- .../att_resources/test_backend_startup.py | 18 +- .../common/att_resources/test_capacity.py | 8 +- .../att_resources/test_compute_binding.py | 2 +- .../common/att_resources/test_layer_groups.py | 2 + .../test_layer_state_selection.py | 188 ++++++++++++++++++ .../att_resources/test_manager_helpers.py | 34 ++-- .../att_resources/test_model_resources.py | 6 +- .../test_model_storage_definition.py | 2 +- .../att_resources/test_mtp_state_commit.py | 16 +- .../att_resources/test_network_transport.py | 6 +- .../att_resources/test_no_attention_draft.py | 2 +- .../att_resources/test_role_transfers.py | 16 +- .../att_resources/test_runtime_transport.py | 14 +- .../att_resources/test_state_creation.py | 14 +- .../att_resources/test_state_identity.py | 22 +- .../att_resources/test_state_queries.py | 10 +- unit_tests/models/glm5_next/test_cache.py | 6 +- unit_tests/models/glm5_next/test_kernels.py | 7 +- unit_tests/models/glm5_next/test_layers.py | 12 +- unit_tests/models/glm5_next/test_pd_cache.py | 4 +- .../test_hybrid_cpu_state_cache.py | 2 +- .../model_infer/test_multimodal_chunking.py | 4 +- 58 files changed, 668 insertions(+), 193 deletions(-) create mode 100644 unit_tests/common/att_resources/test_layer_state_selection.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 9df001f10a..c8cd6e3137 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -225,19 +225,23 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 `AttResourceManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `set_att_state(model_name, state)` 绑定。`model_name` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.resource_id`。`FullAttentionConfig`（含量化与 MLA 子类）、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `resource_id` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_creation()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 -管理器拥有的 State 成员名与 `state.resource_id` 完全一致，统一使用 `

__state`。draft 按加载顺序从 0 编号，最多 36 个；例如 `draft0_kda_state` 与 `draft1_kda_state` 是两个独立对象。三类 State 基类统一声明抽象接口 `get_state_type() -> str`，各具体方案必须实现，返回算法及存储方案标记，不再保存 `state_type` 属性。GQA、MLA、量化与 K-pool 实现根据 config 返回标记，KDA/GDN 分别返回 `"kda"` / `"gdn"`；定制方案通过重写接口区分。现有标记字符串保持不变，资源绑定后返回值必须稳定，不随容量或设备变化；CPU/DP/PD 布局直接使用相同的 `resource_id` 关联资源，无需转换命名格式。同一管理器内每个 `(role, state_type)` 最多一个实例，启用时必须有且只有一个。同一模型同方案的全部层放入同一个 State，不按层拆成同类型列表，不通过别名重复持有同一 State。 +管理器成员名与 `state.resource_id` 完全一致，统一使用 `{model_name}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState`、`AttentionTailState` 类型表达，名称中不再重复 `full_att`、`bounded_att`、`att_tail`。 + +三类 State 都必须实现 `get_state_type()`，只返回中间的算法及方案标记，使用小写 snake_case。方案标记区分量化位宽、group size、scale 方式/精度，KDA/GDN 分别标注 Conv 和 SSM 精度，K-pool 标注压缩比例与存储精度。层号、容量、设备、TP rank 和请求信息不进入名称。 + +同一角色、同一标记只能注册一个 State，包含该方案的全部逻辑层；不同存储方案可以共存，通过算法字符串、层归属及可选的 State 类别筛选。名字用于资源定位，不能替代完整布局兼容性检查。标记注册后保持稳定，meta 定义、实际存储和 CPU/DP/PD 进程使用相同名称，不提供旧名称别名。 | 方案 | State 成员名及 `resource_id` | | --- | --- | -| 普通 BF16 KV | `main_full_att_bf16_state` | -| INT8，group size 8，BF16 scale | `main_full_att_int8_group8_bf16_scale_state` | -| 静态 FP8，per-head scale | `main_full_att_fp8_static_per_head_state` | +| 普通 BF16 KV | `main_gqa_bf16_state` | +| INT8，group size 8，BF16 scale | `main_gqa_int8_group8_bf16_scale_state` | +| 静态 FP8，per-head scale | `main_gqa_fp8_static_per_head_state` | | MLA BF16 | `draft0_mla_bf16_state` | -| 带 indexer 字段的 packed MLA FP8 | `main_mla_packed_fp8_state` | -| KDA / GDN | `main_kda_state` / `draft0_gdn_state` | -| K-pool C4 indexer tail | `main_indexer_kpool_c4_tail_state` | +| 带 indexer 字段的 packed MLA FP8 | `main_mla_packed_fp8_group128_state` | +| KDA / GDN | `main_kda_conv_bf16_ssm_fp32_state` / `draft0_gdn_conv_bf16_ssm_fp32_state` | +| K-pool C4 indexer tail | `main_kpool_c4_bf16_state` | -普通 KV 不加 `dense`；MLA 直接使用 `mla` 标记，不叠加 `full_att_mla`。普通浮点存储区分 `bf16`、`fp16`、`fp32`；量化方案区分 INT4/INT8、group size 与 scale dtype，静态 FP8 区分 per-head/per-tensor。`packed` 表示 KV 与 indexer 字段联合存储，与纯 MLA 区分。Indexer 使用算法及压缩比例标记，例如 `indexer_kpool_c4_tail`，不以通用 `att_tail` 掩盖具体实现。完整 shape、标定身份等仍保存在 State 配置及存储描述中。 +普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；K-pool 使用 `kpool_cN_*` 标明压缩比例与存储精度。完整 shape、标定身份等仍保存在 State 配置中。 以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： @@ -247,17 +251,17 @@ Glm5NextAttResourceManager ├── bounded_att_big_page_allocator ├── bounded_att_small_page_allocator ├── main_mla_packed_bf16_state -├── main_kda_state -├── main_indexer_kpool_c4_tail_state +├── main_kda_conv_bf16_ssm_fp32_state +├── main_kpool_c4_bf16_state ├── draft0_mla_packed_bf16_state -└── draft0_indexer_kpool_c4_tail_state +└── draft0_kpool_c4_bf16_state ``` -当前 GLM draft 不使用 KDA，因此不创建 `draft0_kda_state`。Qwen 的线性状态为 `main_gdn_state`；Full Attention 按实际配置创建 `main_full_att_bf16_state` 等成员。MLA 与普通 KV 都可使用同一个 `full_att_allocator` 分配域。 +当前 GLM draft 不使用 KDA，因此不创建 `draft0_kda_conv_bf16_ssm_fp32_state`。Qwen 的线性状态为 `main_gdn_conv_bf16_ssm_fp32_state`；Full Attention 按实际配置创建 `main_gqa_bf16_state` 等成员。MLA 与普通 KV 都可使用同一个 `full_att_allocator` 分配域。 `vanilla_no_att`、`eagle_no_att` 等仅执行 FFN 的 draft 不创建 KV State，也不初始化 attention backend；其计算引用可以为 `None`，仍共享管理器的 allocator 与请求映射。缺少实际需要的 draft attention State 则应在初始化时断言失败。 -因此同一角色可以同时拥有 `main_full_att_bf16_state` 和 `main_full_att_int8_group8_bf16_scale_state`；这是两种不同方案，各自覆盖对应的层，聚合复制及传输必须覆盖二者。现有内置模型按配置选择自己的方案；未来混合量化模型在管理器中显式定义多个规范成员，基类通过聚合查询将它们传给相同 helper，不需要增加通用 State 注册表。 +因此同一角色可以同时拥有 `main_gqa_bf16_state` 和 `main_gqa_int8_group8_bf16_scale_state`；这是两种不同方案，各自覆盖对应的层，聚合复制及传输必须覆盖二者。现有内置模型按配置选择自己的方案；未来混合量化模型在管理器中显式定义多个规范成员，基类通过聚合查询将它们传给相同 helper，不需要增加通用 State 注册表。 main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 State 同时命名为 main 和 draft。当前启动流程为各模型分别创建 State 和存储；GLM 的每个模型也分别创建 indexer ring。三种 allocator 的名称不带模型角色，同一分配域只初始化或绑定一个实例。 @@ -340,7 +344,7 @@ CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式 无需模型特有初始化的场景直接使用 `AttResourceManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `set_att_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 -模型特有的 State 创建与存储绑定保留在 `lightllm/models//att_resource_manager.py`。Qwen3Next 创建 `main_gdn_state`；Glm5Next 创建 `main_kda_state` 和 `main_indexer_kpool_c4_tail_state`，压缩比例标记随配置变化。main/draft 的 KV State 按各自方案命名，draft 自行添加需要的 bounded/indexer State。模型管理器继承基类的数据操作和传输布局实现，仅在确有差异时覆盖，不重复实现相同的聚合转发。 +模型特有的 State 创建与存储绑定保留在 `lightllm/models//att_resource_manager.py`。Qwen3Next 创建 `main_gdn_conv_bf16_ssm_fp32_state`；Glm5Next 创建 `main_kda_conv_bf16_ssm_fp32_state` 和 `main_kpool_c4_bf16_state`，压缩比例标记随配置变化。main/draft 的 KV State 按各自方案命名，draft 自行添加需要的 bounded/indexer State。模型管理器继承基类的数据操作和传输布局实现，仅在确有差异时覆盖，不重复实现相同的聚合转发。 token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttResourceManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttResourceManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 resource_id 匹配时可构建临时映射，它不决定模型的资源组织方式。 @@ -448,6 +452,16 @@ State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get 同组逻辑层共享该 State 的全部配套字段，必须使用兼容的布局和参数。分组只表达数据共享关系，不指定写入者或执行顺序；模型负责跳过共享消费者的 KV 写入，避免覆盖生产层数据。普通模型每个逻辑层单独一组；Gemma4 将后部共享层加入同类型的最后一个非共享层所在组。推理层按自己的逻辑层号访问，CPU 服务与模型使用同一套配置构建物理分组。DP/CPU/PD 传输按相同的物理层顺序搬运，main/draft 仍各自拥有独立 State。 +#### State 层归属与模型装配 + +三类 State 均提供 `has_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 + +`AttResourceManager.get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_name` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`att_type` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`has_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `att_type=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 + +该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_states_created=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。`init_model_full_attention()` 只创建和注册 State，`prepare_att_resources()` 不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_att_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 + +推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层 tail 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 + #### 继承与内部 storage 的关系 建议以 `FullAttentionState` 作为公共接口，由具体布局子类实现差异。子类内部的 storage 持有 buffer 和布局信息，operator 可以复用兼容的写入、量化及复制 kernel。布局选择集中在具体子类及其配套实现，公共调用方使用稳定接口。 @@ -494,6 +508,8 @@ class FullAttentionState: # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_layer_view(self, logical_layer_id: int) -> ViewT: ... @@ -559,6 +575,7 @@ class FullAttentionState: | 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | | `initialize_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | | `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | +| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | | `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | | `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | @@ -669,6 +686,7 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 | `set_att_state(model_name, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | | `finish_state_creation()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_states_created = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | | `initialize_resources(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：标记 State 创建完成，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | +| `get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_states_created` 并使用 `lru_cache` 缓存 | | `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | | `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | | `get_attention_tail_states(role=None)` | 返回指定角色或全部 `AttentionTailState` 的 tuple | @@ -683,9 +701,9 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 | `create_dp_ipc_copy()` | DP shared 导出：保留运行态和共享 CPU 大小页池，仅排除 allocator | | `register_page_pools_pin_memory()` | 在当前 GPU 上统一注册全部 bounded State 的大小页池，已注册存储可重复调用 | -聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `set_att_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `resource_id` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数，也不提供单对象选择接口。模型推理直接访问具体 State 成员。 +聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `set_att_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `resource_id` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数。`get_att_state` 提供创建完成后的单对象缓存查询，推理层通过该接口直接取得当前层使用的资源，不另存成员名。 -State 装配期间由管理器内部遍历 `__dict__`，公共查询会断言失败，避免缓存不完整的 main/draft 集合。装配完成后先调用 `finish_state_creation()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 +State 装配期间由管理器内部遍历 `__dict__`，公共聚合查询及 `get_att_state` 均会断言失败，避免缓存不完整的 main/draft 集合；模型准备阶段只创建和注册 State，不查询推理资源。装配完成后先调用 `finish_state_creation()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `initialize_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded/indexer 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 @@ -955,13 +973,13 @@ CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，具体 ### 6.3 模型推理使用具体资源视图 -Model 和 InferStateInfo 通过 `att_resource_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_full_att_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_resource_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 +Model 和 InferStateInfo 通过 `att_resource_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_gqa_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_resource_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与 tail 的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 -普通 KV 层可支持多种存储精度。模型准备 State 时，将该层使用的规范成员名保存为 `layer.full_att_state_name`，关联的 indexer 同样记录该名称；推理通过 `getattr(att_resource_manager, self.full_att_state_name)` 直接访问成员。这里仅保存字符串，不持有 State 别名，不建立资源注册表，也不在 forward 中遍历候选 State。不同层可以记录不同方案的成员名，由具体模型负责确定对应关系。 +普通 KV 层可支持多种存储精度。推理层使用 `get_att_state(model_name=infer_state.att_resource_role, att_type="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState)` 取得资源。MLA 推理层和 GLM indexer 使用 `att_type="mla"`，indexer 使用自己的 `layer_idx`。通用写入模板使用空字符串和 `FullAttentionState` 类型，不推断算法。匹配结果由管理器缓存，推理层和 indexer 不保存重复成员名，也不需要模型启动时逐层绑定。不同精度的 State 由逻辑层归属区分。 -Gemma4 共享 KV 层使用被共享源层所属 State 的成员名，并通过源层号取得视图。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按自身压缩比例访问对应角色的 tail 成员。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 +Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 State 将共享消费者映射到源物理层。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按自身压缩比例访问对应角色的 tail 成员。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 成员访问只读取 Python 属性，不复制 buffer，也不执行 GPU→CPU 同步。KV 写入直接调用该 State 的方法，packed KV/indexer 通过其层视图访问。DFlash 的特殊 KV 写入上下文同样设置管理器和 draft 标记。需要主机标量的 backend 在 State 初始化时准备标定镜像，避免 forward 中逐层执行 GPU→CPU 同步。 @@ -985,7 +1003,7 @@ Gemma4 共享 KV 层使用被共享源层所属 State 的成员名，并通过 | --- | --- | --- | | 本地存储布局与分片事实 | 具体 State 的配置、存储参数及内部 Tensor | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | | 传输字段及本地 TP 切片 | 具体 State 的 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` | 字段选择、全局形状、State 内字段的起止区间和各 rank 的读写切片 | -| 组合页布局及资源间分区 | 模型专属 `AttResourceManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_full_att_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | +| 组合页布局及资源间分区 | 模型专属 `AttResourceManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | | 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `StateTransferLayout`；按资源与 CPU/PD 路径分别保存 | | 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | @@ -1056,7 +1074,7 @@ class TransferRegion: @dataclass(frozen=True) class StateTransferLayout: - resource_id: str # 含模型角色，如 main_full_att_fp16_state + resource_id: str # 含模型角色，如 main_gqa_fp16_state tokens_per_page: int | None # KV 页为 token 容量，纯状态页为 None byte_start: int # 本 State 的页内起点，包含 byte_end: int # 本 State 的页内终点，不包含，含自身对齐与预留空间 @@ -1097,7 +1115,7 @@ k_region = TransferRegion( v_region = replace(k_region, field="v", byte_start=16384, byte_end=32768) layout = StateTransferLayout( - resource_id="main_full_att_fp16_state", + resource_id="main_gqa_fp16_state", tokens_per_page=8, byte_start=0, byte_end=32768, @@ -1179,8 +1197,8 @@ DP 只复制对应 TP rank 的完整分片，不支持通过 State 接口拼接 假设 main 的 full-attention 逻辑层为 `[3, 7]`，本地压缩存储层为 `[0, 1]`；一页 16 个 token，全局 4 个 KV heads，每 head 64 个 FP16 元素。以 K、V 分区的示例格式组织 `kv` 页： ```text -byte 0 : main_full_att_fp16_state K [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB -byte 16384 : main_full_att_fp16_state V [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 0 : main_gqa_fp16_state K [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 16384 : main_gqa_fp16_state V [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB byte 32768 : draft0_mla_fp8_state 的数据区域，按 draft 自身格式定义 ``` @@ -1334,6 +1352,8 @@ class BoundedAttentionState(ABC): # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== @abstractmethod + def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod @@ -1412,6 +1432,7 @@ CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_pag | --- | --- | | 初始化 | 构造函数按 `request_capacity` 和 `mtp_step` 创建运行态及层映射，大小页池先为 `None`；不创建 allocator，不分配请求号或页槽位 | | `initialize_page_pools(...)` | 容量规划完成后创建大小页池，覆盖共享 allocator 的槽位域；容量为零的页池保持 `None`，不重建运行态或重置 MTP 游标 | +| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | | `reset_bounded_att_state(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | | `commit_bounded_att_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | @@ -1485,7 +1506,7 @@ state.commit_bounded_att_state(req_indices, accepted_steps) ### 6.8 AttentionTailState：原始尾部窗口的操作接口 -`AttentionTailState(ABC)` 约束注意力相关模块的尾部运行态操作，具体实现负责自身窗口、位置规则和搬运 kernel。例如 `KPoolAttentionTailState` 定义 GLM K-pool 的尾部布局。通用分类使用 `att_tail`，具体类及 `indexer_kpool_c4_tail` 等资源标记保留算法信息。该分类不包含算子临时 workspace、完整递推状态或随历史增长的索引 key cache；对齐后可重置是 K-pool 的具体规则，不作为所有 tail 实现的共同假设。它不继承 `FullAttentionState` 或 `BoundedAttentionState`，也不持有 allocator、请求 KV 索引表或独立大小页 pool。 +`AttentionTailState(ABC)` 约束注意力相关模块的尾部运行态操作，具体实现负责自身窗口、位置规则和搬运 kernel。例如 `KPoolAttentionTailState` 定义 GLM K-pool 的尾部布局。通用分类使用 `att_tail`，具体类及 `kpool_c4_bf16` 等资源标记保留算法信息。该分类不包含算子临时 workspace、完整递推状态或随历史增长的索引 key cache；对齐后可重置是 K-pool 的具体规则，不作为所有 tail 实现的共同假设。它不继承 `FullAttentionState` 或 `BoundedAttentionState`，也不持有 allocator、请求 KV 索引表或独立大小页 pool。 #### KPoolAttentionTailState 布局与状态边界 @@ -1517,6 +1538,8 @@ class AttentionTailState(ABC): # 二、Attention tail：计算视图与无快照重置 # ======================================================================== @abstractmethod + def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... @abstractmethod @@ -1546,6 +1569,7 @@ class AttentionTailState(ABC): | 接口 | 语义 | | --- | --- | | 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | +| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | | `reset_attention_tail_state(req_indices)` | 清零指定请求的 tail，适用于新请求或可重置的缓存前缀；K-pool 清零完整 ring，其他实现清零对应数据与控制信息。模型配置和缓存流程负责保证边界，接口不接收长度 | @@ -1583,7 +1607,7 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 | 单模型公共接口 | 封装的内部操作 | | --- | --- | -| `prepare_att_resources()` | 在本模型 KV 内存分区内调用 `_init_att_resource_manager`，根据本模型配置准备自身全部 State | +| `prepare_att_resources()` | 在本模型 KV 内存分区内调用 `_init_att_resource_manager` 准备并注册自身全部 State | | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。ReqManager 仍由 main/draft 共用，其构造不依赖 attention 资源已经初始化。 @@ -1595,7 +1619,7 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由单模型公共阶段接口调用，构造函数不调用运行时初始化。main 在资源 hook 中选择具体管理器，各模型分别在自己的 hook 中准备所需 attention State，子类可以继承适用的模型族实现。原生 MTP 如需复用 main 的资源配置，也在自身资源准备阶段读取；ReqManager 初始化只处理请求映射共享。 -需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。分配完成后的 backend 绑定、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方可按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 +需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方可按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_storage_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/att_resources/att_resource_manager.py index d242434ae3..12b274d055 100644 --- a/lightllm/common/att_resources/att_resource_manager.py +++ b/lightllm/common/att_resources/att_resource_manager.py @@ -60,11 +60,17 @@ def __init__(self, request_capacity: int): def set_att_state( self, model_name: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState ) -> None: - """按模型角色和具体方案绑定已创建的 State，保留原对象及其存储。""" + """按模型角色和具体方案注册 State，保留原对象及其存储。 + + 成员名和 resource_id 均为 {model_name}_{state.get_state_type()}_state。 + model_name 为 main 或 draft0..draft35；State 只返回中间的算法/存储方案标记， + 不自行添加角色或 _state。各 State 基类的 get_state_type 注释规定具体命名格式。 + 同角色、同标记不能重复注册；成员名是资源身份，不代表完整布局兼容性。 + """ assert not self._states_created, "State creation is already finished" assert model_name in MODEL_ROLES, f"unknown model role {model_name}" state_type = state.get_state_type() - assert state_type.isidentifier(), f"invalid State type {state_type}" + assert state_type.isidentifier() and state_type.islower(), f"invalid State type {state_type}" name = f"{model_name}_{state_type}_state" assert name not in self.__dict__, f"State {name} is already bound" assert all(state is not value for value in self.__dict__.values()), "a State cannot belong to multiple model roles" @@ -72,6 +78,40 @@ def set_att_state( state.resource_id = name setattr(self, name, state) + @lru_cache(maxsize=None) + def get_att_state( + self, + model_name: str, + att_type: str, + logical_layer_id: int, + *, + state_class: type[FullAttentionState | BoundedAttentionState | AttentionTailState] | None = None, + ) -> FullAttentionState | BoundedAttentionState | AttentionTailState: + """按模型角色、算法/方案字符串和逻辑层选择唯一 State，缓存查询结果。 + + model_name 为 main 或 draft0..draft35，按完整角色前缀匹配，避免 draft1 匹配 draft10。 + att_type 在 get_state_type() 返回的字符串中做包含匹配，如 gqa、gqa_int8、mla_packed。 + 通用推理模板可传空字符串，并通过 state_class 指定资源类别；具体算法传自己的标记。 + 同层存在多类资源时，state_class 可进一步区分；最终必须恰好匹配一个 State。 + 全部模型 State 创建完成后才允许查询；后续原地分配存储不改变缓存对象身份。 + """ + assert self._states_created, "State creation must finish before querying" + assert model_name in MODEL_ROLES, f"unknown model role {model_name}" + matches = [ + state + for name, state in self.__dict__.items() + if name.startswith(model_name + "_") + and isinstance(state, (FullAttentionState, BoundedAttentionState, AttentionTailState)) + and att_type in state.get_state_type() + and (state_class is None or isinstance(state, state_class)) + and state.has_logical_layer(logical_layer_id) + ] + assert len(matches) == 1, ( + f"expected exactly one attention State: {model_name=}, {att_type=}, {logical_layer_id=}, " + f"matches={[state.resource_id for state in matches]}" + ) + return matches[0] + def initialize_resources( self, args: StartArgs | EasyDict, diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/att_resources/att_tail/base_state.py index 33add9cc6f..2127c98c3f 100644 --- a/lightllm/common/att_resources/att_tail/base_state.py +++ b/lightllm/common/att_resources/att_tail/base_state.py @@ -21,13 +21,30 @@ class AttentionTailState(ABC): @abstractmethod def get_state_type(self) -> str: - """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + K-pool 格式为 kpool_c{压缩比例}_{存储精度}，例如 kpool_c4_bf16， + 对应 main_kpool_c4_bf16_state；同压缩比例的不同存储精度使用不同名称。 + 新实现使用具体算法及必要方案参数，不仅返回 indexer 或 tail 这种通用名称。 + """ ... # ======================================================================== # 二、Attention tail：计算视图与无快照重置 # ======================================================================== + @abstractmethod + def has_logical_layer(self, logical_layer_id: int) -> bool: + """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + ... + @abstractmethod def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回尾部状态计算所需的存储视图。""" diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py index 10410bdb76..69c2fcd8f1 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -5,6 +5,7 @@ import torch +from ..utils import dtype_name from ..types import AxisSlice, TransferRegion from .base_state import AttentionTailState from .transfer_utils import att_tail_buffer_copies @@ -77,13 +78,17 @@ def __init__( ) def get_state_type(self) -> str: - """返回包含 K-pool 压缩比例的 tail 方案标记。""" - return f"indexer_kpool_c{self.config.pool_size}_tail" + """返回 K-pool 压缩比例及存储精度。""" + return f"kpool_c{self.config.pool_size}_{dtype_name(dtype=self.config.dtype)}" # ======================================================================== # 二、Attention tail：计算视图与无快照重置 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """返回该层所有请求的原生 ring 视图，末轴包含 key 和 gate。""" return self.tail_buffer[self._logical_to_storage_layer[logical_layer_id]] diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/att_resources/bounded_att/base_state.py index 3756fe2e82..6b47f39e04 100644 --- a/lightllm/common/att_resources/bounded_att/base_state.py +++ b/lightllm/common/att_resources/bounded_att/base_state.py @@ -24,7 +24,20 @@ class BoundedAttentionState(ABC): @abstractmethod def get_state_type(self) -> str: - """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + KDA/GDN 格式为 {kda|gdn}_conv_{Conv精度}_ssm_{SSM精度}。 + 两份状态精度分别标明，不能只写一个 dtype，例如 kda_conv_bf16_ssm_fp32， + 对应 draft0_kda_conv_bf16_ssm_fp32_state。不同精度方案可以在同一角色下共存。 + 其他 bounded 算法使用自己的算法名及方案参数，不复用 kda/gdn 标记。 + """ ... @abstractmethod @@ -43,6 +56,11 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== + @abstractmethod + def has_logical_layer(self, logical_layer_id: int) -> bool: + """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + ... + @abstractmethod def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回运行态和算法控制信息的计算视图。""" diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py index 624d961134..917b106583 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -7,6 +7,7 @@ from lightllm.utils.log_utils import init_logger +from ..utils import dtype_name from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig @@ -60,8 +61,11 @@ def __init__( self.small_pages_pool = None def get_state_type(self) -> str: - """返回 GDN 算法标记。""" - return "gdn" + """返回 GDN 算法及 Conv/SSM 存储精度。""" + return ( + f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) def initialize_page_pools( self, @@ -89,6 +93,10 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" storage_layer_index = self._logical_to_storage_layer[logical_layer_id] diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py index a800eed33e..42c85ba1b7 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -7,6 +7,7 @@ from lightllm.utils.log_utils import init_logger +from ..utils import dtype_name from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig @@ -60,8 +61,11 @@ def __init__( self.small_pages_pool = None def get_state_type(self) -> str: - """返回 KDA 算法标记。""" - return "kda" + """返回 KDA 算法及 Conv/SSM 存储精度。""" + return ( + f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) def initialize_page_pools( self, @@ -89,6 +93,10 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" storage_layer_index = self._logical_to_storage_layer[logical_layer_id] diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index 39d49766e7..c693e53b3a 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -85,7 +85,23 @@ class FullAttentionState(ABC): # 具体存储和量化方案标识，用于 main/draft State 成员命名。 @abstractmethod def get_state_type(self) -> str: - """返回具体算法与存储方案标记，用于生成资源成员名；绑定后须保持稳定。""" + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + - GQA 浮点 KV：gqa_{bf16|fp16|fp32}。 + - GQA 分组量化：gqa_int{4|8}_group{分组大小}_{scale精度}_scale。 + - GQA 静态 FP8：gqa_fp8_static_{per_tensor|per_head}。 + - MLA 浮点：mla_{精度}；KV/indexer 联合存储：mla_packed_{精度}。 + - MLA packed FP8：mla_packed_fp8_group{分组大小}，当前内核固定为 group128。 + 例如 gqa_int8_group128_bf16_scale 对应 main_gqa_int8_group128_bf16_scale_state。 + 新量化实现必须区分影响方案选择的位宽、分组、scale 方式及精度，不编码标定数值。 + """ ... @abstractmethod @@ -113,6 +129,11 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + @abstractmethod + def has_logical_layer(self, logical_layer_id: int) -> bool: + """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + ... + @abstractmethod def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index a87a844982..6611e6b4a7 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -63,7 +63,7 @@ def __init__( def get_state_type(self) -> str: """返回 GQA KV 精度对应的存储方案标记。""" - return "full_att_" + dtype_name(dtype=self.config.dtype) + return "gqa_" + dtype_name(dtype=self.config.dtype) def initialize_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" @@ -84,6 +84,10 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" storage_layer_index = self._logical_to_storage_layer[logical_layer_id] diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index 180618e96e..84186c6af8 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -75,7 +75,7 @@ def __init__( def get_state_type(self) -> str: """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" - return f"full_att_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.dtype)}_scale" + return f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.dtype)}_scale" def initialize_storage(self, token_capacity: int) -> None: """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" @@ -101,6 +101,10 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """返回该层 K/V 编码数据及各自的分组 scale 视图。""" storage_layer_index = self._logical_to_storage_layer[logical_layer_id] diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index 9a8b776008..162dbb378a 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -72,6 +72,10 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" return self.storage[self._logical_to_storage_layer[logical_layer_id]] diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index 8f8142285f..2c61f4fe49 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -93,7 +93,7 @@ def __init__( def get_state_type(self) -> str: """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" - encoding = "fp8" if self.config.fp8_nope_dim else dtype_name(dtype=self.config.dtype) + encoding = f"fp8_group{self.config.quant_group_size}" if self.config.fp8_nope_dim else dtype_name(dtype=self.config.dtype) return "mla_packed_" + encoding def initialize_storage(self, token_capacity: int) -> None: @@ -126,6 +126,10 @@ def indexer_offset(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" row = self.storage[self._logical_to_storage_layer[logical_layer_id]] diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index 6097f433fb..8762d7a6ca 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -86,7 +86,7 @@ def __init__( def get_state_type(self) -> str: """返回包含 per-tensor 或 per-head 标定方式的 FP8 方案标记。""" - return "full_att_fp8_static_" + self.config.scale_mode + return "gqa_fp8_static_" + self.config.scale_mode def initialize_storage(self, token_capacity: int) -> None: """绑定 KV 与真实标定参数，并预备 FlashInfer 使用的 CPU scale。""" @@ -114,6 +114,10 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== + def has_logical_layer(self, logical_layer_id: int) -> bool: + """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" storage_layer_index = self._logical_to_storage_layer[logical_layer_id] diff --git a/lightllm/common/att_resources/types.py b/lightllm/common/att_resources/types.py index 3972edd4f5..c9cb2c49ac 100644 --- a/lightllm/common/att_resources/types.py +++ b/lightllm/common/att_resources/types.py @@ -182,7 +182,7 @@ class StateTransferLayout: ... ), ... ) >>> layout = StateTransferLayout( - ... resource_id="main_full_att_fp16_state", + ... resource_id="main_gqa_fp16_state", ... tokens_per_page=8, ... byte_start=16384, byte_end=49152, ... regions=(k_region, replace(k_region, field="v", byte_start=32768, byte_end=49152)), diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index 0629e2215e..d88b258d05 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -41,7 +41,7 @@ def init_state(self): batch_size = self.infer_state.batch_size state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) offline_scales: torch.Tensor = state.scales @@ -87,7 +87,7 @@ def _fp8_prefill_att( cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) static_q_scales = state.q_scales[storage_layer_index] @@ -123,7 +123,7 @@ def init_state(self): state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) super().init_state() @@ -179,7 +179,7 @@ def _fp8_decode_att( state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_head_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 022eb16cb6..174387aa33 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -68,7 +68,7 @@ def _fp8_prefill_att( v = v.unsqueeze(1).view(torch.float8_e4m3fn) state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", ) storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) k_descale = state.cpu_scales[storage_layer_index][0] @@ -124,7 +124,7 @@ def _fp8_decode_att( v = v.unsqueeze(1).view(torch.float8_e4m3fn) state = getattr( self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_full_att_fp8_static_per_tensor_state", + f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", ) storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index fde9c74c7e..9b8f81ef27 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.bounded_att import GDNState import dataclasses import weakref import torch @@ -29,7 +30,10 @@ class LinearAttBackend(BaseAttBackend, ABC): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) role = model.att_resource_role - self.linear_state = getattr(model.att_resource_manager, f"{role}_gdn_state") + (self.linear_state,) = ( + state for state in model.att_resource_manager.get_bounded_attention_states(role) + if isinstance(state, GDNState) + ) self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) self.prefill_kernel = self.get_prefill_kernel() diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 3807107666..fc4a59fdee 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,6 +10,7 @@ import torch +from lightllm.common.att_resources.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, @@ -38,7 +39,10 @@ class KDALinearAttBackend(BaseAttBackend): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) role = model.att_resource_role - self.linear_state = getattr(model.att_resource_manager, f"{role}_kda_state") + (self.linear_state,) = ( + state for state in model.att_resource_manager.get_bounded_attention_states(role) + if isinstance(state, KDAState) + ) config = model.config["linear_attn_config"] self.num_heads = config["num_heads"] self.head_dim = config["head_dim"] diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index b063aaef00..d64692b16a 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -137,15 +137,6 @@ def init_model_full_attention(model: TpPartBaseModel) -> None: if definition is None: return model.att_resource_manager.set_att_state(model.att_resource_role, definition) - # 层和 indexer 只记录规范成员名，推理时直接访问管理器中已分配的 State。 - name = definition.resource_id - # 绑定所有逻辑层，包括共享消费者；State 自己将逻辑层号映射到物理层。 - logical_layer_ids = {logical_layer_id for logical_layer_group in definition.config.logical_layer_groups for logical_layer_id in logical_layer_group} - for layer in model.layers_infer: - if layer.layer_num_ in logical_layer_ids: - layer.full_att_state_name = name - if hasattr(layer, "indexer"): - layer.indexer.full_att_state_name = name def init_model_bounded_attention(model: TpPartBaseModel) -> None: @@ -157,9 +148,9 @@ def init_model_bounded_attention(model: TpPartBaseModel) -> None: ) if definition is None: return - from lightllm.common.att_resources.bounded_att import create_gdn_state, create_kda_state + from lightllm.common.att_resources.bounded_att import GDNState, KDAState, create_gdn_state, create_kda_state - create_state = {"gdn": create_gdn_state, "kda": create_kda_state}[definition.get_state_type()] + create_state = {GDNState: create_gdn_state, KDAState: create_kda_state}[type(definition)] state = create_state( config=definition.config, request_capacity=model.max_req_num + 1, mtp_step=model.args.mtp_step, device="cuda" ) diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index 9a2183326b..dedc153180 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import os import torch import torch.distributed as dist @@ -33,7 +34,12 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[to raise Exception("need to impl") def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight): - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) state.write_full_att_kv( logical_layer_id=self.layer_num_, token_indices=infer_state.mem_index, diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index a018399bdb..5709657ffd 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import torch from typing import Tuple from lightllm.common.basemodel import TransformerLayerInferTpl @@ -28,7 +29,12 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( @@ -44,7 +50,12 @@ def _context_attention_kernel( def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index a5a061cbd2..61411d4eaf 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import os import torch import triton @@ -103,7 +104,12 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) kv = state.get_layer_view(logical_layer_id=self.layer_num_) out = infer_state.decode_att_state.decode_att( @@ -120,7 +126,12 @@ def _decompress_kv( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) compressed_kv = state.get_layer_view(logical_layer_id=self.layer_num_) total_token_num = infer_state.total_token_num diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index 43f9998a52..1122d6badc 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo @@ -99,7 +100,12 @@ def _context_attention_kernel( }, ) - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) mla_out = infer_state.prefill_att_state.prefill_att( q=q_all, k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, @@ -142,7 +148,12 @@ def _token_attention_kernel( }, ) - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) o_tensor = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, @@ -198,7 +209,12 @@ def _get_indices( q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) destindex_copy_indexer_ks( K_fp8=k_fp8, K_scale=k_scale, diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2010a674ee..bd7720c81e 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import math import torch import torch.nn as nn @@ -183,7 +184,12 @@ def _att_control(self): def _get_layer_kv(self, infer_state: InferStateInfo): # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k_raw, _v_raw = state.get_layer_view(logical_layer_id=self.layer_num_) # _k_raw / _v_raw shape (S, cache_slot_num, cache_slot_dim). Use .view # (not .reshape) so any non-contiguous layout from a future mem_manager diff --git a/lightllm/models/glm5_next/att_resource_manager.py b/lightllm/models/glm5_next/att_resource_manager.py index 5c43191b44..e9333bf7b0 100644 --- a/lightllm/models/glm5_next/att_resource_manager.py +++ b/lightllm/models/glm5_next/att_resource_manager.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.utils import dtype_name from lightllm.common.att_resources.allocators import StatePageAllocator from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig @@ -83,7 +84,7 @@ def __init__( if logical_layer_ids: state = create_kda_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) self.set_att_state(role, state) - self.att_tail_type = f"indexer_kpool_c{cache_config.index_kpool}_tail" + self.att_tail_type = f"kpool_c{cache_config.index_kpool}_{dtype_name(dtype=cache_config.full_att_dtype)}" own_full_att_logical_layer_ids = ( tuple( range( @@ -147,7 +148,11 @@ def bind_storage( for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "draft State must only contain draft layers" - state = self.main_kda_state + state = getattr( + self, + f"main_kda_conv_{dtype_name(dtype=self.cache_config.conv_state_dtype)}" + f"_ssm_{dtype_name(dtype=self.cache_config.ssm_state_dtype)}_state", + ) state.initialize_page_pools( big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=state.conv_state.is_cuda, diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index a0afccb1c0..e06dbc2fe6 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import torch import triton @@ -38,10 +39,17 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) raw = torch.cat((k, gate), -1) role = infer_state.att_resource_role - tail = getattr( - infer_state.att_resource_manager, f"{role}_indexer_kpool_c{self.pool_size}_tail_state" + tail = infer_state.att_resource_manager.get_att_state( + model_name=role, + att_type=f"kpool_c{self.pool_size}", + logical_layer_id=self.layer_idx, ).get_layer_view(logical_layer_id=self.layer_idx) - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_idx, + state_class=FullAttentionState, + ) packed_buffer = state.get_layer_view(logical_layer_id=self.layer_idx).indexer_k compress_pools( raw=raw, diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 0a02b295e6..4374e4f51b 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -2,6 +2,8 @@ from __future__ import annotations +from lightllm.common.att_resources.full_att import FullAttentionState + import torch from lightllm.common.basemodel import TransformerLayerInferTpl @@ -159,7 +161,12 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): layer_weight=layer_weight, ) del infer_state.get_topk_indices_params - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) return infer_state.prefill_att_state.prefill_att( q=q, k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, @@ -241,7 +248,12 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) return infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index f698747191..4fe03c943d 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import torch from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo @@ -72,7 +73,12 @@ def _context_attention_kernel( window_size = (-1, -1) use_sliding_window = False - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( @@ -100,7 +106,12 @@ def _token_attention_kernel( window_size = (-1, -1) use_sliding_window = False - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 36ea15e0b6..19a6a8c688 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState import torch import triton import torch.distributed as dist @@ -44,7 +45,12 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( @@ -62,7 +68,12 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor) diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index 20afacf6f8..e1873b308c 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.att_resources.full_att import FullAttentionState from lightllm.common.basemodel.attention import AttControl from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv @@ -13,7 +14,12 @@ def __init__(self, layer_num, network_config): self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) def _token_attention_kernel(self, q, infer_state, layer_weight): - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) k, v = state.get_layer_view(logical_layer_id=self.layer_num_) k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) @@ -35,7 +41,12 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): - state = getattr(infer_state.att_resource_manager, self.full_att_state_name) + state = infer_state.att_resource_manager.get_att_state( + model_name=infer_state.att_resource_role, + att_type="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) storage_head_num = state.config.head_num storage_head_dim = state.config.head_dim storage_kv = cache_kv.view(cache_kv.shape[0], 2 * storage_head_num, storage_head_dim) diff --git a/lightllm/models/qwen3next/att_resource_manager.py b/lightllm/models/qwen3next/att_resource_manager.py index 63438488ab..48e07c7c29 100644 --- a/lightllm/models/qwen3next/att_resource_manager.py +++ b/lightllm/models/qwen3next/att_resource_manager.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.utils import dtype_name from lightllm.common.att_resources.allocators import StatePageAllocator from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig @@ -78,7 +79,11 @@ def bind_storage( for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group ), "draft State must only contain draft layers" - state = self.main_gdn_state + state = getattr( + self, + f"main_gdn_conv_{dtype_name(dtype=self.cache_config.conv_state_dtype)}" + f"_ssm_{dtype_name(dtype=self.cache_config.ssm_state_dtype)}_state", + ) state.initialize_page_pools( big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=state.conv_state.is_cuda, diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 13263d56f4..0d0c247960 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -107,7 +107,7 @@ def _full_attention_options(cls, config, args): def _init_att_backend1(self): role = self.att_resource_role - if getattr(self.att_resource_manager, f"{role}_gdn_state", None) is None: + if not self.att_resource_manager.get_bounded_attention_states(role): self.prefill_att_backend1 = None self.decode_att_backend1 = None return diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 8a55f86145..da9a805c41 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -205,7 +205,7 @@ def decode(request: torch.Tensor): logits.append(output.logits.clone()) if draft_models and hybrid: ( - resources.main_kda_state if hasattr(resources, "main_kda_state") else resources.main_gdn_state + resources.main_kda_conv_bf16_ssm_fp32_state if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") else resources.main_gdn_conv_bf16_ssm_fp32_state ).commit_bounded_att_state(request, [accepted_width]) return logits diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index e7fc6f238b..defa3e59ca 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -105,7 +105,7 @@ def restore(): + full.v_storage[:, :tokens].numel() * full.v_storage.element_size() ) pool = ( - resources.main_kda_state if hasattr(resources, "main_kda_state") else resources.main_gdn_state + resources.main_kda_conv_bf16_ssm_fp32_state if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") else resources.main_gdn_conv_bf16_ssm_fp32_state ).big_pages_pool state_bytes = args.pages * sum(t[0].numel() * t.element_size() for t in (pool.conv_state, pool.ssm_state)) median = statistics.median(times) diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index f41e1d9924..ed78f7fedd 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -23,19 +23,19 @@ def __init__(self, device: str | torch.device | int = "cpu", *, with_tail: bool super().__init__(request_capacity=3) self.set_att_state("main", make_state(device=device)) draft_state = GDNState( - replace(self.main_kda_state.config, global_v_heads=2), - request_capacity=3, mtp_step=self.main_kda_state.mtp_step, device=device, + replace(self.main_kda_conv_fp32_ssm_fp32_state.config, global_v_heads=2), + request_capacity=3, mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, device=device, ) draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") self.set_att_state("draft0", draft_state) self.set_att_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( - self.main_full_att_fp32_state.config, - page_size=self.main_full_att_fp32_state.page_size, device=device, + self.main_gqa_fp32_state.config, + page_size=self.main_gqa_fp32_state.page_size, device=device, ) - prepared_full_state.initialize_storage(self.main_full_att_fp32_state.capacity) + prepared_full_state.initialize_storage(self.main_gqa_fp32_state.capacity) self.set_att_state("draft0", prepared_full_state) - self.att_tail_type = f"indexer_kpool_c{tail_pool_size}_tail" + self.att_tail_type = f"kpool_c{tail_pool_size}_fp32" if with_tail: tail = make_tail(device=device) state = KPoolAttentionTailState( diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 49b27d72e8..1def48fb89 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -62,9 +62,8 @@ def owned_states(attributes: Mapping[str, object], state_class: type | tuple[typ def main_linear_state(resources: AttResourceManager): - if hasattr(resources, "main_kda_state"): - return resources.main_kda_state - return resources.main_gdn_state + (state,) = resources.get_bounded_attention_states("main") + return state def prepare_model_states(model: TpPartBaseModel, *, logical_layer_ids: Sequence[int] | None = None): diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index 98122c4b9d..e5e6a34208 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -127,6 +127,18 @@ def test_backend_allocates_all_roles_before_warmup( mapping = SimpleNamespace(bind_hold_page=lambda indices: None) def mark_backends(model: TpPartBaseModel): + # 第一个模型初始化 backend 前，全部 State 已创建完成，可通过缓存查询取得资源。 + for item in models: + for state in item.att_resource_manager.get_full_attention_states(item.att_resource_role): + logical_layer_ids = {layer_id for group in state.config.logical_layer_groups for layer_id in group} + for layer in item.layers_infer: + if layer.layer_num_ in logical_layer_ids: + assert item.att_resource_manager.get_att_state( + model_name=item.att_resource_role, + att_type="", + logical_layer_id=layer.layer_num_, + state_class=type(state), + ) is state model.prefill_att_backend = model.decode_att_backend = object() model.prefill_att_backend1 = model.decode_att_backend1 = None model.backends_bound = True @@ -181,8 +193,8 @@ def memory_info(device: str | torch.device | int): manager = models[0].att_resource_manager states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) - assert manager.draft0_gdn_state.conv_state.is_cuda - assert manager.draft0_gdn_state.big_pages_pool is None + assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.conv_state.is_cuda + assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.big_pages_pool is None assert manager.full_att_allocator is None assert not warmed return manager.get_full_att_token_bytes() * (32 + args.page_size), 0 @@ -248,7 +260,7 @@ def start_runtime(): ) assert all(not state.mtp_state_index.any() for state in draft_states) if main_linear: - assert manager.main_gdn_state.mtp_state_index[0] == draft_count + assert manager.main_gdn_conv_bf16_ssm_fp32_state.mtp_state_index[0] == draft_count for i, state in enumerate(draft_states): assert len(state.config.logical_layer_ids) == 1 manager.commit_bounded_att_state( diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index 4515812be5..86f1f20534 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -401,7 +401,7 @@ def _get_transfer_regions(self, page_tokens: int): manager.initialize_storage(8, 4, big_page_tokens=8, device="cpu") assert get_page_bytes() == (80 if transfer == "cpu" else 72) layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() - assert layouts["main_full_att_bf16_state"].byte_end == 72 + assert layouts["main_gqa_bf16_state"].byte_end == 72 @pytest.mark.parametrize("with_tail", [False, True]) @@ -435,7 +435,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( manager.finish_state_creation() assert manager.get_pd_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) assert manager.full_att_allocator is None - assert read_kv_storage(manager.main_full_att_fp32_state) is None + assert read_kv_storage(manager.main_gqa_fp32_state) is None assert not hasattr(manager, "pd_transfer_layouts") @@ -613,7 +613,7 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation getter = getattr(manager, f"get_{transfer}_transfer_layouts") before = getter() assert getter() is before - state = manager.main_gdn_state if transfer == "pd_bounded_att_or_tail" else manager.main_full_att_bf16_state + state = manager.main_gdn_conv_bf16_ssm_bf16_state if transfer == "pd_bounded_att_or_tail" else manager.main_gqa_bf16_state before_layout = before[state.resource_id] monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) @@ -623,7 +623,7 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation after_layout = after[state.resource_id] assert after is not before and getter() is after assert after_layout is not before_layout - full_state = manager.main_full_att_bf16_state + full_state = manager.main_gqa_bf16_state assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() if transfer == "pd_bounded_att_or_tail": diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index 0001baac9d..9817e23265 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -47,7 +47,7 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, heads = 1 if kind == "packed_glm" else 2 * state.config.head_num dim = state.config.kv_dim if kind == "packed_glm" else state.config.head_dim kv = torch.full((2, heads, dim), 0.5, dtype=torch.bfloat16, device="cuda") - layer = SimpleNamespace(layer_num_=8, full_att_state_name=state.resource_id) + layer = SimpleNamespace(layer_num_=8) TransformerLayerInferTpl._post_cache_kv(layer, kv, infer, None) if capture: graph = torch.cuda.CUDAGraph() diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 91efbee2c5..9f8acdaf4f 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -31,6 +31,8 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - template = definition(kind) config = replace(template.config, logical_layer_groups=((0, 4), (8, 11))) state = rebuild_full_state(template, config=config, token_capacity=8, device="cpu") + assert all(state.has_logical_layer(logical_layer_id) for logical_layer_id in (0, 4, 8, 11)) + assert not state.has_logical_layer(1) assert all(buffer.shape[0] == 2 for buffer in full_state_test_buffers(state).values()) assert state.get_token_bytes() == template.get_token_bytes() assert allocated_bytes(state, token_only=True) == state.get_token_bytes() * (8 + state.page_size) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py new file mode 100644 index 0000000000..eec81f7195 --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -0,0 +1,188 @@ +"""按角色、层归属及方案查询资源，不在推理层保存绑定状态。""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.full_att import FullAttentionConfig, FullAttentionState, GQAFullAttentionState +from lightllm.common.att_resources.bounded_att import KDAState, GDNState +from unit_tests.common.att_resources.test_bounded_state import make_state +from unit_tests.common.att_resources.test_tail_state import make_tail +from unit_tests.common.att_resources.test_capacity import definition +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.basemodel_utils import init_model_full_attention +from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer + + +def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args, model.config = SimpleNamespace(), {} + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.layers_infer = [SimpleNamespace(layer_num_=0, indexer=SimpleNamespace())] + model.att_resource_manager = AttResourceManager(request_capacity=1) + state = GQAFullAttentionState( + layout_config=FullAttentionConfig( + logical_layer_groups=((0,),), global_head_num=1, head_dim=128, dtype=torch.bfloat16, + ), device="meta", + ) + monkeypatch.setattr( + "lightllm.common.basemodel.basemodel_utils.create_full_attention_definition", lambda *args, **kwargs: state, + ) + + init_model_full_attention(model) + + assert model.att_resource_manager.main_gqa_bf16_state is state + assert not hasattr(model.layers_infer[0], "full_att_state_name") + assert not hasattr(model.layers_infer[0].indexer, "full_att_state_name") + + +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str) -> None: + manager = AttResourceManager(request_capacity=1) + for model_role in ("main", "draft0"): + for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): + manager.set_att_state(model_role, GQAFullAttentionState( + layout_config=FullAttentionConfig( + logical_layer_groups=groups, global_head_num=1, head_dim=128, dtype=dtype, + ), device="meta", + )) + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.is_mtp_draft_model = role != "main" + model.mtp_previous_draft_models = [] + model.max_req_num = 0 + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + model.main_model = SimpleNamespace(att_resource_manager=manager) + model._init_att_resource_manager = lambda: setattr(model, "att_resource_manager", manager) + model.layers_infer = [] + for logical_layer_id in range(4): + layer = TransformerLayerInfer.__new__(TransformerLayerInfer) + layer.layer_num_ = logical_layer_id + model.layers_infer.append(layer) + + model.prepare_att_resources() + + assert all(not hasattr(layer, "full_att_state_name") for layer in model.layers_infer) + manager.finish_state_creation() + state = manager.get_att_state(role, "gqa", 0) + assert state.resource_id == f"{role}_gqa_bf16_state" + assert manager.get_att_state(role, "gqa", 2) is state + assert manager.get_att_state(role, "gqa", 3).resource_id == f"{role}_gqa_fp16_state" + + +@pytest.mark.parametrize("finalized", [False, True]) +def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> None: + from dataclasses import replace + + manager = AttResourceManager(request_capacity=1) + plain = definition("gqa") + quantized = definition("int8") + quantized = type(quantized)( + layout_config=replace(quantized.config, logical_layer_groups=((10, 12),)), device="meta", + ) + manager.set_att_state("main", plain) + manager.set_att_state("main", quantized) + manager.set_att_state("draft0", definition("gqa")) + logical_layer_id = plain.config.logical_layer_groups[0][0] + if not finalized: + with pytest.raises(AssertionError, match="State creation must finish"): + manager.get_att_state("main", "gqa", logical_layer_id) + manager.finish_state_creation() + + assert manager.get_att_state("main", "gqa", logical_layer_id) is plain + assert manager.get_att_state("main", "gqa", 10) is quantized + assert manager.get_att_state("main", "int8", 10) is quantized + assert manager.get_att_state("main", "gqa", 12) is quantized + assert manager.get_att_state("draft0", "gqa", logical_layer_id) is not plain + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_att_state("main", "gqa", 99) + + +def test_selection_requires_finished_creation_and_reuses_cache() -> None: + manager = AttResourceManager(request_capacity=1) + plain = definition("gqa") + manager.set_att_state("main", plain) + logical_layer_id = plain.config.logical_layer_groups[0][0] + with pytest.raises(AssertionError, match="State creation must finish"): + manager.get_att_state("main", "gqa", logical_layer_id) + manager.finish_state_creation() + assert manager.get_att_state("main", "gqa", logical_layer_id) is plain + hits = manager.get_att_state.cache_info().hits + assert manager.get_att_state("main", "gqa", logical_layer_id) is plain + assert manager.get_att_state.cache_info().hits == hits + 1 + + +def test_selection_rejects_ambiguous_finished_members() -> None: + manager = AttResourceManager(request_capacity=1) + plain = definition("gqa") + manager.set_att_state("main", plain) + manager.set_att_state("main", definition("int8")) + manager.finish_state_creation() + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_att_state("main", "gqa", plain.config.logical_layer_groups[0][0]) + + +@pytest.mark.parametrize("state_class", [KDAState, GDNState]) +def test_bounded_state_layer_ownership(state_class: type[KDAState] | type[GDNState]) -> None: + state = make_state(cls=state_class) + manager = AttResourceManager(request_capacity=3) + manager.set_att_state("main", state) + manager.finish_state_creation() + assert state.has_logical_layer(1) and state.has_logical_layer(3) + assert not state.has_logical_layer(0) + assert manager.get_att_state("main", state.get_state_type(), 3) is state + + +def test_mla_selection_excludes_same_layer_tail() -> None: + from dataclasses import replace + + manager = AttResourceManager(request_capacity=3) + tail = make_tail() + template = definition("mla") + state = type(template)( + layout_config=replace(template.config, logical_layer_groups=((3, 7),)), device="meta", + ) + manager.set_att_state("main", state) + manager.set_att_state("main", tail) + manager.finish_state_creation() + assert tail.has_logical_layer(3) and tail.has_logical_layer(7) + assert not tail.has_logical_layer(0) + assert manager.get_att_state("main", "mla", 7) is state + assert manager.get_att_state("main", "kpool", 7) is tail + assert manager.get_att_state("main", "", 7, state_class=FullAttentionState) is state + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_att_state("main", "", 7) + + + +@pytest.mark.parametrize("role", ["draft1", "draft10"]) +def test_query_matches_complete_model_role(role: str) -> None: + manager = AttResourceManager(request_capacity=1) + for model_name in ("draft1", "draft10"): + manager.set_att_state(model_name, definition("gqa")) + manager.finish_state_creation() + assert manager.get_att_state(role, "gqa", 0).resource_id == f"{role}_gqa_bf16_state" + + +@pytest.mark.parametrize("algorithm", ["kda", "gdn", "kpool"]) +def test_precision_variants_have_distinct_names_and_can_coexist(algorithm: str) -> None: + from dataclasses import replace + + manager = AttResourceManager(request_capacity=3) + template = make_tail() if algorithm == "kpool" else make_state(cls=KDAState if algorithm == "kda" else GDNState) + states = [] + for dtype in (torch.bfloat16, torch.float32): + config = replace(template.config, **({"dtype": dtype} if algorithm == "kpool" else {"conv_dtype": dtype})) + state = type(template)(config=config, request_capacity=3, device="cpu") + manager.set_att_state("main", state) + states.append(state) + manager.finish_state_creation() + logical_layer_id = template.config.logical_layer_ids[0] + for state in states: + assert manager.get_att_state("main", state.get_state_type(), logical_layer_id) is state + assert getattr(manager, state.resource_id) is state + assert states[0].resource_id != states[1].resource_id + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_att_state("main", algorithm, logical_layer_id) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 84480681a8..e43f32b5e1 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -29,9 +29,9 @@ def composition(device: str | torch.device | int = "cpu", *, with_tail: bool = T def layouts_for(manager: AttResourceManager, *, kind: str = "cpu"): layouts = [] if kind != "att_state": - layouts += [layout_for(state) for state in (manager.main_full_att_fp32_state, manager.draft0_mla_fp32_state)] + layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] if kind != "kv": - layouts += [page_layout(state) for state in (manager.main_kda_state, manager.draft0_gdn_state)] + layouts += [page_layout(state) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)] if kind == "att_state": layouts += [tail_layout(state) for state in manager.get_attention_tail_states("main")] offset = 0 @@ -52,12 +52,12 @@ def layouts_for(manager: AttResourceManager, *, kind: str = "cpu"): def init_source(manager: AttResourceManager): - for i, state in enumerate((manager.main_full_att_fp32_state, manager.draft0_mla_fp32_state)): + for i, state in enumerate((manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.copy_( torch.arange(storage_buffer.numel(), device=storage_buffer.device).reshape_as(storage_buffer) + i ) - for i, state in enumerate((manager.main_kda_state, manager.draft0_gdn_state)): + for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): state.conv_state[:, 0].fill_(11 + i) state.conv_state[:, 1].fill_(23 + i) state.ssm_state[:, :3].fill_(31 + i) @@ -67,7 +67,7 @@ def init_source(manager: AttResourceManager): state.tail_buffer[:, 1].fill_(9) slots = manager.bounded_att_big_page_allocator.alloc(2) manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) - if manager.main_kda_state.conv_state.is_cuda: + if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: torch.cuda.synchronize() return slots @@ -76,14 +76,14 @@ def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.Mon manager = composition() slots = init_source(manager) manager.restore_bounded_att_pages([2], page_kind="big", slot_ids=[slots[0]]) - for i, state in enumerate((manager.main_kda_state, manager.draft0_gdn_state)): + for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): assert (state.conv_state[:, 2, :, :2] == 11 + i).all() assert (state.ssm_state[:, 6] == 31 + i).all() def fail(*args: object, **kwargs: object): raise RuntimeError("simulated copy failure") - monkeypatch.setattr(manager.draft0_gdn_state, "save_bounded_att_pages", fail) + monkeypatch.setattr(manager.draft0_gdn_conv_fp32_ssm_fp32_state, "save_bounded_att_pages", fail) with pytest.raises(RuntimeError): manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) assert manager.bounded_att_big_page_allocator.free_count == 1 @@ -125,13 +125,13 @@ def test_reset_tail_clears_all_tail_states_and_preserves_linear(): init_source(manager) # Preserve accepted-state cursors as well as both algorithms' runtime and # cached pages. A full reset here would silently corrupt a prefix restore. - for state in (manager.main_kda_state, manager.draft0_gdn_state): + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state): state.mtp_state_index[1] = 2 tensors = [ - *full_state_test_buffers(manager.main_full_att_fp32_state).values(), + *full_state_test_buffers(manager.main_gqa_fp32_state).values(), *full_state_test_buffers(manager.draft0_mla_fp32_state).values(), ] - for state in (manager.main_kda_state, manager.draft0_gdn_state): + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state): tensors.extend(( state.conv_state, state.ssm_state, state.mtp_state_index, state.big_pages_pool.conv_state, state.big_pages_pool.ssm_state, @@ -176,19 +176,19 @@ def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | tor req_idx=2, ) for src, dst in zip( - (source.main_kda_state, source.draft0_gdn_state), (target.main_kda_state, target.draft0_gdn_state) + (source.main_kda_conv_fp32_ssm_fp32_state, source.draft0_gdn_conv_fp32_ssm_fp32_state), (target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state) ): for src_slot, dst_slot in zip(slots, target_slots): assert torch.equal(src.big_pages_pool.conv_state[src_slot], dst.big_pages_pool.conv_state[dst_slot]) assert torch.equal(src.big_pages_pool.ssm_state[src_slot], dst.big_pages_pool.ssm_state[dst_slot]) assert (tail.tail_buffer[:, :2] == 37).all() assert not tail.tail_buffer[:, 2].any() - for i, state in enumerate((target.main_kda_state, target.draft0_gdn_state)): + for i, state in enumerate((target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state)): assert (state.conv_state[:, 2, :, :2] == 23 + i).all() assert (state.ssm_state[:, 6] == 47 + i).all() for src, dst in zip( - (source.main_full_att_fp32_state, source.draft0_mla_fp32_state), - (target.main_full_att_fp32_state, target.draft0_mla_fp32_state), + (source.main_gqa_fp32_state, source.draft0_mla_fp32_state), + (target.main_gqa_fp32_state, target.draft0_mla_fp32_state), ): assert torch.equal(read_kv_storage(src)[:, :8], read_kv_storage(dst)[:, 4:12]) @@ -298,12 +298,12 @@ def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.d src_req_idx=1, dst_req_idx=2, ) for a, b in zip( - (target.main_full_att_fp32_state, target.draft0_mla_fp32_state), - (dp_target.main_full_att_fp32_state, dp_target.draft0_mla_fp32_state), + (target.main_gqa_fp32_state, target.draft0_mla_fp32_state), + (dp_target.main_gqa_fp32_state, dp_target.draft0_mla_fp32_state), ): assert torch.equal(read_kv_storage(a)[:, [5, 7]], read_kv_storage(b)[:, [5, 7]]) for a, b in zip( - (target.main_kda_state, target.draft0_gdn_state), (dp_target.main_kda_state, dp_target.draft0_gdn_state) + (target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state), (dp_target.main_kda_conv_fp32_ssm_fp32_state, dp_target.draft0_gdn_conv_fp32_ssm_fp32_state) ): assert torch.equal(a.conv_state[:, 2], b.conv_state[:, 2]) assert torch.equal(a.ssm_state[:, 6], b.ssm_state[:, 6]) diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index e7390ce2f5..4c86507c08 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -128,7 +128,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: i torch.cuda.synchronize() for rank, (resources, full) in enumerate(targets): if glm: - resources.main_indexer_kpool_c4_tail_state.tail_buffer.fill_(13) + resources.main_kpool_c4_bf16_state.tail_buffer.fill_(13) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] resources.read_req_from_cpu( buffer, @@ -148,8 +148,8 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: i dst_fields[name].tensor[:, 8:16].contiguous().view(torch.uint8), ) if glm: - assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].any() - assert (resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, [0, 2]] == 13).all() + assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, 1].any() + assert (resources.main_kpool_c4_bf16_state.tail_buffer[:, [0, 2]] == 13).all() assert not buffer[0].any() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index f3b6fd7160..e8771c686d 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -125,7 +125,7 @@ def test_model_and_cpu_service_use_identical_storage_geometry( assert state.k_storage.shape == state.v_storage.shape == (storage_layers, 20, local_heads, dim) if storage_layers == 1: assert state.config.logical_layer_groups == ((0, 1),) - assert model.layers_infer[0].full_att_state_name == model.layers_infer[1].full_att_state_name + assert resources.get_att_state("main", "gqa", 0) is resources.get_att_state("main", "gqa", 1) assert state.get_layer_view(logical_layer_id=0)[0].data_ptr() == state.get_layer_view(logical_layer_id=1)[0].data_ptr() assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) assert resources.full_att_allocator.size == 16 diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index f72d7dc3cd..dce30194c1 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -38,19 +38,19 @@ def test_commit_selects_state_for_one_role(): torch.tensor([3], dtype=torch.int32, device="cuda"), role="main", ) - for state in (manager.main_kda_state, manager.main_gdn_state): + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.main_gdn_conv_fp32_ssm_fp32_state): assert state.mtp_state_index.tolist() == [2, 1, 1] selected = state.mtp_state_index[0].item() assert (state.ssm_state[:, selected] == 12).all() - assert manager.draft0_gdn_state.mtp_state_index.tolist() == [1, 1, 1] + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [1, 1, 1] manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([1], dtype=torch.int32, device="cuda"), role="draft0", ) - assert manager.draft0_gdn_state.mtp_state_index.tolist() == [0, 1, 1] - assert manager.main_kda_state.mtp_state_index.tolist() == [2, 1, 1] - assert manager.main_gdn_state.mtp_state_index.tolist() == [2, 1, 1] + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 1] + assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [2, 1, 1] + assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [2, 1, 1] @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -104,6 +104,6 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i ) assert accept_lengths.tolist() == [3, 1, 2] assert accepted_index.tolist() == expected_mask - assert manager.main_kda_state.mtp_state_index.tolist() == [0, 1, 2] - assert manager.main_gdn_state.mtp_state_index.tolist() == [0, 1, 2] - assert manager.draft0_gdn_state.mtp_state_index.tolist() == [0, 0, 0] + assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 4689a3d29e..26dd22f929 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -143,7 +143,7 @@ def _snapshot( state = main_linear_state(resources) result["conv"] = state.conv_state[:, req_idx, ..., cursor : cursor + state.config.conv_width].float().cpu().numpy() result["ssm"] = state.ssm_state[:, req_idx * 3 + cursor].cpu().numpy() - tail = resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, req_idx] + tail = resources.main_kpool_c4_bf16_state.tail_buffer[:, req_idx] result["tail"] = torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() return result @@ -189,7 +189,7 @@ def _worker(connection: Connection, backend: str, sender: bool): field.tensor.copy_(values) state = main_linear_state(resources) for i, tensor in enumerate( - (state.conv_state, state.ssm_state, resources.main_indexer_kpool_c4_tail_state.tail_buffer) + (state.conv_state, state.ssm_state, resources.main_kpool_c4_bf16_state.tail_buffer) ): tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) expected = _snapshot(resources, full, [3, 0, 7, 1, 9], 1, 0) @@ -254,7 +254,7 @@ def _worker(connection: Connection, backend: str, sender: bool): assert main_linear_state(resources).mtp_state_index[2] == 0 assert not main_linear_state(resources).conv_state[:, :2].any() assert not main_linear_state(resources).ssm_state[:, :6].any() - assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, :2].any() + assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, :2].any() assert not read_kv_storage(full)[:, [0, 1, 2, 3]].any() assert (pool.conv_state == 41).all() assert (pool.ssm_state == 43).all() diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index 0393431051..e47073939c 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -71,6 +71,6 @@ def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: s resources.finish_state_creation() (state,) = resources.get_full_attention_states("draft0") assert not hasattr(draft, "full_att_state") - assert state.resource_id == "draft0_full_att_bf16_state" + assert state.resource_id == "draft0_gqa_bf16_state" assert state.config.logical_layer_groups == ((2,),) assert read_kv_storage(state) is None diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 5d40e8dbca..23378f26e5 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -24,13 +24,13 @@ def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( source, _ = make_resources(glm, 1, 0, 2, draft_layers=2) target, _ = make_resources(glm, 1, 0, 2, draft_layers=2) for index, role in enumerate(("main", "draft0")): - state = getattr(source, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + state = getattr(source, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): if "token" in field.axes: field.tensor.fill_(11 + index) if glm: - getattr(source, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 0].fill_(7 + index) - getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 2].fill_(19) + getattr(source, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 0].fill_(7 + index) + getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 2].fill_(19) main_linear_state(source).conv_state.fill_(3) main_linear_state(source).ssm_state.fill_(5) page = torch.empty(source.get_pd_page_bytes(), dtype=torch.uint8, device="cuda") @@ -68,12 +68,12 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: src_req_idx=0, dst_req_idx=2, ) for index, role in enumerate(("main", "draft0")): - state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): if "token" in field.axes: assert (field.tensor[:, 8:12] == 11 + index).all() if glm: - tail = getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer + tail = getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer expected_tail = torch.zeros_like(tail[:, 1]) for position in range(4, 8): expected_tail[:, position % tail.shape[2]].fill_(7 + index) @@ -102,10 +102,10 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: req_idx=1, ) for index, role in enumerate(("main", "draft0")): - state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_full_att_bf16_state") + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): if "token" in field.axes: assert (field.tensor[:, 8:16] == 11 + index).all() if glm: - assert not getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 1].any() - assert not getattr(target, f"{role}_indexer_kpool_c4_tail_state").tail_buffer[:, 2].any() + assert not getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 1].any() + assert not getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 2].any() diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 5cdbea0018..a184d2d8a3 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -118,8 +118,8 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( main_linear_state(source).conv_state[:, 1, ..., :2].fill_(13) main_linear_state(source).ssm_state[:, 3].fill_(17) if glm: - source.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].fill_(23) - target.main_indexer_kpool_c4_tail_state.tail_buffer[:, 2].fill_(29) + source.main_kpool_c4_bf16_state.tail_buffer[:, 1].fill_(23) + target.main_kpool_c4_bf16_state.tail_buffer[:, 2].fill_(29) torch.cuda.synchronize(source_device) # 同时验证运行态与指定大页复制，保留其他页及小页快照。 for pool in (main_linear_state(target).big_pages_pool, main_linear_state(target).small_pages_pool): @@ -179,9 +179,9 @@ def barrier(group: str): assert (main_linear_state(target).small_pages_pool.conv_state == 41).all() assert (main_linear_state(target).small_pages_pool.ssm_state == 43).all() if glm: - assert not target.main_indexer_kpool_c4_tail_state.tail_buffer[:, 2].any() - assert (source.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1] == 23).all() - assert not target.main_indexer_kpool_c4_tail_state.tail_buffer[:, :2].any() + assert not target.main_kpool_c4_bf16_state.tail_buffer[:, 2].any() + assert (source.main_kpool_c4_bf16_state.tail_buffer[:, 1] == 23).all() + assert not target.main_kpool_c4_bf16_state.tail_buffer[:, :2].any() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -214,7 +214,7 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( for storage_buffer in full_state_test_buffers(full).values(): storage_buffer[:, 4].zero_() main_linear_state(resources).conv_state.fill_(11) - resources.main_indexer_kpool_c4_tail_state.tail_buffer.fill_(7) + resources.main_kpool_c4_bf16_state.tail_buffer.fill_(7) for state in resources.get_bounded_attention_states(): for pool in (state.big_pages_pool, state.small_pages_pool): if pool is not None: @@ -237,7 +237,7 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( error = parent.recv() assert error is None, error assert (read_kv_storage(full)[:, 4] == 9).all() - assert not resources.main_indexer_kpool_c4_tail_state.tail_buffer[:, 1].any() + assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, 1].any() assert (main_linear_state(resources).conv_state == 11).all() assert (main_linear_state(resources).big_pages_pool.conv_state[slot[0]] == 11).all() assert not main_linear_state(resources).big_pages_pool.ssm_state[slot[0]].any() diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index 983fa4c6e5..c1980b3339 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -21,11 +21,11 @@ @pytest.mark.parametrize( "kv_type, expected_type, storage_dtype, width", [ - ("None", "full_att_bf16", torch.bfloat16, 128), - ("int4kv", "full_att_int4_group16_bf16_scale", torch.int8, 64), - ("int8kv", "full_att_int8_group16_bf16_scale", torch.int8, 128), - ("fp8kv_spt", "full_att_fp8_static_per_tensor", torch.uint8, 128), - ("fp8kv_sph", "full_att_fp8_static_per_head", torch.uint8, 128), + ("None", "gqa_bf16", torch.bfloat16, 128), + ("int4kv", "gqa_int4_group16_bf16_scale", torch.int8, 64), + ("int8kv", "gqa_int8_group16_bf16_scale", torch.int8, 128), + ("fp8kv_spt", "gqa_fp8_static_per_tensor", torch.uint8, 128), + ("fp8kv_sph", "gqa_fp8_static_per_head", torch.uint8, 128), ], ) def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( @@ -66,7 +66,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( (False, "None", 576, 576, "mla_bf16"), (True, "None", 576, 648, "mla_packed_bf16"), (True, "None", 512, 584, "mla_packed_bf16"), - (True, "fp8kv_dsa", 576, 800, "mla_packed_fp8"), + (True, "fp8kv_dsa", 576, 800, "mla_packed_fp8_group128"), ], ) def test_mla_creation_computes_packed_width_from_logical_dimensions( @@ -107,7 +107,7 @@ def test_kpool_creation_preserves_pool_and_mtp_dimensions(): state = create_kpool_state(config=config, request_capacity=3, mtp_step=2, device="cpu") assert state.config is config and state.resource_id is None assert state.tail_buffer.shape == (2, 3, 6, 256) - assert state.get_state_type() == "indexer_kpool_c4_tail" + assert state.get_state_type() == "kpool_c4_bf16" def test_create_rejects_encoding_from_another_logical_family(): diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 4bc1d07c7c..5fe91caad1 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -62,8 +62,8 @@ def test_numbered_drafts_keep_resource_id_order_and_unique_members(): assert [state.resource_id.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ "draft0", "draft10", "draft1", "draft2", "draft35", "main", ] - assert manager.get_full_attention_states("draft1") == (manager.draft1_full_att_bf16_state,) - assert manager.get_full_attention_states("draft10") == (manager.draft10_full_att_bf16_state,) + assert manager.get_full_attention_states("draft1") == (manager.draft1_gqa_bf16_state,) + assert manager.get_full_attention_states("draft10") == (manager.draft10_gqa_bf16_state,) def test_custom_state_type_controls_registered_resource_name(): @@ -72,14 +72,14 @@ def test_custom_state_type_controls_registered_resource_name(): class CustomGQAState(GQAFullAttentionState): def get_state_type(self) -> str: - return "full_att_custom_fp32" + return "gqa_custom_fp32" state = full_state(CustomGQAState) manager = AttResourceManager(request_capacity=1) manager.set_att_state("draft0", state) manager.finish_state_creation() - assert manager.draft0_full_att_custom_fp32_state is state - assert state.resource_id == "draft0_full_att_custom_fp32_state" + assert manager.draft0_gqa_custom_fp32_state is state + assert state.resource_id == "draft0_gqa_custom_fp32_state" assert manager.get_full_attention_states("draft0") == (state,) assert not hasattr(state, "state_type") @@ -209,13 +209,13 @@ def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_la overlapping = allocated_state("static_head") manager.set_att_state("main", overlapping) manager.finish_state_creation() - assert manager.main_full_att_bf16_state is main - assert manager.draft0_full_att_bf16_state is draft - assert manager.main_full_att_int8_group8_bf16_scale_state is quant + assert manager.main_gqa_bf16_state is main + assert manager.draft0_gqa_bf16_state is draft + assert manager.main_gqa_int8_group8_bf16_scale_state is quant assert manager.get_full_attention_states("draft0") == (draft,) expected = [main, quant] if overlap: - assert manager.main_full_att_fp8_static_per_head_state is overlapping + assert manager.main_gqa_fp8_static_per_head_state is overlapping expected.append(overlapping) assert manager.get_full_attention_states("main") == tuple(sorted(expected, key=lambda state: state.resource_id)) @@ -225,8 +225,8 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c cpu_args.cpu_cache_token_page_size = 4 full = allocated_state("gqa", device="cuda") quant = allocated_state("int8", device="cuda") - assert full.get_state_type() == "full_att_bf16" - assert quant.get_state_type() == "full_att_int8_group8_bf16_scale" + assert full.get_state_type() == "gqa_bf16" + assert quant.get_state_type() == "gqa_int8_group8_bf16_scale" for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): field.tensor.fill_(5) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 375e58cf98..4ed5aa0326 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -66,10 +66,10 @@ def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): results = {(query.__name__, role): query(role) for query in queries for role in (None, "main", "draft0")} assert manager._get_states.cache_info().currsize == 9 assert manager.get_full_attention_states() == ( - manager.draft0_mla_fp32_state, manager.main_full_att_fp32_state, + manager.draft0_mla_fp32_state, manager.main_gqa_fp32_state, ) - assert manager.get_bounded_attention_states("main") == (manager.main_kda_state,) - assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_state,) + assert manager.get_bounded_attention_states("main") == (manager.main_kda_conv_fp32_ssm_fp32_state,) + assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_conv_fp32_ssm_fp32_state,) assert manager.get_attention_tail_states("draft0") == () hits = manager._get_states.cache_info().hits @@ -79,7 +79,7 @@ def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): with pytest.raises(AssertionError, match="unknown model role"): query("target") assert manager._get_states.cache_info().hits == hits + 9 - manager.main_kda_state.conv_state.fill_(7) + manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.fill_(7) assert (manager.get_bounded_attention_states("main")[0].conv_state == 7).all() @@ -172,7 +172,7 @@ def test_clearing_query_cache_releases_manager_and_states(): manager.get_bounded_attention_states() manager.get_attention_tail_states() manager_ref = weakref.ref(manager) - state_ref = weakref.ref(manager.main_kda_state) + state_ref = weakref.ref(manager.main_kda_conv_fp32_ssm_fp32_state) del manager gc.collect() assert manager_ref() is not None and state_ref() is not None diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 414a1d766a..6d467fbc83 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -108,8 +108,8 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t small_page_capacity=2, ) full = main_full - linear = resources.main_kda_state - tail = resources.main_indexer_kpool_c4_tail_state + linear = resources.main_kda_conv_bf16_ssm_fp32_state + tail = resources.main_kpool_c4_bf16_state assert not hasattr(req, "req_to_conv_state") att_kv = full.get_layer_view(3).kv assert att_kv.shape == (17, 1, 512) @@ -124,7 +124,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t assert torch.equal(index_bytes, expected_index_bytes) assert tail.get_layer_view(3).shape == (4, 4 + mtp_step, 256) if mtp_step: - assert resources.draft0_indexer_kpool_c4_tail_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) + assert resources.draft0_kpool_c4_bf16_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) assert draft_full.get_layer_view(4).kv.shape == att_kv.shape kind = "small" if small_page else "big" allocator = resources.bounded_att_small_page_allocator if small_page else resources.bounded_att_big_page_allocator diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index cd3af18217..05d509feaa 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -590,8 +590,10 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): infer = SimpleNamespace( att_resource_role="main", att_resource_manager=SimpleNamespace( - main_indexer_kpool_c4_tail_state=SimpleNamespace(get_layer_view=lambda logical_layer_id: tail), - main_mla_packed_bf16_state=full_att_state, + get_att_state=lambda att_type, **kwargs: ( + SimpleNamespace(get_layer_view=lambda logical_layer_id: tail) + if "kpool" in att_type else full_att_state + ), ), is_draft_model=False, is_prefill=True, @@ -616,7 +618,6 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5, "index_kpool": 4}, 1, ) - indexer.full_att_state_name = "main_mla_packed_bf16_state" _, full = indexer._get_indices(hidden, hidden, infer, state, weights) assert full[-1, 2048].item() == tokens - 1 # Always-selected incomplete tail. assert full[0, 0].item() == 0 and (full[0, 1:] == -1).all() diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 455b90b1c9..0399b490a8 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -170,9 +170,9 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) state = InferStateInfo() state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - layer.full_att_state_name = "main_mla_packed_bf16_state" + state.att_resource_role = "main" state.att_resource_manager = SimpleNamespace( - main_mla_packed_bf16_state=SimpleNamespace( + get_att_state=lambda **kwargs: SimpleNamespace( get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv) ) ) @@ -423,8 +423,8 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): for model in (main, *drafts): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) main.prepare_att_resources() - assert main.att_resource_manager.main_kda_state.conv_state.is_cuda - assert main.att_resource_manager.main_indexer_kpool_c4_tail_state.tail_buffer.is_cuda + assert main.att_resource_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.is_cuda + assert main.att_resource_manager.main_kpool_c4_bf16_state.tail_buffer.is_cuda for step, draft in enumerate(drafts): draft.prepare_att_resources() assert draft.att_resource_manager is main.att_resource_manager @@ -432,11 +432,11 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): state = getattr(draft.att_resource_manager, f"draft{step}_mla_packed_bf16_state") assert state is not main.att_resource_manager.main_mla_packed_bf16_state assert not hasattr(draft, "full_att_state") - assert state.storage.device.type == "meta" + assert state.storage is None assert state.config.logical_layer_groups == ((4 + step,),) assert state.resource_id == f"draft{step}_mla_packed_bf16_state" assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None - assert getattr(draft.att_resource_manager, f"draft{step}_indexer_kpool_c4_tail_state").tail_buffer.is_cuda + assert getattr(draft.att_resource_manager, f"draft{step}_kpool_c4_bf16_state").tail_buffer.is_cuda # NoPE state initialization must not require a model-owned RoPE cache. for model in (main, *drafts): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index a6074a4661..777cd6a5f9 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -97,8 +97,8 @@ def create(tp, target_mtp_step=None): SimpleNamespace( full=main, full_states=tuple(state for state in (main, draft) if state is not None), - linear=resources.main_kda_state, - tail=resources.main_indexer_kpool_c4_tail_state, + linear=resources.main_kda_conv_bf16_ssm_fp32_state, + tail=resources.main_kpool_c4_bf16_state, tails=resources.get_attention_tail_states(), resources=resources, ) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 54c859875c..a988f3fa28 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -106,7 +106,7 @@ def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( full, resources, requests, config = linear_cache disable_chunked_mm_input = get_env_start_args().disable_chunked_mm_input first_page_has_state = not disable_chunked_mm_input - linear = resources.main_gdn_state + linear = resources.main_gdn_conv_bf16_ssm_fp32_state big = linear.big_pages_pool small = linear.small_pages_pool big_states = {} diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index c7ff8758b8..699ff79ada 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -39,7 +39,7 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): resources = composition() slots = init_source(resources) - tail = resources.main_indexer_kpool_c4_tail_state + tail = resources.main_kpool_c4_fp32_state resources.reset_bounded_att_state([1]) assert (tail.tail_buffer[:, 1] == 9).all() resources.restore_bounded_att_pages([1], page_kind="big", slot_ids=[slots[-1]]) @@ -67,7 +67,7 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): assert state.mtp_state_index.tolist() == [2, 0, 2] assert state.conv_state[:, 0].any() pool = state.big_pages_pool - expected_conv, expected_ssm = (23, 47) if state is resources.main_kda_state else (24, 48) + expected_conv, expected_ssm = (23, 47) if state is resources.main_kda_conv_fp32_ssm_fp32_state else (24, 48) assert (pool.conv_state[slots[-1]] == expected_conv).all() assert (pool.ssm_state[slots[-1]] == expected_ssm).all() assert not tail.tail_buffer[:, 1].any() From 08847e7af25f600fb6161b64a9e4f155aa188a4c Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 02:42:06 +0000 Subject: [PATCH 04/56] refactor: create and bind shared request manager during startup --- docs/cache_manager_design.md | 12 ++++++---- lightllm/common/basemodel/basemodel.py | 24 +++++++------------ lightllm/models/deepseek_mtp/model.py | 4 ---- lightllm/models/glm4_moe_lite_mtp/model.py | 3 --- lightllm/models/glm5_next_mtp/model.py | 3 --- lightllm/models/mistral_mtp/model.py | 4 ---- lightllm/models/qwen3_5_mtp/model.py | 4 ---- lightllm/models/qwen3_dflash/model.py | 3 --- lightllm/models/qwen3_eagle/model.py | 3 --- lightllm/models/qwen3_moe_mtp/model.py | 4 ---- .../model_infer/mode_backend/base_backend.py | 9 ++++++- test/att_resources/validate_model_cache.py | 9 ++++++- .../static_inference/static_benchmark.py | 9 ++++++- .../att_resources/test_backend_startup.py | 22 +++++++++++++++-- .../test_model_initialization.py | 11 +++++++-- unit_tests/models/glm5_next/test_layers.py | 2 +- 16 files changed, 69 insertions(+), 57 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index c8cd6e3137..32fbf29b36 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1603,23 +1603,25 @@ P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_s ### 7.1 初始化与绑定 -BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己的接口，各 draft 调用各自的接口，模型方法不接收其他模型或组织跨模型循环。`ModeBackend.init_model()` 直接组织 main/draft 构造、资源准备、共享分配和各模型运行时初始化，并在全部初始化成功后发布就绪标志。benchmark 和资源验证脚本在各自调用点展开相同流程，不另设跨模型运行时启动封装。模型构造保留配置、权重、请求管理器和计算层的准备。 +BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己的接口，各 draft 调用各自的接口，模型方法不接收其他模型或组织跨模型循环。`ModeBackend.init_model()` 直接组织 main/draft 构造、共享请求管理器创建与绑定、资源准备、共享分配和各模型运行时初始化，并在全部初始化成功后发布就绪标志。benchmark 和资源验证脚本在各自调用点展开相同流程，不另设跨模型运行时启动封装。模型构造保留配置、权重和计算层的准备，不创建或读取请求管理器。 | 单模型公共接口 | 封装的内部操作 | | --- | --- | +| `get_req_manager_max_seq_len()` | 返回请求映射所需长度容量，取本模型 `batch_max_tokens`、`max_seq_length` 与 0 的最大值；未配置的值按 0 处理，不表示模型最大上下文能力 | +| `bind_req_manager(req_manager)` | 绑定调用方创建的共享实例，不判断 main/draft，也不通过 `main_model` 获取资源 | | `prepare_att_resources()` | 在本模型 KV 内存分区内调用 `_init_att_resource_manager` 准备并注册自身全部 State | | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | -1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。ReqManager 仍由 main/draft 共用，其构造不依赖 attention 资源已经初始化。 +1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。main 的资源 hook 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，其资源 hook 通过 `set_att_state` 只添加本模型的 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。静态 FP8 标定偏移由模型侧配置 helper 计算：main 从 0 开始，draft 累加 main 及前序 draft 的实际 full-attention 层数；混合模型排除 bounded-attention 层，原生 MTP 使用自己的层定义，无 attention 的 draft 不占标定行。公共准备接口及资源 hook 均不接收标定偏移，也不返回层数。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 -3. backend 检查 main 和所有 draft 使用相同请求容量且引用同一个管理器，随后调用 `manager.initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)`。管理器先标记 State 创建完成；每个模型的成员已经直接保存在共享实例中，无需合并管理器或 State。不同 draft 即使方案相同，也可拥有不同的配置、存储和游标。 +3. backend 检查 main 和所有 draft 引用同一个 attention 资源管理器，随后调用 `manager.initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)`。管理器先标记 State 创建完成；每个模型的成员已经直接保存在共享实例中，无需合并管理器或 State。不同 draft 即使方案相同，也可拥有不同的配置、存储和游标。 4. `finish_state_creation()` 完成后，通过管理器接口计算容量并初始化实际存储。FullAttentionState 原地分配 buffer；不同 State 可以具有不同物理布局，但共享 allocator 的索引域及容量；同一分配域只创建一个 allocator。实际存储分配完成后再配置 CPU/PD 布局。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 -BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由单模型公共阶段接口调用，构造函数不调用运行时初始化。main 在资源 hook 中选择具体管理器，各模型分别在自己的 hook 中准备所需 attention State，子类可以继承适用的模型族实现。原生 MTP 如需复用 main 的资源配置，也在自身资源准备阶段读取；ReqManager 初始化只处理请求映射共享。 +BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由单模型公共阶段接口调用，构造函数不调用运行时初始化。main 在资源 hook 中选择具体管理器，各模型分别在自己的 hook 中准备所需 attention State，子类可以继承适用的模型族实现。原生 MTP 如需复用 main 的资源配置，也在自身资源准备阶段读取；ReqManager 的创建与共享关系由 backend 组织，模型只提供容量需求和公共绑定接口。 -需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方可按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 +需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_storage_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index df4730d69d..c7e8228619 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -108,8 +108,6 @@ def __init__(self, kvargs): enable_weight_cpu_backup = self.args.enable_weight_cpu_backup with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): self._init_weights() - with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - self._init_req_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -129,6 +127,14 @@ def att_resource_role(self) -> str: assert draft_index < MAX_DRAFT_MODELS, "at most 36 draft models are supported" return f"draft{draft_index}" + def get_req_manager_max_seq_len(self) -> int: + """返回请求到 token 索引映射所需的长度容量，不表示模型支持的最大上下文长度。""" + return max(0, self.batch_max_tokens or 0, self.max_seq_length or 0) + + def bind_req_manager(self, req_manager: ReqManager) -> None: + """绑定启动流程创建的共享请求管理器；main/draft 共用请求槽位和 token 索引映射。""" + self.req_manager = req_manager + def prepare_att_resources(self) -> None: """Prepare this model's attention States from its own configuration.""" assert getattr(self, "att_resource_manager", None) is None, "resources already initialized" @@ -240,20 +246,6 @@ def _check_mem_size(self): return - def _init_req_manager(self): - if self.is_mtp_draft_model: - self.req_manager = self.main_model.req_manager - return - create_max_seq_len = 0 - - if self.batch_max_tokens is not None: - create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens) - if self.max_seq_length is not None: - create_max_seq_len = max(create_max_seq_len, self.max_seq_length) - - self.req_manager = ReqManager(self.max_req_num, create_max_seq_len) - return - def _init_infer_layer(self, start_layer_index=0): self.pre_infer = self.pre_layer_infer_class(network_config=self.config) self.post_infer = self.post_layer_infer_class(network_config=self.config) diff --git a/lightllm/models/deepseek_mtp/model.py b/lightllm/models/deepseek_mtp/model.py index c5c4a9aa61..c151d8efcd 100644 --- a/lightllm/models/deepseek_mtp/model.py +++ b/lightllm/models/deepseek_mtp/model.py @@ -28,10 +28,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.pre_post_weight = self.pre_and_post_weight_class( diff --git a/lightllm/models/glm4_moe_lite_mtp/model.py b/lightllm/models/glm4_moe_lite_mtp/model.py index 245af80505..3e5c34e1df 100644 --- a/lightllm/models/glm4_moe_lite_mtp/model.py +++ b/lightllm/models/glm4_moe_lite_mtp/model.py @@ -37,9 +37,6 @@ def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - def _init_weights(self, start_layer_index=None): assert start_layer_index is None diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 294c350e91..9cecde71ff 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -59,9 +59,6 @@ def _init_some_value(self): self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - def _init_att_resource_manager(self): from lightllm.models.glm5_next.att_resource_manager import init_model_att_tail from lightllm.common.basemodel.basemodel_utils import init_model_full_attention diff --git a/lightllm/models/mistral_mtp/model.py b/lightllm/models/mistral_mtp/model.py index d85b60754e..e8458846f9 100644 --- a/lightllm/models/mistral_mtp/model.py +++ b/lightllm/models/mistral_mtp/model.py @@ -41,10 +41,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.config["n_layer"] = 1 diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py index 61c920f23b..0c31b7c853 100644 --- a/lightllm/models/qwen3_5_mtp/model.py +++ b/lightllm/models/qwen3_5_mtp/model.py @@ -37,10 +37,6 @@ def _resource_config(cls, model_config): config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1) return config - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - def _init_config(self): super()._init_config() # MTP draft model: reuses the main model's config, but overrides the following: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 61ad02d9c4..3b2432c77c 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -46,9 +46,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - def _init_att_backend(self): super()._init_att_backend() # FA3 is currently the only backend that supports non-causal block decode. diff --git a/lightllm/models/qwen3_eagle/model.py b/lightllm/models/qwen3_eagle/model.py index 490d06fba3..42d273b31f 100644 --- a/lightllm/models/qwen3_eagle/model.py +++ b/lightllm/models/qwen3_eagle/model.py @@ -46,9 +46,6 @@ def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - def _init_weights(self, start_layer_index=None): self.pre_post_weight = self.pre_and_post_weight_class( self.data_type, network_config=self.config, quant_cfg=self.quant_cfg diff --git a/lightllm/models/qwen3_moe_mtp/model.py b/lightllm/models/qwen3_moe_mtp/model.py index e5d5f06781..da6a1dc76f 100644 --- a/lightllm/models/qwen3_moe_mtp/model.py +++ b/lightllm/models/qwen3_moe_mtp/model.py @@ -33,10 +33,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None mtp_index = len(self.mtp_previous_draft_models) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 0e68642414..f94899143e 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -150,13 +150,20 @@ def init_model(self, kvargs): from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.att_resources import AttResourceManager + from lightllm.common.req_manager import ReqManager models = (self.model, *self.draft_models) assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(item.max_req_num == self.model.max_req_num for item in models), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=self.model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) for item in models: item.prepare_att_resources() manager: AttResourceManager = self.model.att_resource_manager - assert all(item.max_req_num == self.model.max_req_num for item in models), "models must share the request domain" assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.initialize_resources( diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index da9a805c41..3b477e3e83 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -293,13 +293,20 @@ def main(): from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.att_resources import AttResourceManager + from lightllm.common.req_manager import ReqManager models = (model, *draft_models) assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) for item in models: item.prepare_att_resources() manager: AttResourceManager = model.att_resource_manager - assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.initialize_resources( diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 1fd3a2cc24..86246f3835 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -1274,13 +1274,20 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.att_resources import AttResourceManager + from lightllm.common.req_manager import ReqManager models = (model, *draft_models) assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) for item in models: item.prepare_att_resources() manager: AttResourceManager = model.att_resource_manager - assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.initialize_resources( diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e5e6a34208..c826fbe1bc 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -125,6 +125,15 @@ def test_backend_allocates_all_roles_before_warmup( models, warmed, padded, captured, prefill_captured, ready = [], set(), set(), set(), set(), [] checked = [] mapping = SimpleNamespace(bind_hold_page=lambda indices: None) + request_allocations = [] + + def create_req_manager(max_request_num: int, max_sequence_length: int): + assert len(models) == 1 + draft_count + assert all(not hasattr(item, "req_manager") for item in models) + request_allocations.append((max_request_num, max_sequence_length)) + return mapping + + monkeypatch.setattr("lightllm.common.req_manager.ReqManager", create_req_manager) def mark_backends(model: TpPartBaseModel): # 第一个模型初始化 backend 前，全部 State 已创建完成，可通过缓存查询取得资源。 @@ -145,6 +154,7 @@ def mark_backends(model: TpPartBaseModel): def warmup(model: TpPartBaseModel): assert model.backends_bound + assert all(item.req_manager is mapping for item in models) assert all(item.att_resource_manager is models[0].att_resource_manager for item in models) assert model.att_resource_manager.full_att_allocator.size == 32 assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) @@ -171,11 +181,10 @@ def check_max_len(model: TpPartBaseModel): model.is_mtp_draft_model = bool(index) model.args, model.config = args, draft_config.copy() if index else config.copy() model.max_req_num, model.max_total_token_num, model.mem_fraction = 2, None, 0.9 - model.max_seq_length, model.batch_max_tokens = 32, 16 + model.max_seq_length, model.batch_max_tokens = (32 if index else 24), 16 model.tp_world_size_, model.tp_rank_ = 1, 0 model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(offset, offset + model.config["n_layer"])] offset += len(model.layers_infer) - model.req_manager = mapping model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) if index: model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() @@ -190,6 +199,7 @@ def check_max_len(model: TpPartBaseModel): models.append(model) def memory_info(device: str | torch.device | int): + assert all(item.req_manager is mapping for item in models) manager = models[0].att_resource_manager states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) @@ -228,6 +238,13 @@ def start_runtime(): dict(args=args, world_size=1, load_way="HF", weight_dir=str(tmp_path), max_total_token_num=None) ) else: + assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + req_manager = create_req_manager( + max_request_num=models[0].max_req_num, + max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), + ) + for item in models: + item.bind_req_manager(req_manager=req_manager) for index, item in enumerate(models): item.prepare_att_resources() assert all(other.att_resource_manager is models[0].att_resource_manager for other in models[: index + 1]) @@ -245,6 +262,7 @@ def start_runtime(): } start_runtime() + assert request_allocations == [(2, 32)] assert captured == {id(model) for model in models} assert checked == [models[0]] assert ready == ([True] if via_backend else []) diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index a1f254edf3..c78b4f7f53 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -108,7 +108,6 @@ def check_runtime_before_profile(device: str | torch.device | int): return resources.get_full_att_token_bytes() * (32 + args.page_size) + reserve, 0 monkeypatch.setattr(torch.cuda, "mem_get_info", check_runtime_before_profile) - model._init_req_manager() prepare_model_states(model) drafts = [] if mtp: @@ -119,11 +118,19 @@ def check_runtime_before_profile(device: str | torch.device | int): draft.tp_rank_, draft.tp_world_size_ = 0, 1 draft.max_req_num, draft.max_total_token_num = model.max_req_num, None draft.config = dict(model.config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) - draft._init_req_manager() + draft.max_seq_length, draft.batch_max_tokens = model.max_seq_length, model.batch_max_tokens prepare_model_states(draft, logical_layer_ids=(4,)) assert draft.att_resource_manager is model.att_resource_manager assert read_kv_storage(owned_states(draft.att_resource_manager.__dict__, FullAttentionState)[0]) is None drafts.append(draft) + models = (model, *drafts) + req_manager = ReqManager( + max_request_num=model.max_req_num, + max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), + ) + for item in models: + item.bind_req_manager(req_manager=req_manager) + assert all(item.req_manager is req_manager for item in models) model.att_resource_manager.initialize_resources( args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 0399b490a8..ddaf9fc602 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -404,7 +404,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): draft._init_weights() draft._init_infer_layer() draft._init_some_value() - draft._init_req_manager() + draft.bind_req_manager(req_manager=main.req_manager) draft.max_req_num = main.max_req_num assert main.config["mhc"] and not draft.config["mhc"] assert draft.config["layer_types"] is not main.config["layer_types"] From df9e5ec20e97010556c7e266b2eaaf77cddf27db Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 02:53:44 +0000 Subject: [PATCH 05/56] refactor: remove model revision namespace from disk cache --- docs/cache_manager_design.md | 4 +-- .../multi_level_kv_cache/disk_cache_worker.py | 3 +- lightllm/utils/att_resource_config.py | 17 --------- lightllm/utils/kv_cache_utils.py | 4 +-- .../model_infer/test_cpu_resource_layout.py | 35 ++----------------- 5 files changed, 6 insertions(+), 57 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 32fbf29b36..af891f31d7 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1652,7 +1652,7 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 `basemodel_utils.py` 提供 `get_resource_config(model_config, *, model_class=None)` 和 `get_full_attention_options(config, args, *, model_class=None)`，统一处理字段别名、text config 和默认 KV 维度。它们是普通工具函数，BaseModel 不再声明对应的 `_resource_config`、`_full_attention_options` 方法；具体模型仅在存在专有规则时提供可选 hook，并调用工具函数取得默认配置。传入 `model_class` 时，工具函数选择该模型的专有实现，否则使用通用逻辑。 -同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域；磁盘缓存命名空间包含 main/draft 模型版本。静态量化标定按本模型实际 full-attention 层区间选取，CPU 服务和推理端复用相同的布局构建流程。 +同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域。静态量化标定按本模型实际 full-attention 层区间选取，CPU 服务和推理端复用相同的布局构建流程。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 @@ -1717,7 +1717,7 @@ P/D 使用相同模型和 main/draft 资源集合，允许不同 TP 与 `mtp_ste ### 8.3 CPU/Disk 与 DP fetch -CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布局保存数据，不预留格式页头或状态长度、版本字段。各资源区域由模型配置与 State 字段确定，TP 副本由明确的单个写入者更新共享区域。CPU 服务绑定时核对页大小，读回长度由任务参数提供，前缀与槽位的关联由调用方维护。磁盘文件按 main/draft 模型路径及文件版本隔离。 +CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布局保存数据，不预留格式页头或状态长度、版本字段。各资源区域由模型配置与 State 字段确定，TP 副本由明确的单个写入者更新共享区域。CPU 服务绑定时核对页大小，读回长度由任务参数提供，前缀与槽位的关联由调用方维护。磁盘缓存使用配置目录下的固定文件名 `cache_file`。 GPU→CPU 保存状态页完成后，才能在 CPU 上重排或复制这些快照；主机侧读取必须等待保存完成，不能只向另一个 CUDA stream 排入依赖。CPU 页的 ready 发布仍由调用流程在全部参与资源与 TP 分片传输完成后执行。 diff --git a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py index fd9534dccb..542ddbd877 100644 --- a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py +++ b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py @@ -54,8 +54,7 @@ def __init__( if not cache_dir: cache_dir = os.path.join(tempfile.gettempdir(), f"lightllm_disk_cache_{get_unique_server_name()}") os.makedirs(cache_dir, exist_ok=True) - namespace = cpu_cache_client.kv_cache_tensor_meta.cache_namespace - cache_file = os.path.join(cache_dir, f"cache_file_{namespace}" if namespace else "cache_file") + cache_file = os.path.join(cache_dir, "cache_file") self.max_concurrent_write_tasks = max_concurrent_write_tasks self._page_major_tensor = self._prepare_tensor(cpu_cache_client.cpu_kv_cache_tensor) diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/att_resource_config.py index d8bcd9c5ea..d0796ec082 100644 --- a/lightllm/utils/att_resource_config.py +++ b/lightllm/utils/att_resource_config.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING import hashlib import json -from pathlib import Path import torch @@ -204,19 +203,3 @@ def build_cpu_cache_layouts(): args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num ), "hybrid CPU pages must match the large state-page token boundary" return manager.get_cpu_transfer_layouts() - - -def cpu_cache_namespace(): - """按 main/draft 模型路径及文件版本隔离磁盘缓存。""" - args = get_env_start_args() - directories = [args.model_dir, *(args.mtp_draft_model_dir or [])] - revisions = [] - for directory in directories: - directory = Path(directory).resolve() - artifacts = [] - for pattern in ("*.safetensors", "*.bin", "*.index.json", "config.json"): - for path in sorted(directory.glob(pattern)): - stat = path.stat() - artifacts.append((path.name, stat.st_size, stat.st_mtime_ns)) - revisions.append((str(directory), artifacts)) - return hashlib.sha256(repr(revisions).encode()).hexdigest() diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 923fb9aada..9f55465c7c 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -51,7 +51,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - from lightllm.utils.att_resource_config import build_cpu_cache_layouts, cpu_cache_namespace + from lightllm.utils.att_resource_config import build_cpu_cache_layouts layouts = build_cpu_cache_layouts() page_bytes = 0 @@ -62,7 +62,6 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": cpu_cache_meta = CpuKVCacheMeta( page_num=0, page_bytes=page_bytes, - cache_namespace=cpu_cache_namespace(), ) cpu_cache_page_num = int( @@ -81,7 +80,6 @@ class CpuKVCacheMeta: page_num: int page_bytes: int - cache_namespace: str = "" def __post_init__(self): if self.page_num < 0 or self.page_bytes <= 0: diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index 837296705c..6acb945931 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -21,7 +21,6 @@ from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.att_resource_config import cpu_cache_namespace from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_unique_server_name from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta @@ -277,29 +276,6 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( assert not buffer[:2].any() -def test_cpu_namespace_isolates_model_revision(runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): - args = StartArgs(model_dir=str(tmp_path)) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - get_env_start_args.cache_clear() - (tmp_path / "config.json").write_text("{}") - initial = cpu_cache_namespace() - (tmp_path / "model.safetensors").write_bytes(b"first revision") - updated = cpu_cache_namespace() - assert updated != initial - (tmp_path / "model.safetensors").write_bytes(b"second revision") - assert cpu_cache_namespace() != updated - - -def test_cpu_namespace_includes_draft_revision(runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): - draft_dir = tmp_path / "draft" - draft_dir.mkdir() - args = StartArgs(model_dir=str(tmp_path), mtp_draft_model_dir=[str(draft_dir)]) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - get_env_start_args.cache_clear() - before = cpu_cache_namespace() - (draft_dir / "model.safetensors").write_bytes(b"new draft weights") - assert cpu_cache_namespace() != before - @pytest.mark.parametrize("hybrid", [False, True]) def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( @@ -415,7 +391,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() -def test_disk_preserves_resource_pages_and_isolates_new_model_revision( +def test_disk_preserves_resource_pages( runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -449,11 +425,4 @@ def test_disk_preserves_resource_pages_and_isolates_new_model_revision( for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.zero_() resources.read_req_from_cpu(client.cpu_kv_cache_tensor, pages, torch.arange(16, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts()) - # 模型版本改变后使用独立磁盘文件，即使页大小及 token 哈希相同也不命中旧数据。 - (tmp_path / "model.safetensors").write_bytes(b"new model revision") - next_meta = dataclasses.replace(meta, cache_namespace=cpu_cache_namespace()) - other_client = cpu_cache_factory(next_meta) - other = DiskCacheWorker(0.125, other_client, str(tmp_path / "disk")) - assert not any(other.service.query([11, 22])) - assert any((tmp_path / "disk").glob(f"cache_file_{meta.cache_namespace}*")) - assert any((tmp_path / "disk").glob(f"cache_file_{next_meta.cache_namespace}*")) + assert any((tmp_path / "disk").glob("cache_file*")) From c44a10decc4221afc68f6de26f2ae18a622acb22 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 03:11:44 +0000 Subject: [PATCH 06/56] refactor: simplify attention configuration and PD page transfers --- docs/cache_manager_design.md | 2 +- lightllm/common/basemodel/basemodel_utils.py | 49 +++++-- .../decode_node_impl/decode_trans_process.py | 35 ++--- .../prefill_trans_process.py | 26 ++-- .../mode_backend/pd/resource_io.py | 106 +++++++++----- lightllm/utils/att_resource_config.py | 134 +++++------------- .../att_resources/test_network_transport.py | 28 ++-- .../att_resources/test_runtime_transport.py | 34 ++--- unit_tests/models/glm5_next/test_pd_cache.py | 28 ++-- 9 files changed, 219 insertions(+), 223 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index af891f31d7..d8df49bd01 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1652,7 +1652,7 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 `basemodel_utils.py` 提供 `get_resource_config(model_config, *, model_class=None)` 和 `get_full_attention_options(config, args, *, model_class=None)`，统一处理字段别名、text config 和默认 KV 维度。它们是普通工具函数，BaseModel 不再声明对应的 `_resource_config`、`_full_attention_options` 方法；具体模型仅在存在专有规则时提供可选 hook，并调用工具函数取得默认配置。传入 `model_class` 时，工具函数选择该模型的专有实现，否则使用通用逻辑。 -同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域。静态量化标定按本模型实际 full-attention 层区间选取，CPU 服务和推理端复用相同的布局构建流程。 +同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具，直接构造 `FullAttentionConfig` 并调用算法级 `create_gqa_state` / `create_mla_state`，不再经由额外的 State 构建转发函数；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域。静态量化标定按明确的物理行起点和本模型 full-attention 物理层数选取，不再根据可选起点推断 main/draft 合并层数；CPU 服务和推理端复用相同的布局构建流程。`att_resource_config.py` 仅保留标定加载和 CPU 服务布局构建。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index d64692b16a..71636e1d2c 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -3,10 +3,13 @@ from __future__ import annotations from typing import TYPE_CHECKING +import hashlib +import json from lightllm.common.build_utils import repair_config if TYPE_CHECKING: from collections.abc import Sequence + from lightllm.common.att_resources.full_att import FullAttentionState from easydict import EasyDict from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.server.core.objs.start_args_type import StartArgs @@ -65,23 +68,51 @@ def create_full_attention_definition( tp_rank: int = 0, tp_world_size: int = 1, model_logical_layer_ids: Sequence[int] | None = None, - calibration_storage_layer_start: int | None = None, -): - """Build a meta State using the model's configuration and storage hooks.""" - from lightllm.utils.att_resource_config import build_full_attention_state + calibration_storage_layer_start: int = 0, +) -> FullAttentionState | None: + """按模型配置直接创建算法 State；静态 FP8 使用当前模型的物理层标定区间。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state, create_mla_state + from lightllm.utils.att_resource_config import load_static_fp8_calibration + from lightllm.utils.torch_dtype_utils import get_torch_dtype config = get_resource_config(model_config, model_class=model_class) options = get_full_attention_options(config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids) if not options["logical_layer_groups"]: return None - state = build_full_attention_state( - args, - **options, + mla = options.get("mla", False) + packed = options.get("packed", False) + dtype = options.get("dtype") + layout_config = FullAttentionConfig( + logical_layer_groups=options["logical_layer_groups"], + global_head_num=options["heads"], + head_dim=options["dim"], + dtype=get_torch_dtype(args.data_type) if dtype is None else dtype, tp_rank=tp_rank, tp_world_size=tp_world_size, - calibration_storage_layer_start=calibration_storage_layer_start, + replicated=options.get("replicated", False) or mla or packed, + ) + if mla or packed: + return create_mla_state( + config=layout_config, kv_type=args.llm_kv_type, packed=packed, page_size=args.page_size, device="meta", + ) + + scales = q_scales = None + calibration_id = "" + if args.llm_kv_type in ("fp8kv_sph", "fp8kv_spt"): + calibration, scales, q_scales = load_static_fp8_calibration( + args=args, + per_head=args.llm_kv_type == "fp8kv_sph", + storage_layer_count=len(layout_config.logical_layer_groups), + heads=layout_config.head_num, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + storage_layer_start=calibration_storage_layer_start, + ) + calibration_id = hashlib.sha256(json.dumps(calibration, sort_keys=True).encode()).hexdigest() + return create_gqa_state( + config=layout_config, kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + scales=scales, q_scales=q_scales, calibration_id=calibration_id, ) - return state def create_bounded_attention_definition( diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index 672cb59ac8..a14b4210e6 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -173,22 +173,13 @@ def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): copy_page( - self.resource_managers, - self.kv_move_buffer, - write=True, + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_buffer=False, token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), - page_index=0, - dp_index=dp_index, - tp_size=self.dp_world_size, - ) - copy_page( - self.resource_managers, - self.kv_move_buffer, - write=False, - token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), - page_index=0, - dp_index=dp_index, - tp_size=self.dp_world_size, + pd_page_index=0, + dp_rank_in_node=dp_index, + tp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -409,16 +400,16 @@ def read_page_to_mems_loop(self): with torch.cuda.stream(stream=self.copy_cuda_stream): copy_start_event.record(self.copy_cuda_stream) copy_page( - self.resource_managers, - self.kv_move_buffer, - write=False, + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_buffer=False, token_indices=torch.tensor( trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device ), - page_index=trans_task.dst_page_index, - dp_index=trans_task.decode_dp_index, - tp_size=self.dp_world_size, - kind=trans_task.page_kind, + pd_page_index=trans_task.dst_page_index, + dp_rank_in_node=trans_task.decode_dp_index, + tp_world_size=self.dp_world_size, + task_kind=trans_task.page_kind, req_idx=trans_task.req_idx, state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index 2fdb114225..c28f5a9b06 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -148,13 +148,13 @@ def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): copy_page( - self.resource_managers, - self.kv_move_buffer, - write=True, + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_buffer=True, token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), - page_index=0, - dp_index=dp_index, - tp_size=self.dp_world_size, + pd_page_index=0, + dp_rank_in_node=dp_index, + tp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -207,16 +207,16 @@ def local_copy_kv_loop(self): # 将kv 数据拷贝到 page 上，然后传输给 decode node，让其进行读取。 with torch.cuda.stream(stream=self.copy_cuda_stream): copy_page( - self.resource_managers, - self.kv_move_buffer, - write=True, + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_buffer=True, token_indices=torch.tensor( trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device ), - page_index=trans_task.src_page_index, - dp_index=trans_task.prefill_dp_index, - tp_size=self.dp_world_size, - kind=trans_task.page_kind, + pd_page_index=trans_task.src_page_index, + dp_rank_in_node=trans_task.prefill_dp_index, + tp_world_size=self.dp_world_size, + task_kind=trans_task.page_kind, req_idx=trans_task.req_idx, state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index 30c54d633a..6a47d6b98f 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -18,40 +18,80 @@ def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, *, def copy_page( - resources: Sequence[AttResourceManager], - buffer: torch.Tensor, + node_att_managers: Sequence[AttResourceManager], + pd_page_buffers: torch.Tensor, *, - page_index: int, - dp_index: int, - tp_size: int, + pd_page_index: int, + dp_rank_in_node: int, + tp_world_size: int, token_indices: torch.Tensor, - kind: Literal["kv", "att_state"] = "kv", + task_kind: Literal["kv", "att_state"] = "kv", req_idx: int | None = None, state_token_len: int | None = None, - write: bool, -): - ranks = resources[dp_index * tp_size : (dp_index + 1) * tp_size] - if dp_index < 0 or tp_size <= 0 or len(ranks) != tp_size or not 0 <= page_index < len(buffer): - raise ValueError("invalid PD group or page index") - rank_layouts = [] - for resource in ranks: - if kind == "kv": - layouts = resource.get_pd_full_att_transfer_layouts() - elif kind == "att_state": - layouts = resource.get_pd_bounded_att_or_tail_transfer_layouts() - else: - raise ValueError(f"unknown PD task kind {kind}") - if not layouts or resource.get_pd_page_bytes() != buffer.shape[1]: - raise ValueError("PD page layout does not match this group") - rank_layouts.append(layouts) - for resource, layouts in zip(ranks, rank_layouts): - kwargs = {"kind": kind, "layouts": layouts} - if kind == "kv": - kwargs["token_indices"] = token_indices - else: - kwargs["req_idx"] = req_idx - kwargs["state_token_len"] = state_token_len - if write: - resource.write_to_pd_buffer(buffer, pd_page_index=page_index, **kwargs) - else: - resource.read_from_pd_buffer(buffer, pd_page_index=page_index, **kwargs) + write_to_pd_buffer: bool, +) -> None: + """选择节点内指定 DP 的 TP 组，检查整组布局后提交 PD 页面读写。 + + node_att_managers 按节点内 DP 编号、组内 TP rank 顺序平铺。 + pd_page_buffers 是 PD 页集合，pd_page_index 选择本次传输的页。 + write_to_pd_buffer=True 将 State 写入发送页；False 从接收页恢复 State。 + """ + tp_att_managers = node_att_managers[dp_rank_in_node * tp_world_size : (dp_rank_in_node + 1) * tp_world_size] + assert ( + dp_rank_in_node >= 0 + and tp_world_size > 0 + and len(tp_att_managers) == tp_world_size + and 0 <= pd_page_index < len(pd_page_buffers) + ), "invalid PD group or page index" + tp_transfer_layouts = [] + for att_manager in tp_att_managers: + match task_kind: + case "kv": + state_layouts = att_manager.get_pd_full_att_transfer_layouts() + case "att_state": + state_layouts = att_manager.get_pd_bounded_att_or_tail_transfer_layouts() + case _: + assert False, f"unknown PD task kind {task_kind}" + assert state_layouts and att_manager.get_pd_page_bytes() == pd_page_buffers.shape[1], ( + "PD page layout does not match this group" + ) + tp_transfer_layouts.append(state_layouts) + + for att_manager, state_layouts in zip(tp_att_managers, tp_transfer_layouts): + match task_kind: + case "kv": + if write_to_pd_buffer: + att_manager.write_to_pd_buffer( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + token_indices=token_indices, + ) + else: + att_manager.read_from_pd_buffer( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + token_indices=token_indices, + ) + case "att_state": + if write_to_pd_buffer: + att_manager.write_to_pd_buffer( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + req_idx=req_idx, + state_token_len=state_token_len, + ) + else: + att_manager.read_from_pd_buffer( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + req_idx=req_idx, + state_token_len=state_token_len, + ) diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/att_resource_config.py index d0796ec082..9135c275f4 100644 --- a/lightllm/utils/att_resource_config.py +++ b/lightllm/utils/att_resource_config.py @@ -1,22 +1,18 @@ -"""Shared State definitions for model initialization and CPU page metadata.""" +"""静态 FP8 标定加载与 CPU cache 服务布局构建。""" from __future__ import annotations from typing import TYPE_CHECKING -import hashlib import json import torch from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.full_att import ( - FullAttentionConfig, FullAttentionState, create_gqa_state, create_mla_state, -) -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.torch_dtype_utils import get_torch_dtype +from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from collections.abc import Mapping + from lightllm.common.att_resources.types import StateTransferLayout from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs @@ -29,8 +25,8 @@ def load_static_fp8_calibration( heads: int, tp_rank: int = 0, tp_world_size: int = 1, - storage_layer_start: int | None = None, -): + storage_layer_start: int = 0, +) -> tuple[dict[str, object], torch.Tensor, torch.Tensor | None]: """读取模型和 CPU cache 共用的静态标定参数。 storage_layer_count 是当前 State 的物理层数；storage_layer_start 是其在标定表中的 @@ -39,58 +35,42 @@ def load_static_fp8_calibration( """ path = args.kv_quant_calibration_config_path if path is None: - if per_head: - raise ValueError("fp8kv_sph requires kv_quant_calibration_config_path with q_calibration") + assert not per_head, "fp8kv_sph requires kv_quant_calibration_config_path with q_calibration" return {"scale": "unit"}, torch.ones((storage_layer_count, 2), dtype=torch.float32), None with open(path) as f: cfg = json.load(f) calibrated_storage_layer_count = cfg["num_layers"] - expected_draft_storage_layer_count = get_added_mtp_kv_layer_num() if storage_layer_start is None else 0 - if storage_layer_start is not None: - if storage_layer_start < 0 or storage_layer_start + storage_layer_count > calibrated_storage_layer_count: - raise ValueError("calibration does not cover this model's attention layers") - elif per_head and ("num_target_layers" in cfg or "num_draft_layers" in cfg): - calibrated_target_storage_layer_count = cfg["num_target_layers"] - calibrated_draft_storage_layer_count = cfg["num_draft_layers"] - if ( - calibrated_draft_storage_layer_count < 0 - or calibrated_target_storage_layer_count + calibrated_draft_storage_layer_count != calibrated_storage_layer_count - or calibrated_target_storage_layer_count != storage_layer_count - expected_draft_storage_layer_count - or ( - expected_draft_storage_layer_count > 0 - and calibrated_draft_storage_layer_count != expected_draft_storage_layer_count - ) - ): - raise ValueError("calibration target/draft layer layout does not match the runtime") - elif calibrated_storage_layer_count != storage_layer_count: - raise ValueError( - f"calibration num_layers {calibrated_storage_layer_count} does not match runtime {storage_layer_count}" - ) + # 当前模型只读取自己的物理层区间；main/draft 偏移由装配流程明确提供。 + storage_layer_end = storage_layer_start + storage_layer_count + assert 0 <= storage_layer_start < storage_layer_end <= calibrated_storage_layer_count, ( + "calibration does not cover this model's attention layers" + ) if not per_head: from lightllm.utils.config_utils import get_model_architectures - if cfg["architectures"] != get_model_architectures(args.model_dir): - raise ValueError("calibration architecture does not match the model") + assert cfg["architectures"] == get_model_architectures(args.model_dir), ( + "calibration architecture does not match the model" + ) mode = "per_head" if per_head else "per_tensor" - if cfg.get("quant_type", mode) != mode: - raise ValueError("calibration quantization mode does not match the State") + assert cfg.get("quant_type", mode) == mode, "calibration quantization mode does not match the State" limits = torch.finfo(torch.float8_e4m3fn) - if cfg.get("qmin", limits.min) != limits.min or cfg.get("qmax", limits.max) != limits.max: - raise ValueError("calibration range does not match float8_e4m3fn") + assert cfg.get("qmin", limits.min) == limits.min and cfg.get("qmax", limits.max) == limits.max, ( + "calibration range does not match float8_e4m3fn" + ) - def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]): + def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]) -> torch.Tensor: tensor = torch.tensor(values["scales"], dtype=torch.float32) - if list(values["scales_shape"]) != list(shape) or list(tensor.shape) != list(shape): - raise ValueError(f"calibration scales and scales_shape must match {shape}") - if not bool((torch.isfinite(tensor) & (tensor > 0)).all()): - raise ValueError("calibration scales must be finite and positive") - first_storage_layer_index = storage_layer_start or 0 - return tensor[first_storage_layer_index : first_storage_layer_index + storage_layer_count] - - def shard(tensor: torch.Tensor, calibrated_heads: int): - if calibrated_heads <= 0 or (heads * tp_world_size) % calibrated_heads: - raise ValueError("calibration heads do not cover the model's KV heads") + assert tuple(values["scales_shape"]) == tensor.shape == shape, ( + f"calibration scales and scales_shape must match {shape}" + ) + assert bool((torch.isfinite(tensor) & (tensor > 0)).all()), "calibration scales must be finite and positive" + return tensor[storage_layer_start:storage_layer_end] + + def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: + assert calibrated_heads > 0 and (heads * tp_world_size) % calibrated_heads == 0, ( + "calibration heads do not cover the model's KV heads" + ) expanded = tensor.repeat_interleave(heads * tp_world_size // calibrated_heads, dim=-1) return expanded[:, tp_rank * heads : (tp_rank + 1) * heads].contiguous() @@ -100,63 +80,14 @@ def shard(tensor: torch.Tensor, calibrated_heads: int): k, v = all_scales.split(calibrated_heads, dim=-1) scales = torch.cat((shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) q_cfg = cfg.get("q_calibration") - if not q_cfg: - raise ValueError("fp8kv_sph calibration config requires q_calibration") + assert q_cfg, "fp8kv_sph calibration config requires q_calibration" q_scales = shard(scale_tensor(q_cfg, (calibrated_storage_layer_count, q_cfg["num_head"])), q_cfg["num_head"]) else: scales, q_scales = scale_tensor(cfg, (calibrated_storage_layer_count, 2)), None return cfg, scales, q_scales -def build_full_attention_state( - args: StartArgs | EasyDict, - *, - heads: int, - dim: int, - logical_layer_groups: tuple[tuple[int, ...], ...], - mla: bool = False, - packed: bool = False, - dtype: torch.dtype | None = None, - replicated: bool = False, - tp_rank: int = 0, - tp_world_size: int = 1, - calibration_storage_layer_start: int | None = None, -) -> FullAttentionState: - """将模型和启动配置转换为算法 create 参数，返回容量规划所需的 meta State。""" - dtype = get_torch_dtype(args.data_type) if dtype is None else dtype - config = FullAttentionConfig( - logical_layer_groups=logical_layer_groups, - global_head_num=heads, - head_dim=dim, - dtype=dtype, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - replicated=replicated or mla or packed, - ) - if mla or packed: - return create_mla_state( - config=config, kv_type=args.llm_kv_type, packed=packed, page_size=args.page_size, device="meta", - ) - scales = q_scales = None - calibration_id = "" - if args.llm_kv_type in ("fp8kv_sph", "fp8kv_spt"): - calibration, scales, q_scales = load_static_fp8_calibration( - args, - per_head=args.llm_kv_type == "fp8kv_sph", - storage_layer_count=len(logical_layer_groups), - heads=config.head_num, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - storage_layer_start=calibration_storage_layer_start, - ) - calibration_id = hashlib.sha256(json.dumps(calibration, sort_keys=True).encode()).hexdigest() - return create_gqa_state( - config=config, kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", - scales=scales, q_scales=q_scales, calibration_id=calibration_id, - ) - - -def build_cpu_cache_layouts(): +def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: """使用启动流程相同的 meta State 定义，由管理器统一生成 CPU 服务页布局。""" from transformers.configuration_utils import PretrainedConfig from lightllm.models import get_draft_model_class @@ -168,6 +99,7 @@ def build_cpu_cache_layouts(): ) from lightllm.utils.llm_utils import get_llm_model_class + # 阶段一：读取 main 和启用 attention 的 draft 配置，不构造实际模型。 args = get_env_start_args() main_config, _ = PretrainedConfig.get_config_dict(args.model_dir) models = [("main", get_llm_model_class(), main_config)] @@ -179,6 +111,7 @@ def build_cpu_cache_layouts(): config, _ = PretrainedConfig.get_config_dict(directory) models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) + # 阶段二：复用模型定义入口创建 meta State，并累计逻辑层号和标定物理行偏移。 manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 logical_layer_start = calibration_storage_layer_start = 0 for role, model_class, config in models: @@ -197,6 +130,7 @@ def build_cpu_cache_layouts(): logical_layer_start += count calibration_storage_layer_start += len(full.config.logical_layer_groups) if full is not None else 0 + # 阶段三：确认混合模型的页边界一致，由管理器聚合 CPU 联合页布局。 manager.finish_state_creation() if manager.get_bounded_attention_states(): assert ( diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 26dd22f929..3c2ad55397 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -197,16 +197,16 @@ def _worker(connection: Connection, backend: str, sender: bool): task = _task(peer, kind) with torch.cuda.stream(copy_stream): copy_page( - [resources], - buffer, - page_index=task.src_page_index, - dp_index=0, - tp_size=1, + node_att_managers=[resources], + pd_page_buffers=buffer, + pd_page_index=task.src_page_index, + dp_rank_in_node=0, + tp_world_size=1, token_indices=torch.tensor(task.mem_indexes, dtype=torch.long, device=device), - kind=kind, + task_kind=kind, req_idx=1, state_token_len=task.end_kv_index, - write=True, + write_to_pd_buffer=True, ) # The production sender publishes a page only after its copy event completes. copy_stream.synchronize() @@ -238,16 +238,16 @@ def _worker(connection: Connection, backend: str, sender: bool): # Match production: unpack on a separate stream after the done notification. with torch.cuda.stream(copy_stream): copy_page( - [resources], - buffer, - page_index=request.dst_page_index, - dp_index=0, - tp_size=1, + node_att_managers=[resources], + pd_page_buffers=buffer, + pd_page_index=request.dst_page_index, + dp_rank_in_node=0, + tp_world_size=1, token_indices=torch.tensor([8, 12, 10, 15, 11], dtype=torch.long, device=device) if kind == "kv" else torch.tensor([], dtype=torch.long, device=device), - kind=kind, + task_kind=kind, req_idx=2, state_token_len=request.end_kv_index, - write=False, + write_to_pd_buffer=False, ) copy_stream.synchronize() expected = _snapshot(resources, full, [8, 12, 10, 15, 11], 2, 0) diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index a184d2d8a3..a7c6af7a42 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -60,10 +60,10 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: page = allocate_transfer_buffer(managers[0], 1, device=device) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - copy_page([managers[0]], page, page_index=0, dp_index=0, tp_size=1, token_indices=source_indices, write=True) + copy_page(node_att_managers=[managers[0]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, write_to_pd_buffer=True) else: - copy_page([managers[0]], page, page_index=0, dp_index=0, tp_size=1, token_indices=source_indices, write=True) - copy_page([managers[1]], page, page_index=0, dp_index=0, tp_size=1, token_indices=pd_indices, write=False) + copy_page(node_att_managers=[managers[0]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, write_to_pd_buffer=True) + copy_page(node_att_managers=[managers[1]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=pd_indices, write_to_pd_buffer=False) assert_token_fields_equal(src, [1, 3], dst, [4, 6]) cpu_pages = torch.empty((1, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") @@ -287,26 +287,26 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe page = allocate_transfer_buffer(source[0][0], 1, device="cuda") for kind in ("kv", "att_state"): copy_page( - [m for m, _ in source], - page, - page_index=0, - dp_index=0, - tp_size=source_tp, + node_att_managers=[m for m, _ in source], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=source_tp, token_indices=torch.tensor([0, 3], dtype=torch.long, device="cuda"), - kind=kind, + task_kind=kind, req_idx=1, - write=True, + write_to_pd_buffer=True, ) copy_page( - [m for m, _ in target], - page, - page_index=0, - dp_index=0, - tp_size=target_tp, + node_att_managers=[m for m, _ in target], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=target_tp, token_indices=torch.tensor([4, 7], dtype=torch.long, device="cuda"), - kind=kind, + task_kind=kind, req_idx=2, - write=False, + write_to_pd_buffer=False, ) for rank, (manager, full) in enumerate(target): for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 777cd6a5f9..7881e58519 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -168,29 +168,29 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem assert source_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() == dest_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() for kind in ("kv", "att_state"): copy_page( - source_resources, - source_buffer, - page_index=0, - dp_index=0, - tp_size=prefill_tp, + node_att_managers=source_resources, + pd_page_buffers=source_buffer, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=prefill_tp, token_indices=torch.tensor(src_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), - kind=kind, + task_kind=kind, req_idx=src_req, state_token_len=length, - write=True, + write_to_pd_buffer=True, ) dest_buffer.copy_(source_buffer) copy_page( - dest_resources, - dest_buffer, - page_index=0, - dp_index=0, - tp_size=decode_tp, + node_att_managers=dest_resources, + pd_page_buffers=dest_buffer, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=decode_tp, token_indices=torch.tensor(dst_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), - kind=kind, + task_kind=kind, req_idx=dst_req, state_token_len=length, - write=False, + write_to_pd_buffer=False, ) torch.cuda.synchronize() From 1d45802c96b2e7db2ac476a17474d11917924704 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 04:08:27 +0000 Subject: [PATCH 07/56] refactor: unify model state factories and defer FP8 calibration --- docs/CN/source/models/add_new_model.md | 6 +- docs/EN/source/models/add_new_model.md | 6 +- docs/cache_manager_design.md | 121 +++++++------ .../att_resources/att_resource_manager.py | 58 +++++- .../full_att/create_gqa_state.py | 12 +- .../full_att/static_fp8_state.py | 108 +++++++---- lightllm/common/basemodel/basemodel.py | 75 +++++++- lightllm/common/basemodel/basemodel_utils.py | 95 +--------- lightllm/models/bloom/model.py | 27 ++- .../models/glm5_next/att_resource_config.py | 32 ++++ .../models/glm5_next/att_resource_manager.py | 170 ------------------ lightllm/models/glm5_next/model.py | 68 ++++--- lightllm/models/glm5_next_mtp/model.py | 41 ++++- lightllm/models/llama/model.py | 27 ++- lightllm/models/mistral/model.py | 25 ++- lightllm/models/mixtral/model.py | 27 ++- .../models/qwen3next/att_resource_manager.py | 101 ----------- lightllm/models/qwen3next/model.py | 54 +++--- lightllm/models/starcoder/model.py | 25 ++- lightllm/models/starcoder2/model.py | 27 ++- .../model_infer/mode_backend/base_backend.py | 2 + lightllm/utils/att_resource_config.py | 103 ++--------- test/att_resources/validate_model_cache.py | 2 + .../static_inference/static_benchmark.py | 2 + test/cpu_cache_kernel/test_speed.py | 32 +++- .../common/att_resources/model_helpers.py | 99 ++++++++++ .../att_resources/test_backend_startup.py | 14 ++ .../common/att_resources/test_capacity.py | 14 +- .../att_resources/test_fp8_compute_binding.py | 9 +- .../test_full_storage_initialization.py | 2 + .../test_layer_state_selection.py | 15 +- .../att_resources/test_model_calibration.py | 134 ++++++++++++-- .../test_model_initialization.py | 5 +- .../att_resources/test_model_resources.py | 24 ++- .../test_model_storage_definition.py | 45 +++++ .../test_quantized_full_state.py | 28 ++- .../att_resources/test_state_creation.py | 10 +- .../att_resources/test_state_queries.py | 1 + unit_tests/models/glm5_next/test_cache.py | 10 +- unit_tests/models/glm5_next/test_pd_cache.py | 10 +- .../model_infer/test_cpu_resource_layout.py | 2 + .../test_hybrid_cpu_state_cache.py | 10 +- 42 files changed, 938 insertions(+), 740 deletions(-) delete mode 100644 lightllm/models/glm5_next/att_resource_manager.py delete mode 100644 lightllm/models/qwen3next/att_resource_manager.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 9792fe5778..abddb2009b 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -124,13 +124,13 @@ class TpPartBaseModel: | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | 读取初始化模型的 config.json, 并进行一些 key 名的同名合法化操作 | | def _verify_params(self)： | 校验参数 | -| def _init_att_resource_manager(self): | 由具体模型创建管理器并准备全部 attention State：KV meta State、GDN/KDA 固定运行态和 indexer tail；由模型侧 helper 根据 main 和前序 draft 的配置解析静态量化标定偏移 | +| `_create_full_att_states(...)` / `_create_bounded_att_states(...)` / `_create_att_tail_states(...)` | 创建并返回当前模型的 full / bounded / tail State 列表；默认返回空列表，公共流程统一注册 | | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | -全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_att_resources()` 准备自身 State，再统一分配共享资源。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_att_resources()` 准备自身 State；随后管理器依次完成 `finish_state_creation()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 -资源准备仍通过私有 hook `_init_att_resource_manager` 交给具体模型或模型族实现。标定偏移由模型侧根据 main 和前序 draft 的 full-attention 层定义计算，构造函数不调用该 hook。hook 内部可复用 `basemodel_utils.init_model_full_attention` 构建普通 KV 定义；helper 不构成独立启动阶段。 +公共 `prepare_att_resources()` 为 main 创建空管理器，draft 引用 main 的实例。三个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded/tail 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 ## 2. 添加 bloom 模型的示例说明 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index ee8e7d9f24..fbb3776a80 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -124,13 +124,13 @@ Commonly used interfaces that need to be inherited and implemented | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | Read the config.json for initializing the model and perform some key name legalization operations | | def _verify_params(self)： | Validate parameters | -| def _init_att_resource_manager(self): | Create the concrete manager and prepare all attention States: KV meta States, fixed GDN/KDA runtime States and indexer tails; derive static quantization calibration offsets in model helpers from the main and preceding draft configurations | +| `_create_full_att_states(...)` / `_create_bounded_att_states(...)` / `_create_att_tail_states(...)` | Return this model's full / bounded / tail State lists; defaults return empty lists and the shared startup stage registers the results | | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | -After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_att_resources()`, then allocates shared resources. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. +After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_att_resources()`, then calls `finish_state_creation()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. -Concrete models or model families still implement the private `_init_att_resource_manager` hook. Model helpers derive calibration offsets from the full-attention layer definitions of the main and preceding drafts; constructors do not call this hook. The hook can reuse `basemodel_utils.init_model_full_attention` to build one KV scheme; this helper is not a separate startup phase. +The public `prepare_att_resources()` stage creates an empty manager for main and shares it with drafts. The three State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded/tail also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. ## 2. Example of Adding Bloom Model diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index d8df49bd01..e2e0ac176c 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -4,7 +4,7 @@ ## 1. 设计约定 -统一在 `lightllm/common/att_resources/` 下组织可复用的资源实现，由 `AttResourceManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。模型可以直接使用该类，仅在资源创建或操作存在差异时提供子类。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 +统一在 `lightllm/common/att_resources/` 下组织可复用的资源实现，由 `AttResourceManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的三类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： @@ -25,7 +25,7 @@ class AttResourceManager: - 内部调用契约使用 `assert`，例如页类型、物理索引范围、Tensor 形状及共享分配域的一致性；用户配置、外部页面内容及格式兼容性继续显式校验，资源不足等运行失败保留相应异常。GPU 上的 MTP 状态选择保留异步断言，避免引入主机同步。 - 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 -本文区分已确定的职责边界和接口草案。main/draft 的资源由该模型组合的具体管理器统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 +本文区分已确定的职责边界和接口草案。main/draft 的资源由共享的 `AttResourceManager` 统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 ### 1.1 接口命名、排列与注释规范 @@ -72,17 +72,17 @@ flowchart TD Transport[DP/PD 调度] -->|调用具体资源操作| Concrete ``` -图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享该组合的具体管理器，分别绑定各自的计算 State。 +图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享共享的 `AttResourceManager`，分别绑定各自的计算 State。 | 层次 | 组件 | 职责 | | --- | --- | --- | | 请求流程 | `InferenceContext`、`ReqManager` | 创建、暂停、恢复和完成请求；由 ReqManager 管理请求槽位、KV 索引映射及采样参数 | | 资源数据操作 | `AttResourceManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | -| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用具体管理器操作资源数据 | +| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用`AttResourceManager`操作资源数据 | | 资源实现 | 模型所需的 State、存储布局、checkpoint pool | 持有 buffer 和布局，实现读写、复制、重置和快照；IPC 导出范围由管理器统一处理 | | 传输调度 | DP/PD 与多级缓存模块 | 管理任务、队列、通信、事件同步及完成通知 | -`AttResourceManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，具体管理器执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 +`AttResourceManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，`AttResourceManager`执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 代码依赖从调用流程向资源层单向传递。radix cache 依赖管理器的通用状态页操作；管理器与 State 不导入或访问 radix 节点、树和策略。资源实现也不反向导入请求调度模块；传输调度使用数据读写与复制接口。 @@ -105,7 +105,7 @@ AttResourceManager # 公共接口与 allocator 入口 └── 绑定推理视图、组织传输所需状态 # 搬运直接调用 State 方法 ``` -例如纯 full-attention 管理器按角色持有具体 KV 方案；混合模型管理器组合不同算法、存储方案的运行态及 indexer tail。main/draft 分别持有各自的 State，共享 allocator 按分配域创建。公共调用方通过接口使用这些组合，通用聚合操作复用基类实现。 +纯 full-attention 与混合模型都使用 `AttResourceManager`，区别在于模型工厂返回的 State 组合。管理器按角色持有各 State，共享 allocator 按分配域创建；公共调用方使用同一组聚合接口。 以下列出可复用资源组件的结构示例。linear attention 直接使用 `KDAState`、`GDNState` 等具体状态对象，由 `BoundedAttentionState` 约定必需操作，不增加资源包装实例。管理器内部的成员遵守统一命名及唯一性约束，checkpoint pool 仅在相应功能启用时存在： @@ -135,7 +135,7 @@ AttentionTailState # 抽象操作契约，不固定内部存储 └── tail_buffer # 原始 key/gate ring，按层、请求组织 ``` -同一共享分配域中的 allocator 只初始化一次，由资源管理器负责其完整生命周期。若 main/draft 分别使用管理器实例，装配时应注入共享 allocator 引用，并明确唯一的创建和销毁责任，以及调用方解除全部引用并向 allocator 单次归还的约定；若采用一个组合管理器，则由该实现协调内部资源。两种装配方式均不能因模型数增加而重复分配或释放同一共享索引。实例装配形式保留为待定项。 +同一共享分配域中的 allocator 只初始化一次，由 main/draft 共用的唯一 `AttResourceManager` 管理生命周期。调用方解除相关引用后，向对应 allocator 单次归还索引，不因模型数增加而重复分配或释放。 | 资源 | 持有者 | 分配单位 / 生命周期 | | --- | --- | --- | @@ -144,16 +144,16 @@ AttentionTailState # 抽象操作契约，不固定内部存储 | 小页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | | 请求 KV 索引表 | `ReqManager` | 记录逻辑 token 到共享物理索引的映射；负责绑定、更新和清理 | | HOLD 保留物理页及索引描述 | `AttResourceManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | -| KV、量化 scale、packed indexer buffer 与层映射 | 具体管理器组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | -| KDA/GDN 的 Conv/SSM、MTP 游标 | 具体管理器持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | +| KV、量化 scale、packed indexer buffer 与层映射 | `AttResourceManager` 组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | +| KDA/GDN 的 Conv/SSM、MTP 游标 | `AttResourceManager`持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | | 大小页 checkpoint buffer | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID 保存数据；不持有独立空闲列表或长度记录 | -| indexer 原始尾部窗口 | 具体管理器组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | +| indexer 原始尾部窗口 | `AttResourceManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | | 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | | 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复关联资源数据，解除引用后直接归还 allocator | ### 3.1 main/draft 共享 full-attention allocator -调用方直接通过 `AttResourceManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，具体管理器恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： +调用方直接通过 `AttResourceManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，`AttResourceManager`恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： ```text 逻辑 token 位置 t @@ -167,13 +167,13 @@ AttentionTailState # 抽象操作契约，不固定内部存储 共享索引表示存储位置的对应关系，不表示 main/draft 在每个位置都已经写入有效 KV。MTP 预留区和回退后的有效长度仍由请求执行流程维护。 -释放时，请求流程与具体管理器确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 +释放时，请求流程与`AttResourceManager`确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 ### 3.2 main/draft 共享 linear checkpoint allocator main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAState`，draft 使用 `GDNState` 或无需 linear 状态。各状态对象持有自身运行态及启用的大页或小页 pool。对同一前缀的一个大页或小页 checkpoint，调用方直接从管理器持有的对应 allocator 分配一次槽位 ID，再将该 ID 传给管理器的聚合数据操作，供参与的两侧 pool 使用。两侧的状态布局、dtype 和单槽位字节数可以不同，各自的存储必须覆盖共享 allocator 可返回的槽位范围。 -例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。具体管理器负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 +例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。`AttResourceManager`负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 @@ -181,11 +181,11 @@ linear 运行态与两侧 indexer tail 均由 `ReqManager` 提供的 `req_idx` ### 3.3 分配状态与数据状态分离 -allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；具体管理器组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 +allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；`AttResourceManager` 组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 -对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由具体管理器维护。 +对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由`AttResourceManager`维护。 ### 3.4 按具体 attention 实现定义状态 @@ -241,12 +241,12 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 | KDA / GDN | `main_kda_conv_bf16_ssm_fp32_state` / `draft0_gdn_conv_bf16_ssm_fp32_state` | | K-pool C4 indexer tail | `main_kpool_c4_bf16_state` | -普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；K-pool 使用 `kpool_cN_*` 标明压缩比例与存储精度。完整 shape、标定身份等仍保存在 State 配置中。 +普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；K-pool 使用 `kpool_cN_*` 标明压缩比例与存储精度。完整 shape、标定方式等仍保存在 State 配置中。 以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： ```text -Glm5NextAttResourceManager +AttResourceManager ├── full_att_allocator ├── bounded_att_big_page_allocator ├── bounded_att_small_page_allocator @@ -265,7 +265,7 @@ Glm5NextAttResourceManager main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 State 同时命名为 main 和 draft。当前启动流程为各模型分别创建 State 和存储；GLM 的每个模型也分别创建 indexer ring。三种 allocator 的名称不带模型角色，同一分配域只初始化或绑定一个实例。 -具体管理器只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_resource_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `set_att_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 +`AttResourceManager`只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_resource_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `set_att_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 `InferStateInfo` 保存 `att_resource_manager`、`is_draft_model` 和资源角色 `att_resource_role`，不额外绑定 `full_att_state`。各层直接访问管理器上对应方案的 State 成员，再通过局部引用取得层视图或写入 KV；FP8 backend 从具体量化方案的成员读取固定参数。每个 draft 的层由自身编号对应的 State 管理；`att_resource_role` 为 `main` 或 `draft0` 到 `draft35`，不能仅根据 `is_draft_model` 选择资源。是否初始化 full-attention backend 根据管理器中该角色实际启用的资源判断；FFN-only 和 linear-only 模型不创建 full-attention backend，linear backend 由相应模型独立初始化。 @@ -325,9 +325,8 @@ lightllm/common/ └── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 lightllm/models// -├── model.py # 选择、装配本模型的资源管理器 -├── att_resource_config.py # 模型维度与资源组合配置，供模型和缓存服务共用 -└── att_resource_manager.py # 模型资源组合、功能数据操作及具体资源绑定 +├── model.py # 三类 State 创建接口，决定本模型的资源组合 +└── att_resource_config.py # 模型配置与可复用的 State 创建工具 ``` `allocators/` 与三类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 @@ -336,7 +335,7 @@ lightllm/models// `KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_pages`、`restore_bounded_att_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 -`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；CPU meta 定义入口为 `create_bounded_attention_definition(...)`。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 +`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `_create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 `FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；其余复制视图函数返回源、目标视图列表，由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 @@ -344,7 +343,7 @@ CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式 无需模型特有初始化的场景直接使用 `AttResourceManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `set_att_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 -模型特有的 State 创建与存储绑定保留在 `lightllm/models//att_resource_manager.py`。Qwen3Next 创建 `main_gdn_conv_bf16_ssm_fp32_state`；Glm5Next 创建 `main_kda_conv_bf16_ssm_fp32_state` 和 `main_kpool_c4_bf16_state`，压缩比例标记随配置变化。main/draft 的 KV State 按各自方案命名，draft 自行添加需要的 bounded/indexer State。模型管理器继承基类的数据操作和传输布局实现，仅在确有差异时覆盖，不重复实现相同的聚合转发。 +模型特有的 State 组合由 `model.py` 中三类创建接口确定，复用工具放在现有的 `att_resource_config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttResourceManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttResourceManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttResourceManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 resource_id 匹配时可构建临时映射，它不决定模型的资源组织方式。 @@ -362,7 +361,7 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 每个算法的创建入口独立放在 `create_gqa_state.py`、`create_mla_state.py`、`create_kda_state.py`、`create_gdn_state.py` 和 `create_kpool_state.py`，通过各分类包导出。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `set_att_state`。 -GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 由调用方提供真实的 `scales` 和 `calibration_id`，per-head 还需 `q_scales`。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 +GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 `lightllm/utils/att_resource_config.py` 负责把启动参数和标定文件转换为这些显式参数，并调用对应算法入口。full-attention 创建时仍为零普通 token 容量，后续由管理器统一分配；bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 @@ -377,7 +376,7 @@ manager.set_att_state("main", state) 根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttResourceManager.initialize_resources()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 -管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过具体管理器操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 +管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttResourceManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 `types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisSlice`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 @@ -391,9 +390,9 @@ manager.set_att_state("main", state) | `bounded_att_big_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 大页 | | `bounded_att_small_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 小页 | -具体管理器初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 +`AttResourceManager`初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 -allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由模型管理器明确描述。 +allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由`AttResourceManager`明确描述。 大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 @@ -416,7 +415,7 @@ CPU 大小页池通过 `bounded_att/utils.py` 的 `create_page_pool_tensor()` CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排除在大页 allocator 的普通分配范围外；参与操作的 main/draft 大页 pool 分别为这些槽位提供存储。成员名为 `bounded_att_cpu_load_big_page_slot`（从 CPU 加载末尾不足一页的状态时使用）和 `bounded_att_cpu_offload_big_page_slot`（为卸载到 CPU 预留，当前尚无调用点），均属于 bounded-attention 大页池。 -具体管理器声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 +`AttResourceManager`声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 ### 5.3 前缀长度与状态页的关联 @@ -428,7 +427,7 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 ## 6. 对外接口与 attention 资源暴露方式 -模型专属管理器负责资源初始化、内部组合及缓存和传输功能所需的数据操作。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，具体管理器组织内部 State 完成操作。 +`AttResourceManager` 负责共享资源初始化及缓存和传输操作，模型工厂决定 State 组合。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，`AttResourceManager` 组织内部 State 完成操作。 ### 6.1 FullAttentionState：公共接口与不同 storage 的完整处理 @@ -458,7 +457,7 @@ State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get `AttResourceManager.get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_name` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`att_type` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`has_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `att_type=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 -该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_states_created=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。`init_model_full_attention()` 只创建和注册 State，`prepare_att_resources()` 不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_att_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 +该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_states_created=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。三类模型工厂只返回 State 列表，`prepare_att_resources()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_att_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层 tail 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 @@ -638,9 +637,9 @@ manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字 #### 索引、异步执行与生命周期 -构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `initialize_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `initialize_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `initialize_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式和静态标定参数的专属检查在对应 State 构造时完成。 +构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `initialize_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `initialize_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `initialize_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式由配置检查；静态标定参数在 `bind_calibration()` 中检查。 -State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 显式接收必需的 `scales` 和可选的 `q_scales`；meta 阶段保留真实标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 +State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 构造时不接收 `scales` 或 `q_scales`；只能通过 `bind_calibration()` 绑定真实的 CPU 标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 `logical_layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 @@ -662,7 +661,7 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 | `write_to_pd_buffer(..., layout)` | 当前 State → `pd_page_index` 指定的 PD 页；按 layout 写入指定区域 | | `read_from_pd_buffer(..., layout)` | `pd_page_index` 指定的 PD 页 → 当前 State；按 layout 恢复对应区域 | -`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，具体管理器按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 +`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，`AttResourceManager`按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttResourceManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 @@ -685,7 +684,8 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 | `AttResourceManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | | `set_att_state(model_name, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | | `finish_state_creation()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_states_created = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | -| `initialize_resources(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：标记 State 创建完成，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | +| `load_full_att_calibration(args)` | 在 `finish_state_creation()` 后、存储分配前加载静态 FP8 标定；文件只读一次，逐 State 绑定物理层区间及 TP 参数；没有静态 FP8 State 时直接返回 | +| `initialize_resources(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：在完成创建标记和标定绑定后，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | | `get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_states_created` 并使用 `lru_cache` 缓存 | | `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | | `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | @@ -894,7 +894,7 @@ main/draft 共享管理器，因此必须显式指定 `role="main"` 或具体 dr #### 多个 State 的参与规则 -具体管理器在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 +`AttResourceManager`在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 例如，一个管理器同时管理两个 KDAState 和一个 GDNState，大页槽位 17 对应： @@ -952,7 +952,7 @@ CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能 `src_manager` 明确来源 DP/TP 存储实例，也可以是传输进程中已绑定共享存储的管理器。源和目标 token 索引等长、逐项对应；`src_req_indices`、`dst_req_indices` 仅接受 `List[int]`，成对提供且数量相等，用来迁移 prefill 基准 bounded 运行态；KDA/GDN 源游标必须为 0。有 tail 时直接清零目标请求的 tail，不传递额外长度。token 索引与请求槽位是不同的寻址维度，无需等长；前缀对应关系及可重置边界由调用方保证。 -只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，具体管理器将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 +只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，`AttResourceManager` 将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 **PD 页读写** @@ -969,7 +969,7 @@ CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引 radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存和恢复状态页时调用上述通用接口，前缀与槽位的对应关系由 radix 调用流程保证；槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 -CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，具体管理器负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 +CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，`AttResourceManager`负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 ### 6.3 模型推理使用具体资源视图 @@ -987,7 +987,7 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S 单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttResourceManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 -具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttResourceManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该模型管理器的具体实现提供。 +具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttResourceManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该`AttResourceManager`的具体实现提供。 跨进程共享通过管理器的 `create_pd_ipc_copy()` / `create_dp_ipc_copy()` 保留 State Tensor 引用并由 IPC 传递句柄，不调用 `initialize_storage()` 重新分配 KV；两者均排除 allocator，PD 副本排除全部大小页池，DP 副本保留共享 CPU 大小页池，不导出整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 @@ -1131,7 +1131,7 @@ D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸， 字段列表按实际物理表示组织。例如 K、V、scale 可分别描述；packed 数据也可整体描述为一个字节区域。别名区域只搬运一次。Conv 的 Q/K 和 V、SSM 可能采用不同 head 映射，应分别描述，不能把一个资源的 `head_start` 套用于所有字段。数据类型、量化组粒度及索引单位由对应格式约束。 -具体管理器从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 +`AttResourceManager`从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 FullAttentionState 的 token 页操作要求 `tokens_per_page` 为正整数。CPU 参数 `cpu_page_buffers` 为完整二维页池；PD 参数 `pd_page_buffers` 同样为页集合，`pd_page_index` 指定物理页。传入完整页集合，不提前切到单页或 State 区域。对于常规 K/V，一个外部地址可以写成： @@ -1212,7 +1212,7 @@ CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索 布局由部署配置和 State 生成；搬运时校验资源名称、字段覆盖、dtype、字节对齐、页内地址边界及层/轴映射范围，TP 切片指定分片范围与副本写入者。各资源的实际写入区域不得相互覆盖；packed 别名按实际物理范围去重。涉及多个布局的描述以实际访问地址为准，不能仅比较 shape。 -数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。具体管理器校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 +数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。`AttResourceManager`校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 ### 6.6 子类实现与直接调用示例 @@ -1259,7 +1259,7 @@ state.read_from_cpu( # 后续计算与写入建立 stream/event 依赖。 ``` -DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由具体管理器组织： +DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由`AttResourceManager` 组织： ```python dst_state.copy_tokens_from_other_dp( @@ -1292,7 +1292,7 @@ d_state.read_from_pd_buffer( ### 6.7 BoundedAttentionState：有界状态的统一操作接口 -`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。模型管理器直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 +`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。`AttResourceManager`直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 ```text BoundedAttentionState # 抽象操作契约 @@ -1473,11 +1473,11 @@ CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字 当前 DP/PD 导出 prefill 基准状态，调用方保证 KDA/GDN 源游标为 0，不传 MTP 候选区，也不支持 MTP 提交后导出。目标可按自己的 MTP 步数初始化工作区；DP 两端配置及 TP 分片一致，按各自 MTP 步数计算基准槽位。PD 则通过全局布局与 TP 切片处理两端分片差异。不能因为 KDA/GDN 的 Conv/SSM 形状相同就跨算法恢复。 -多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；具体管理器完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 +多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；`AttResourceManager`完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 #### 依赖、同步与使用示例 -资源方法接收请求槽位、checkpoint 槽位和明确的资源引用，只有 tail 无快照恢复额外接收前缀长度，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；模型管理器单独协调 tail 在恢复边界上的重置或迁移。 +资源方法接收请求槽位、checkpoint 槽位和明确的资源引用，只有 tail 无快照恢复额外接收前缀长度，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttResourceManager`单独协调 tail 在恢复边界上的重置或迁移。 GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 @@ -1609,23 +1609,34 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 | --- | --- | | `get_req_manager_max_seq_len()` | 返回请求映射所需长度容量，取本模型 `batch_max_tokens`、`max_seq_length` 与 0 的最大值；未配置的值按 0 处理，不表示模型最大上下文能力 | | `bind_req_manager(req_manager)` | 绑定调用方创建的共享实例，不判断 main/draft，也不通过 `main_model` 获取资源 | -| `prepare_att_resources()` | 在本模型 KV 内存分区内调用 `_init_att_resource_manager` 准备并注册自身全部 State | +| `prepare_att_resources()` | main 直接创建空的 `AttResourceManager`，draft 引用 main 的实例；在 KV 内存分区内调用三类 State 工厂并统一注册返回的列表 | | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 -2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。main 的资源 hook 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，其资源 hook 通过 `set_att_state` 只添加本模型的 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。静态 FP8 标定偏移由模型侧配置 helper 计算：main 从 0 开始，draft 累加 main 及前序 draft 的实际 full-attention 层数；混合模型排除 bounded-attention 层，原生 MTP 使用自己的层定义，无 attention 的 draft 不占标定行。公共准备接口及资源 hook 均不接收标定偏移，也不返回层数。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 -3. backend 检查 main 和所有 draft 引用同一个 attention 资源管理器，随后调用 `manager.initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)`。管理器先标记 State 创建完成；每个模型的成员已经直接保存在共享实例中，无需合并管理器或 State。不同 draft 即使方案相同，也可拥有不同的配置、存储和游标。 -4. `finish_state_creation()` 完成后，通过管理器接口计算容量并初始化实际存储。FullAttentionState 原地分配 buffer；不同 State 可以具有不同物理布局，但共享 allocator 的索引域及容量；同一分配域只创建一个 allocator。实际存储分配完成后再配置 CPU/PD 布局。 +2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `_create_full_att_states()`、`_create_bounded_att_states()`、`_create_att_tail_states()`，再通过 `set_att_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 +3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_creation()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 +4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 仍必须有 K/V 和 Q 标定。随后 backend 调用 `initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 -BaseModel 保留 `_init_att_resource_manager()` 等现有私有扩展点，由单模型公共阶段接口调用，构造函数不调用运行时初始化。main 在资源 hook 中选择具体管理器，各模型分别在自己的 hook 中准备所需 attention State，子类可以继承适用的模型族实现。原生 MTP 如需复用 main 的资源配置，也在自身资源准备阶段读取；ReqManager 的创建与共享关系由 backend 组织，模型只提供容量需求和公共绑定接口。 +BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `resource_id`。 + +| 创建接口 | 返回值与职责 | +| --- | --- | +| `_create_full_att_states(...)` | 返回 `list[FullAttentionState]`；只准备布局，不加载静态标定或分配 KV 容量；可返回多种量化或存储方案 | +| `_create_bounded_att_states(...)` | 返回 `list[BoundedAttentionState]`；按请求容量创建运行态，大小页延后分配；CPU 查询使用 `device="meta"` | +| `_create_att_tail_states(...)` | 返回 `list[AttentionTailState]`；按请求容量创建本模型 tail，CPU 联合页不调用此接口 | + +`model_logical_layer_ids` 表示当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；main/draft 保留现有全局编号机制。具体模型自行选择各类 State 的参与层和共享分组，该参数不代表物理存储层数，也不要求每个 State 使用全部层。 + +三个接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_att_resources()` 为 main 直接创建空的 `AttResourceManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttResourceManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 + 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -单一 KV 方案的模型可在资源初始化入口内部复用 `basemodel_utils.init_model_full_attention(...)`：读取当前模型配置和计算层，通过 `get_model_calibration_storage_layer_start(model)` 解析标定起点，创建 meta State，再通过 `set_att_state(model.att_resource_role, state)` 绑定到共享管理器。该 helper 仅用于复用模型侧的构建逻辑，不作为独立启动阶段；管理器基类不提供单类 State 的准备接口。head 信息直接读取 State config；meta 定义不能通过绑定 buffer 视图改名，否则会妨碍后续实际存储分配。存在多种 KV 方案的模型应在自己的 hook 中组织各定义及其标定行映射。模型和 helper 不读取 main 或其他 draft 管理器来推导偏移。 +单一 KV 方案的 `_create_full_att_states()` 可复用 `basemodel_utils.create_full_attention_definition(...)`，将返回的单个定义包装为列表；不存在 full 层时返回空列表。该 helper 不访问管理器，也不构成独立启动阶段。bounded 模型工厂直接调用对应算法 create，不再先生成 meta State 再重建实际 State。GLM main/draft 的 tail 工厂仅从自身配置读取 K-pool 参数，不依赖 main 的 `linear_config`。 -`AttResourceManager.initialize_resources()` 由基类提供通用实现，保留三个阶段：标记 State 创建完成、估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 +`AttResourceManager.initialize_resources()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 自动容量规划使用所有模型实际加载后的剩余显存，不再按 main 权重比例预估 draft 占用。静态量化标定须在规划前加载和校验；KV、动态 scale、packed padding 及 HOLD 页计入预算，静态 KV/Q 标定参数占比较小，不在容量规划中单独扣减。按完整物理页选择可容纳的最大 token 容量，各 rank 再统一取最小容量。共享 allocator 的不同 State 存储分别计算，不能因共享索引而漏算 draft buffer。 @@ -1652,7 +1663,7 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 `basemodel_utils.py` 提供 `get_resource_config(model_config, *, model_class=None)` 和 `get_full_attention_options(config, args, *, model_class=None)`，统一处理字段别名、text config 和默认 KV 维度。它们是普通工具函数，BaseModel 不再声明对应的 `_resource_config`、`_full_attention_options` 方法；具体模型仅在存在专有规则时提供可选 hook，并调用工具函数取得默认配置。传入 `model_class` 时，工具函数选择该模型的专有实现，否则使用通用逻辑。 -同一文件中的 `create_full_attention_definition(model_class, ...)` 和 `create_bounded_attention_definition(model_class, ...)` 只返回包含完整 config 的 meta State，不再单独返回全局 head 数。full-attention 构建使用上述配置工具，直接构造 `FullAttentionConfig` 并调用算法级 `create_gqa_state` / `create_mla_state`，不再经由额外的 State 构建转发函数；linear 构建调用具体模型可选的 `create_bounded_attention_definition` 实现，纯 full-attention 模型由工具返回空定义，基类不保留空实现。CPU 服务读取 main 和各 draft 自己的配置，通过相同的工具入口生成联合页；原生 MTP 和 DFlash 等配置转换由对应模型实现。CPU 页包含所有参与的 linear State，并为各角色安排对应区域。静态量化标定按明确的物理行起点和本模型 full-attention 物理层数选取，不再根据可选起点推断 main/draft 合并层数；CPU 服务和推理端复用相同的布局构建流程。`att_resource_config.py` 仅保留标定加载和 CPU 服务布局构建。 +同一文件中的 `create_full_attention_definition(model_class, ...)` 根据配置直接调用算法级 `create_gqa_state` / `create_mla_state`，供模型 full 工厂复用。CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例或运行推理初始化。返回的 State 列表统一注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可包含同一模型的多个 full/bounded State；静态标定不参与布局计算，CPU 服务无需打开标定文件。原生 MTP、DFlash 等配置转换由对应模型实现。`att_resource_config.py` 仅保留 CPU 服务布局构建；标定加载由管理器集中组织，具体格式解释由静态 FP8 State 完成。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 @@ -1660,7 +1671,7 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)` 和 `reset_attention_tail_state(req_indices)`，分别初始化 bounded 与 tail 运行态。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 -调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_pages`。具体管理器将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 +调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_pages`。`AttResourceManager` 将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 radix 沿用原有判断：大页节点具有状态，小页节点通过 `small_page_buffer_idx is not None` 判断是否可恢复；图片边界仍通过 `clamp_mm_cache_len` 独立检查。匹配得到节点引用后，调用方直接调用 `restore_bounded_att_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_attention_tail_state(req_indices)` 清零 tail，保留已恢复的 linear 状态。调用方还需确认历史 KV、tail 和 bounded 快照均对应恢复前缀，并保证数据就绪。 @@ -1678,7 +1689,7 @@ MTP 接受位置由模型执行流程选择，供后续计算使用；当前大 ## 8. CPU、DP 与 PD 传输设计 -### 8.1 具体 State 实现传输，模型管理器组织组合 +### 8.1 具体 State 实现传输，`AttResourceManager` 组织组合 CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合缓存辅助文件： @@ -1689,7 +1700,7 @@ CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合 命中边界、容量准入和 CPU 页引用释放仍由外层缓存调度负责。以上入口不把 `InferReq` 传入 `AttResourceManager`；资源层只接收显式索引与槽位，并组合各 State 的数据操作。 -外部通过 6.2 节的五个管理器抽象方法执行 CPU/DP/PD 操作。具体管理器将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 +外部通过 6.2 节的管理器接口执行 CPU/DP/PD 操作。`AttResourceManager` 将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 - full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 - `KDAState`、`GDNState` 实现 `BoundedAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移 prefill 基准请求状态；管理器组织参与对象及位置映射。 diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/att_resources/att_resource_manager.py index 12b274d055..c72822694e 100644 --- a/lightllm/common/att_resources/att_resource_manager.py +++ b/lightllm/common/att_resources/att_resource_manager.py @@ -35,7 +35,7 @@ class AttResourceManager: """管理模型组合的注意力资源，提供缓存复用及 CPU/DP/PD 传输所需的聚合操作。 - 默认通过聚合查询操作各类 State；模型子类仅在资源创建或操作有差异时扩展。 + 通过聚合查询操作各类 State；模型通过三类创建接口决定资源组合，main/draft 共享本管理器。 未支持的传输路径必须明确报错。 """ @@ -112,6 +112,52 @@ def get_att_state( ) return matches[0] + def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: + """全部 main/draft State 就绪后统一读取静态标定，并交给具体 State 绑定。 + + 文件沿用物理行顺序：main、draft0..draft35；同角色有多个 State 时按最小逻辑层号 + 排序，State 内保留 logical_layer_groups 顺序。非静态 FP8 的 full State 同样占用 + 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 + """ + import json + from .full_att.static_fp8_state import StaticFP8FullAttentionState + + full_states = self.get_full_attention_states() + calibration_states = [state for state in full_states if isinstance(state, StaticFP8FullAttentionState)] + if not calibration_states: + return + + # 阶段一：只读取一次文件；无文件时由 per_tensor State 绑定单位 scale。 + calibration = None + path = args.kv_quant_calibration_config_path + if path is not None: + with open(path) as file: + calibration = json.load(file) + if any(state.config.scale_mode == "per_tensor" for state in calibration_states): + from lightllm.utils.config_utils import get_model_architectures + + assert calibration["architectures"] == get_model_architectures(args.model_dir), ( + "calibration architecture does not match the model" + ) + + # 阶段二：显式使用模型角色顺序，避免注册先后及 draft10 的字符串排序影响标定行。 + storage_layer_start = 0 + for role in MODEL_ROLES: + states = sorted( + self.get_full_attention_states(role), + key=lambda state: ( + min(layer for group in state.config.logical_layer_groups for layer in group), state.resource_id, + ), + ) + for state in states: + if isinstance(state, StaticFP8FullAttentionState): + state.bind_calibration( + calibration=calibration, + storage_layer_start=storage_layer_start, + ) + storage_layer_start += len(state.config.logical_layer_groups) + logger.info("Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales") + def initialize_resources( self, args: StartArgs | EasyDict, @@ -123,14 +169,12 @@ def initialize_resources( """全部模型完成资源准备后，规划共享容量并初始化缓存存储。 调用前，所有模型的权重和绑定请求容量的运行态已创建，full-attention - 仅持有 meta 定义，大小页池尚未分配。调用方在本方法返回后， + 已完成创建标记和所需的静态标定绑定，大小页池尚未分配。调用方在本方法返回后， 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 """ - # 阶段一：标记全部 State 对象创建完成，后续通过聚合接口访问。 - # main、draft0、draft1 等模型已分别将自己的 State 添加到当前管理器。 - self.finish_state_creation() + assert self._states_created, "State creation must finish before resource initialization" - # 阶段二：规划共享 token 容量，并向下对齐到完整 KV 页。 + # 阶段一：规划共享 token 容量，并向下对齐到完整 KV 页。 # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 # PD 模式还需为传输 buffer 预留空间，布局按全部 full-attention State 计算。 assert self.get_full_attention_states(), "the shared token index domain requires a full-attention State" @@ -164,7 +208,7 @@ def initialize_resources( token_capacity = token_capacity // args.page_size * args.page_size assert token_capacity > 0, "attention capacity must contain at least one complete page" - # 阶段三：管理器为已创建的 State 分配实际存储与共享 allocator。 + # 阶段二：管理器为已创建的 State 分配实际存储与共享 allocator。 # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num self.initialize_storage( diff --git a/lightllm/common/att_resources/full_att/create_gqa_state.py b/lightllm/common/att_resources/full_att/create_gqa_state.py index c31a95c63b..c176ae1daa 100644 --- a/lightllm/common/att_resources/full_att/create_gqa_state.py +++ b/lightllm/common/att_resources/full_att/create_gqa_state.py @@ -19,15 +19,12 @@ def create_gqa_state( page_size: int = 1, device: str | torch.device = "cuda", group_size: int = 8, - scales: torch.Tensor | None = None, - q_scales: torch.Tensor | None = None, - calibration_id: str = "", ) -> GQAFullAttentionState | GroupQuantizedFullAttentionState | StaticFP8FullAttentionState: """按编码参数创建 GQA（含 MHA/MQA）State，普通 token 容量由后续流程分配。 config 描述逻辑 KV 维度和 TP 信息；kv_type 沿用启动参数中的编码名称。 - group_size 仅用于 INT4/INT8；scales、q_scales 和 calibration_id 仅用于静态 FP8。 - 静态标定数据由调用方加载，本入口不读取启动参数、配置文件或绑定模型角色。 + group_size 仅用于 INT4/INT8。静态 FP8 只创建布局，标定参数必须通过后续 + bind_calibration 绑定；本入口不读取配置文件或绑定模型角色。 """ if kv_type == "None": return GQAFullAttentionState(layout_config=config, page_size=page_size, device=device) @@ -45,8 +42,6 @@ def create_gqa_state( ) return GroupQuantizedFullAttentionState(layout_config=quant_config, page_size=page_size, device=device) assert kv_type in ("fp8kv_spt", "fp8kv_sph"), f"unsupported GQA encoding: {kv_type}" - assert scales is not None, "static FP8 requires calibration scales" - assert kv_type != "fp8kv_sph" or q_scales is not None, "per-head FP8 requires Q calibration scales" quant_config = StaticFP8AttentionConfig( logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, @@ -56,12 +51,9 @@ def create_gqa_state( tp_world_size=config.tp_world_size, replicated=config.replicated, scale_mode="per_head" if kv_type == "fp8kv_sph" else "per_tensor", - calibration_id=calibration_id, ) return StaticFP8FullAttentionState( layout_config=quant_config, page_size=page_size, device=device, - scales=scales, - q_scales=q_scales, ) diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index 8762d7a6ca..e9855c9a14 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -1,6 +1,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, List +from collections.abc import Mapping from dataclasses import dataclass from math import prod @@ -25,15 +26,11 @@ class StaticFP8AttentionConfig(FullAttentionConfig): # per_tensor：每层 K/V 各一个固定 scale；per_head：每层、每个本地 K/V head 各一个。 # 标定参数不随 token 增长，也不作为动态 KV 字段写入 CPU/PD 页。 scale_mode: str = "per_tensor" - # 标定数据的身份标记，用于 DP 复制时检查编码是否兼容；实际 scale Tensor 由 State 持有。 - calibration_id: str = "" def __post_init__(self) -> None: - """创建配置时校验通用 TP 约束、标定方式及身份。""" + """创建配置时校验通用 TP 约束及标定方式。""" super().__post_init__() - assert self.scale_mode in ("per_tensor", "per_head") and self.calibration_id, ( - "static FP8 requires a scale mode and calibration identity" - ) + assert self.scale_mode in ("per_tensor", "per_head"), "invalid static FP8 scale mode" class StaticFP8FullAttentionState(FullAttentionState): @@ -50,24 +47,10 @@ def __init__( page_size: int = 1, *, device: str | torch.device = "cuda", - scales: torch.Tensor, - q_scales: torch.Tensor | None = None, ) -> None: - """按静态 FP8 配置准备 KV 存储，并保留传入的标定参数。""" - # 标定 Tensor 由 State 持有，在接收构造参数时校验，不重复放入 config。 - heads = 1 if layout_config.scale_mode == "per_tensor" else layout_config.head_num - assert scales is not None and scales.shape == (len(layout_config.logical_layer_groups), 2 * heads), ( - "calibration scales do not match KV layers and heads" - ) - assert scales.dtype == torch.float32 and bool((torch.isfinite(input=scales) & (scales > 0)).all()), ( - "calibration scales must be finite positive float32 values" - ) - if q_scales is not None: - assert q_scales.shape == (len(layout_config.logical_layer_groups), layout_config.head_num) and ( - q_scales.dtype == torch.float32 - ), "Q calibration does not match layers and heads" - self.scales = scales - self.q_scales = q_scales + """只准备静态 FP8 布局；标定参数由后续 bind_calibration 统一绑定。""" + self.scales: torch.Tensor | None = None + self.q_scales: torch.Tensor | None = None self.cpu_scales = None # 实际存储初始化时准备计算 backend 使用的 CPU 标定值。 self.config = layout_config self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 @@ -84,17 +67,85 @@ def __init__( self.v_storage: torch.Tensor | None = None self.storage_dtype = torch.uint8 + def bind_calibration( + self, + calibration: Mapping[str, object] | None, + *, + storage_layer_start: int, + ) -> None: + """在分配 KV 前绑定标定表中本 State 的物理层区间，并切出当前 TP 的 head。 + + 行起点由管理器按 main、draft0..draft35 的完整 State 集合确定，不是逻辑层号。 + 同组逻辑层共用一行；State 内的行顺序与 logical_layer_groups 的物理层顺序一致。 + 无文件时仅 per_tensor 可使用单位 scale；per_head 必须包含 K/V 和 Q 标定。 + """ + assert self.k_storage is None and self.scales is None, "calibration must be bound once before storage allocation" + config = self.config + storage_layer_count = len(config.logical_layer_groups) + if calibration is None: + assert config.scale_mode == "per_tensor", "per-head FP8 requires a calibration file with Q scales" + scales = torch.ones(size=(storage_layer_count, 2), dtype=torch.float32) + q_scales = None + else: + # 阶段一：检查文件编码和物理行覆盖范围。 + calibrated_layer_count = calibration["num_layers"] + storage_layer_end = storage_layer_start + storage_layer_count + assert 0 <= storage_layer_start < storage_layer_end <= calibrated_layer_count, ( + "calibration does not cover this State's physical layers" + ) + assert calibration.get("quant_type", config.scale_mode) == config.scale_mode, ( + "calibration quantization mode does not match the State" + ) + limits = torch.finfo(torch.float8_e4m3fn) + assert calibration.get("qmin", limits.min) == limits.min and calibration.get("qmax", limits.max) == limits.max, ( + "calibration range does not match float8_e4m3fn" + ) + + def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]) -> torch.Tensor: + tensor = torch.tensor(data=values["scales"], dtype=torch.float32) + assert tuple(values["scales_shape"]) == tensor.shape == shape, "invalid calibration scale shape" + assert bool((torch.isfinite(tensor) & (tensor > 0)).all()), "calibration scales must be finite and positive" + return tensor[storage_layer_start:storage_layer_end] + + def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: + heads = config.head_num + assert calibrated_heads > 0 and (heads * config.tp_world_size) % calibrated_heads == 0, ( + "calibration heads do not cover the model's KV heads" + ) + expanded = tensor.repeat_interleave(repeats=heads * config.tp_world_size // calibrated_heads, dim=-1) + return expanded[:, config.tp_rank * heads : (config.tp_rank + 1) * heads].contiguous() + + # 阶段二：先按物理层取行，再分别切分 K、V、Q 的 head；标定 Tensor 暂留 CPU。 + if config.scale_mode == "per_head": + calibrated_heads = calibration["num_head"] + all_scales = scale_tensor(calibration, (calibrated_layer_count, 2 * calibrated_heads)) + k, v = all_scales.split(split_size=calibrated_heads, dim=-1) + scales = torch.cat(tensors=(shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) + q_config = calibration.get("q_calibration") + assert q_config, "per-head FP8 requires Q calibration" + q_scales = shard( + scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), q_config["num_head"], + ) + else: + scales = scale_tensor(calibration, (calibrated_layer_count, 2)) + q_scales = None + + # 阶段三：绑定标定数值；后续 initialize_storage 将其迁移到实际设备。 + self.scales = scales + self.q_scales = q_scales + def get_state_type(self) -> str: """返回包含 per-tensor 或 per-head 标定方式的 FP8 方案标记。""" return "gqa_fp8_static_" + self.config.scale_mode def initialize_storage(self, token_capacity: int) -> None: """绑定 KV 与真实标定参数，并预备 FlashInfer 使用的 CPU scale。""" + assert self.scales is not None, "static FP8 calibration must be bound before allocation" self.capacity = token_capacity shape = self._storage_shape() self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) - # meta 准备阶段保留真实标定值，实际分配时再迁移到 KV 所在设备。 + # 统一标定阶段保留 CPU 数值，实际分配时再迁移到 KV 所在设备。 if self.k_storage.device.type != "meta": self.scales = self.scales.to(device=self.k_storage.device) if self.q_scales is not None: @@ -235,14 +286,7 @@ def copy_tokens_from_other_dp( dst_token_indices: torch.Tensor, remote_dp_full_att_state: StaticFP8FullAttentionState, ) -> None: - """先校验标定方式与标定身份，再复制兼容的 FP8 KV 编码。""" - assert ( - remote_dp_full_att_state.config.scale_mode, - remote_dp_full_att_state.config.calibration_id, - ) == ( - self.config.scale_mode, - self.config.calibration_id, - ), "incompatible static FP8 calibration" + """复制相同配置的 FP8 KV 编码；各 DP 使用相同的静态标定参数。""" assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( "DP copy requires the same encoding and equivalent TP state configurations" ) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index c7e8228619..015cc71abe 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -8,12 +8,13 @@ import torch import torch.nn.functional as F import triton -from typing import final, List +from typing import final, List, TYPE_CHECKING from tqdm import tqdm from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.req_manager import ReqManager +from lightllm.common.att_resources import AttResourceManager from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, @@ -46,6 +47,13 @@ from .attention import get_prefill_att_backend_class, get_decode_att_backend_class from .attention import BaseAttBackend +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.att_resources.bounded_att import BoundedAttentionState + from lightllm.common.att_resources.att_tail import AttentionTailState + logger = init_logger(__name__) torch.backends.cudnn.enabled = True @@ -136,14 +144,37 @@ def bind_req_manager(self, req_manager: ReqManager) -> None: self.req_manager = req_manager def prepare_att_resources(self) -> None: - """Prepare this model's attention States from its own configuration.""" + """main 创建共享管理器，draft 引用同一实例，再准备当前模型自己的 State。""" assert getattr(self, "att_resource_manager", None) is None, "resources already initialized" if self.is_mtp_draft_model: self.att_resource_manager = self.main_model.att_resource_manager assert not self.att_resource_manager._states_created, "State creation is already finished" assert self.max_req_num + 1 == self.att_resource_manager._request_capacity, "models must share the request domain" with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - self._init_att_resource_manager() + if not self.is_mtp_draft_model: + self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + # 无 attention 的 draft 只共享管理器，不创建自身 State。 + if self.is_mtp_draft_model and self.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + return + + # 当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；各工厂自行选择参与层及共享分组。 + logical_layer_ids = tuple(layer.layer_num_ for layer in self.layers_infer) + full_states = self._create_full_att_states( + model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, + tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, + ) + bounded_states = self._create_bounded_att_states( + model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, + request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, + device="cuda", + ) + tail_states = self._create_att_tail_states( + model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, + request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, + device="cuda", + ) + for state in (*full_states, *bounded_states, *tail_states): + self.att_resource_manager.set_att_state(self.att_resource_role, state) def init_runtime(self) -> None: """Initialize this model's backends, warmup and graphs after shared allocation.""" @@ -221,9 +252,41 @@ def load_weights(self, weight_dict: dict): [weight.verify_load() for weight in self.trans_layers_weight] return - def _init_att_resource_manager(self) -> None: - """Prepare all owned attention States using model-specific configuration.""" - raise NotImplementedError + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """返回本模型的 full-attention 定义列表；默认没有此类资源，不分配 KV 容量。 + + model_logical_layer_ids 是当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列； + 具体模型自行选择参与层并定义 logical_layer_groups，不要求使用全部层。 + """ + return [] + + @classmethod + def _create_bounded_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["BoundedAttentionState"]: + """返回本模型的 bounded 运行态列表；meta 用于布局查询，大小页后续分配。 + + model_logical_layer_ids 描述当前模型拥有的全部全局逻辑层，具体模型自行选择参与层。 + """ + return [] + + @classmethod + def _create_att_tail_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["AttentionTailState"]: + """返回本模型的 tail 运行态列表；默认没有此类资源。 + + model_logical_layer_ids 描述当前模型拥有的全部全局逻辑层，具体模型自行选择参与层。 + """ + return [] def _check_mem_size(self): self.max_total_token_num = self.att_resource_manager.full_att_allocator.size diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index 71636e1d2c..46af0f0822 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -3,8 +3,6 @@ from __future__ import annotations from typing import TYPE_CHECKING -import hashlib -import json from lightllm.common.build_utils import repair_config if TYPE_CHECKING: @@ -68,13 +66,13 @@ def create_full_attention_definition( tp_rank: int = 0, tp_world_size: int = 1, model_logical_layer_ids: Sequence[int] | None = None, - calibration_storage_layer_start: int = 0, ) -> FullAttentionState | None: - """按模型配置直接创建算法 State；静态 FP8 使用当前模型的物理层标定区间。""" + """按模型配置创建算法 State；静态 FP8 标定由全部 State 就绪后的统一阶段绑定。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state, create_mla_state - from lightllm.utils.att_resource_config import load_static_fp8_calibration from lightllm.utils.torch_dtype_utils import get_torch_dtype + if model_class.is_mtp_draft_model and args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + return None config = get_resource_config(model_config, model_class=model_class) options = get_full_attention_options(config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids) if not options["logical_layer_groups"]: @@ -96,93 +94,6 @@ def create_full_attention_definition( config=layout_config, kv_type=args.llm_kv_type, packed=packed, page_size=args.page_size, device="meta", ) - scales = q_scales = None - calibration_id = "" - if args.llm_kv_type in ("fp8kv_sph", "fp8kv_spt"): - calibration, scales, q_scales = load_static_fp8_calibration( - args=args, - per_head=args.llm_kv_type == "fp8kv_sph", - storage_layer_count=len(layout_config.logical_layer_groups), - heads=layout_config.head_num, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - storage_layer_start=calibration_storage_layer_start, - ) - calibration_id = hashlib.sha256(json.dumps(calibration, sort_keys=True).encode()).hexdigest() return create_gqa_state( config=layout_config, kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", - scales=scales, q_scales=q_scales, calibration_id=calibration_id, - ) - - -def create_bounded_attention_definition( - model_class: type[TpPartBaseModel], - model_config: dict[str, object], - args: StartArgs | EasyDict, - *, - model_logical_layer_ids: Sequence[int], - tp_rank: int = 0, - tp_world_size: int = 1, -): - """Use the model's optional bounded-attention definition, including inherited implementations.""" - create_definition = getattr(model_class, "create_bounded_attention_definition", None) - if create_definition is None: - return None - return create_definition( - model_config, args, model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size - ) - - -def get_model_calibration_storage_layer_start(model: TpPartBaseModel) -> int: - """按前序模型的物理层数累计标定表行偏移，不创建 State，也不使用逻辑层号作为行号。""" - if not model.is_mtp_draft_model: - return 0 - storage_layer_start = 0 - for previous in (model.main_model, *model.mtp_previous_draft_models): - if previous.is_mtp_draft_model and previous.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): - continue - config = get_resource_config(previous.config, model_class=type(previous)) - options = get_full_attention_options( - config, - previous.args, - model_class=type(previous), - model_logical_layer_ids=tuple(layer.layer_num_ for layer in previous.layers_infer), - ) - storage_layer_start += len(options["logical_layer_groups"]) - return storage_layer_start - - -def init_model_full_attention(model: TpPartBaseModel) -> None: - """按当前模型的配置和层号创建 full-attention meta State，直接挂到共享管理器。""" - if model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): - return - definition = create_full_attention_definition( - type(model), - model.config, - model.args, - tp_rank=model.tp_rank_, - tp_world_size=model.tp_world_size_, - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - calibration_storage_layer_start=get_model_calibration_storage_layer_start(model), - ) - if definition is None: - return - model.att_resource_manager.set_att_state(model.att_resource_role, definition) - - -def init_model_bounded_attention(model: TpPartBaseModel) -> None: - """按当前模型的配置创建请求运行态，并直接挂到共享管理器；大小页稍后分配。""" - definition = create_bounded_attention_definition( - type(model), model.config, model.args, - 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_, - ) - if definition is None: - return - from lightllm.common.att_resources.bounded_att import GDNState, KDAState, create_gdn_state, create_kda_state - - create_state = {GDNState: create_gdn_state, KDAState: create_kda_state}[type(definition)] - state = create_state( - config=definition.config, request_capacity=model.max_req_num + 1, mtp_step=model.args.mtp_step, device="cuda" ) - model.att_resource_manager.set_att_state(model.att_resource_role, state) diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index d44e6f63da..0685e0f9b5 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer @@ -7,6 +8,12 @@ from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class BloomTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = BloomPreAndPostLayerWeight @@ -24,13 +31,19 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention - - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _init_config(self): super()._init_config() diff --git a/lightllm/models/glm5_next/att_resource_config.py b/lightllm/models/glm5_next/att_resource_config.py index 571ccdd084..eec434053b 100644 --- a/lightllm/models/glm5_next/att_resource_config.py +++ b/lightllm/models/glm5_next/att_resource_config.py @@ -9,6 +9,8 @@ from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.att_resources.att_tail import AttentionTailState from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs @@ -63,3 +65,33 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) + + +def create_model_att_tail_states( + model_class: type[TpPartBaseModel], model_config: dict[str, object], args: StartArgs | EasyDict, *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int, + tp_rank: int, tp_world_size: int, device: str, +) -> list[AttentionTailState]: + """由模型配置创建 K-pool tail 列表，不读取模型实例或注册管理器。""" + from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + from lightllm.common.basemodel.basemodel_utils import get_full_attention_options, get_resource_config + + config = get_resource_config(model_config, model_class=model_class) + options = get_full_attention_options( + config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids, + ) + if not options["logical_layer_groups"]: + return [] + assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" + state = create_kpool_state( + config=KPoolStateConfig( + logical_layer_ids=tuple( + logical_layer_id for logical_layer_group in options["logical_layer_groups"] + for logical_layer_id in logical_layer_group + ), + pool_size=config["index_kpool"], head_dim=config["index_head_dim"], + dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ) + return [state] diff --git a/lightllm/models/glm5_next/att_resource_manager.py b/lightllm/models/glm5_next/att_resource_manager.py deleted file mode 100644 index e9333bf7b0..0000000000 --- a/lightllm/models/glm5_next/att_resource_manager.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Glm5Next attention resources with explicit model-owned State objects.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.utils import dtype_name -from lightllm.common.att_resources.allocators import StatePageAllocator -from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig -from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig - -if TYPE_CHECKING: - import torch - from collections.abc import Sequence - from lightllm.common.att_resources.allocators.full_att_kv_allocator import FullAttKVAllocator - from lightllm.common.att_resources.full_att.base_state import FullAttentionState - from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig - from lightllm.models.glm5_next.model import Glm5NextTpPartModel - - -def init_model_att_tail(model: Glm5NextTpPartModel) -> None: - """按当前模型自己的 full-attention 层创建 tail，直接加入共享管理器。""" - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options, get_resource_config - - config = get_resource_config(model.config, model_class=type(model)) - options = get_full_attention_options( - config, model.args, model_class=type(model), - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - ) - if not options["logical_layer_groups"]: - return - cfg = model.linear_config - state = create_kpool_state( - config=KPoolStateConfig( - logical_layer_ids=tuple( - logical_layer_id for logical_layer_group in options["logical_layer_groups"] - for logical_layer_id in logical_layer_group - ), pool_size=cfg.index_kpool, - head_dim=cfg.index_head_dim, dtype=cfg.full_att_dtype, - tp_rank=model.tp_rank_, tp_world_size=model.tp_world_size_, - ), - request_capacity=model.max_req_num + 1, mtp_step=model.args.mtp_step, device="cuda", - ) - model.att_resource_manager.set_att_state(model.att_resource_role, state) - - -class Glm5NextAttResourceManager(AttResourceManager): - # ======================================================================== - # 一、模型资源：创建 State 与绑定存储 - # ======================================================================== - - def __init__( - self, - cache_config: Glm5NextResourceConfig, - request_capacity: int, - mtp_step: int, - *, - tp_rank: int = 0, - device: str | torch.device | int = "cuda", - role: str = "main", - bounded_logical_layer_ids: Sequence[int] | None = None, - full_att_logical_layer_ids: Sequence[int] | None = None, - ): - super().__init__(request_capacity=request_capacity) - self.cache_config = cache_config - self.cpu_tail_restore = True - logical_layer_ids = ( - tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) - if bounded_logical_layer_ids is None - else tuple(bounded_logical_layer_ids) - ) - config = ConvSSMStateConfig( - logical_layer_ids=logical_layer_ids, - global_k_heads=cache_config.global_linear_k_heads, - global_v_heads=cache_config.global_linear_v_heads, - key_dim=cache_config.head_linear_k_dim, - value_dim=cache_config.head_linear_v_dim, - conv_width=cache_config.conv_kernel_size - 1, - conv_dtype=cache_config.conv_state_dtype, - ssm_dtype=cache_config.ssm_state_dtype, - tp_rank=tp_rank, - tp_world_size=cache_config.tp_world_size, - ) - if logical_layer_ids: - state = create_kda_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) - self.set_att_state(role, state) - self.att_tail_type = f"kpool_c{cache_config.index_kpool}_{dtype_name(dtype=cache_config.full_att_dtype)}" - own_full_att_logical_layer_ids = ( - tuple( - range( - cache_config.full_attention_interval - 1, - cache_config.all_layer_num, - cache_config.full_attention_interval, - ) - ) - if full_att_logical_layer_ids is None - else tuple(full_att_logical_layer_ids) - ) - tail_logical_layers_by_role = [(role, own_full_att_logical_layer_ids)] - if role == "main" and cache_config.draft_full_att_kv_layer_num: - tail_logical_layers_by_role.append( - ( - "draft0", - tuple( - range( - cache_config.all_layer_num, - cache_config.all_layer_num + cache_config.draft_full_att_kv_layer_num, - ) - ), - ) - ) - for role, logical_layer_ids in tail_logical_layers_by_role: - if logical_layer_ids: - state = create_kpool_state( - config=KPoolStateConfig( - logical_layer_ids=logical_layer_ids, - pool_size=cache_config.index_kpool, - head_dim=cache_config.index_head_dim, - dtype=cache_config.full_att_dtype, - tp_rank=tp_rank, - tp_world_size=cache_config.tp_world_size, - ), - request_capacity=request_capacity, mtp_step=mtp_step, device=device, - ) - self.set_att_state(role, state) - - def bind_storage( - self, - full_state: FullAttentionState, - full_allocator: FullAttKVAllocator, - *, - big_page_capacity: int, - small_page_capacity: int, - draft_full_state: FullAttentionState | None = None, - ): - """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" - assert self.full_att_allocator is None, "attention storage is already bound" - assert all( - logical_layer_id < self.cache_config.all_layer_num - for logical_layer_group in full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group - ), "main State must only contain main layers" - assert (draft_full_state is not None) == bool( - self.cache_config.draft_full_att_kv_layer_num - ), "draft State must match the model configuration" - if draft_full_state is not None: - assert all( - logical_layer_id >= self.cache_config.all_layer_num - for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group - ), "draft State must only contain draft layers" - - state = getattr( - self, - f"main_kda_conv_{dtype_name(dtype=self.cache_config.conv_state_dtype)}" - f"_ssm_{dtype_name(dtype=self.cache_config.ssm_state_dtype)}_state", - ) - state.initialize_page_pools( - big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, - pin_memory=state.conv_state.is_cuda, - ) - big_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) if big_page_capacity else None - small_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None - self.set_att_state("main", full_state) - if draft_full_state is not None: - self.set_att_state("draft0", draft_full_state) - self.full_att_allocator = full_allocator - self.bounded_att_big_page_allocator = big_allocator - self.bounded_att_small_page_allocator = small_allocator - self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 if big_allocator else None - self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 if big_allocator else None - self.finish_state_creation() diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index c4f9af4c22..a00f450304 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -1,5 +1,6 @@ import json import os +from typing import TYPE_CHECKING import torch import triton @@ -17,6 +18,14 @@ from .layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.att_resources.bounded_att import BoundedAttentionState + from lightllm.common.att_resources.att_tail import AttentionTailState + + class Glm5NextTpPartModel(TpPartBaseModel): is_hybrid_att_model = True pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight @@ -26,15 +35,21 @@ class Glm5NextTpPartModel(TpPartBaseModel): transformer_layer_infer_class = Glm5NextTransformerLayerInfer @classmethod - def create_bounded_attention_definition(cls, model_config, args, *, model_logical_layer_ids, tp_rank=0, tp_world_size=1): + def _create_bounded_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["BoundedAttentionState"]: + """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig from lightllm.common.basemodel.basemodel_utils import get_resource_config config = get_resource_config(model_config, model_class=cls) + assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) if not logical_layer_ids: - return None + return [] state = create_kda_state( config=ConvSSMStateConfig( logical_layer_ids=logical_layer_ids, @@ -48,9 +63,9 @@ def create_bounded_attention_definition(cls, model_config, args, *, model_logica tp_rank=tp_rank, tp_world_size=tp_world_size, ), - request_capacity=1, device="meta", + request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, ) - return state + return [state] def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: @@ -84,23 +99,34 @@ def _init_some_value(self): self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - def _init_att_resource_manager(self): - from .att_resource_manager import Glm5NextAttResourceManager, init_model_att_tail - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention, init_model_bounded_attention - - self.linear_config = Glm5NextResourceConfig.from_model_config(self.config, self.args, include_draft=False) - if self.args.linear_att_hash_page_size % self.linear_config.index_kpool: - raise ValueError("GLM cache pages must align to index_kpool") - if self.args.mtp_step and self.args.enable_prefill_decode_mixed: - raise ValueError("MTP linear attention does not support mixed prefill/decode") - if not self.is_mtp_draft_model: - self.att_resource_manager = Glm5NextAttResourceManager( - self.linear_config, self.max_req_num + 1, self.args.mtp_step, - tp_rank=self.tp_rank_, bounded_logical_layer_ids=(), full_att_logical_layer_ids=(), - ) - init_model_bounded_attention(self) - init_model_att_tail(self) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] + + @classmethod + def _create_att_tail_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["AttentionTailState"]: + """按本模型的 full-attention 层创建 K-pool tail，不依赖其他创建接口的副作用。""" + from lightllm.models.glm5_next.att_resource_config import create_model_att_tail_states + + return create_model_att_tail_states( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, request_capacity=request_capacity, + tp_rank=tp_rank, tp_world_size=tp_world_size, device=device, + ) @classmethod def _full_attention_options(cls, config, args): diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 9cecde71ff..c07511ee21 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from copy import deepcopy from lightllm.common.basemodel import TpPartBaseModel @@ -10,6 +11,13 @@ from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.att_resources.att_tail import AttentionTailState + + class Glm5NextMTPModel(TpPartBaseModel): is_mtp_draft_model = True pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight @@ -59,13 +67,34 @@ def _init_some_value(self): self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - def _init_att_resource_manager(self): - from lightllm.models.glm5_next.att_resource_manager import init_model_att_tail - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] + + @classmethod + def _create_att_tail_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["AttentionTailState"]: + """按本模型的 full-attention 层创建 K-pool tail，不依赖其他创建接口的副作用。""" + from lightllm.models.glm5_next.att_resource_config import create_model_att_tail_states - self.linear_config = self.main_model.linear_config - init_model_att_tail(self) - init_model_full_attention(self) + return create_model_att_tail_states( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, request_capacity=request_capacity, + tp_rank=tp_rank, tp_world_size=tp_world_size, device=device, + ) @classmethod def _resource_config(cls, model_config): diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 1aa3037304..5e494d0772 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import os import json import torch @@ -15,6 +16,12 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class LlamaTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -32,13 +39,19 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention - - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _init_config(self): super()._init_config() diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index d4db886ae2..40669e9fea 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch from lightllm.common.basemodel import TpPartBaseModel from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight @@ -9,6 +10,12 @@ from lightllm.utils.envs_utils import get_env_start_args +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class MistralTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -26,13 +33,19 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _verify_params(self): assert self.load_way in ["HF"], "mistral only supports HF format to load Now!" diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 05cd662a03..9e79f99887 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch import numpy as np from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -13,6 +14,12 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class MixtralTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -34,13 +41,19 @@ def _init_config(self): super()._init_config() return - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention - - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _verify_params(self): assert self.load_way in ["HF"], "mixtral only supports HF format to load Now!" diff --git a/lightllm/models/qwen3next/att_resource_manager.py b/lightllm/models/qwen3next/att_resource_manager.py deleted file mode 100644 index 48e07c7c29..0000000000 --- a/lightllm/models/qwen3next/att_resource_manager.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Qwen3Next attention resources with explicit model-owned State objects.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.utils import dtype_name -from lightllm.common.att_resources.allocators import StatePageAllocator -from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig - -if TYPE_CHECKING: - import torch - from collections.abc import Sequence - from lightllm.common.att_resources.allocators.full_att_kv_allocator import FullAttKVAllocator - from lightllm.common.att_resources.full_att.base_state import FullAttentionState - from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig - - -class Qwen3NextAttResourceManager(AttResourceManager): - # ======================================================================== - # 一、模型资源：创建 State 与绑定存储 - # ======================================================================== - - def __init__( - self, - cache_config: Qwen3NextResourceConfig, - request_capacity: int, - mtp_step: int, - *, - tp_rank: int = 0, - device: str | torch.device | int = "cuda", - role: str = "main", - logical_layer_ids: Sequence[int] | None = None, - ): - super().__init__(request_capacity=request_capacity) - self.cache_config = cache_config - logical_layer_ids = ( - tuple(i for i in range(cache_config.all_layer_num) if (i + 1) % cache_config.full_attention_interval) - if logical_layer_ids is None - else tuple(logical_layer_ids) - ) - config = ConvSSMStateConfig( - logical_layer_ids=logical_layer_ids, - global_k_heads=cache_config.global_linear_k_heads, - global_v_heads=cache_config.global_linear_v_heads, - key_dim=cache_config.head_linear_k_dim, - value_dim=cache_config.head_linear_v_dim, - conv_width=cache_config.conv_kernel_size - 1, - conv_dtype=cache_config.conv_state_dtype, - ssm_dtype=cache_config.ssm_state_dtype, - tp_rank=tp_rank, - tp_world_size=cache_config.tp_world_size, - ) - if logical_layer_ids: - state = create_gdn_state(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) - self.set_att_state(role, state) - - def bind_storage( - self, - full_state: FullAttentionState, - full_allocator: FullAttKVAllocator, - *, - big_page_capacity: int, - small_page_capacity: int, - draft_full_state: FullAttentionState | None = None, - ): - """按已确定的 KV 容量绑定存储，创建本模型的大小页池与共享分配器。""" - assert self.full_att_allocator is None, "attention storage is already bound" - assert all( - logical_layer_id < self.cache_config.all_layer_num - for logical_layer_group in full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group - ), "main State must only contain main layers" - assert (draft_full_state is not None) == bool( - self.cache_config.draft_full_att_kv_layer_num - ), "draft State must match the model configuration" - if draft_full_state is not None: - assert all( - logical_layer_id >= self.cache_config.all_layer_num - for logical_layer_group in draft_full_state.config.logical_layer_groups for logical_layer_id in logical_layer_group - ), "draft State must only contain draft layers" - - state = getattr( - self, - f"main_gdn_conv_{dtype_name(dtype=self.cache_config.conv_state_dtype)}" - f"_ssm_{dtype_name(dtype=self.cache_config.ssm_state_dtype)}_state", - ) - state.initialize_page_pools( - big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, - pin_memory=state.conv_state.is_cuda, - ) - big_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) if big_page_capacity else None - small_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None - self.set_att_state("main", full_state) - if draft_full_state is not None: - self.set_att_state("draft0", draft_full_state) - self.full_att_allocator = full_allocator - self.bounded_att_big_page_allocator = big_allocator - self.bounded_att_small_page_allocator = small_allocator - self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 if big_allocator else None - self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 if big_allocator else None - self.finish_state_creation() diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 0d0c247960..6937660d44 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -1,5 +1,5 @@ import torch -from typing import Optional +from typing import Optional, TYPE_CHECKING import triton from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( @@ -11,10 +11,14 @@ ) from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.att_resources.bounded_att import BoundedAttentionState + logger = init_logger(__name__) @@ -36,15 +40,22 @@ def __init__(self, kvargs) -> None: super().__init__(kvargs) @classmethod - def create_bounded_attention_definition(cls, model_config, args, *, model_logical_layer_ids, tp_rank=0, tp_world_size=1): + def _create_bounded_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, + tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + ) -> list["BoundedAttentionState"]: + """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig from lightllm.common.basemodel.basemodel_utils import get_resource_config config = get_resource_config(model_config, model_class=cls) + assert config["num_attention_heads"] % tp_world_size == 0 + assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) if not logical_layer_ids: - return None + return [] state = create_gdn_state( config=ConvSSMStateConfig( logical_layer_ids=logical_layer_ids, @@ -58,9 +69,9 @@ def create_bounded_attention_definition(cls, model_config, args, *, model_logica tp_rank=tp_rank, tp_world_size=tp_world_size, ), - request_capacity=1, device="meta", + request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, ) - return state + return [state] def _init_triton(self): def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: @@ -79,22 +90,19 @@ def _init_config(self): super()._init_config() self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - def _init_att_resource_manager(self): - from .att_resource_manager import Qwen3NextAttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention, init_model_bounded_attention - - assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 - start_args: StartArgs = get_env_start_args() - self.linear_config = Qwen3NextResourceConfig.from_model_config(self.config, start_args, include_draft=False) - - if start_args.mtp_step and start_args.enable_prefill_decode_mixed: - raise ValueError("MTP linear attention does not support mixed prefill/decode") - if not self.is_mtp_draft_model: - self.att_resource_manager = Qwen3NextAttResourceManager( - self.linear_config, self.max_req_num + 1, start_args.mtp_step, tp_rank=self.tp_rank_, logical_layer_ids=() - ) - init_model_bounded_attention(self) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] @classmethod def _full_attention_options(cls, config, args): diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index d874a5abad..4bb9ec1888 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from lightllm.models.starcoder.layer_infer.transformer_layer_infer import StarcoderTransformerLayerInfer from lightllm.models.starcoder.layer_infer.pre_layer_infer import StarcoderPreLayerInfer from lightllm.models.starcoder.layer_weights.transformer_layer_weight import StarcoderTransformerLayerWeight @@ -8,6 +9,12 @@ from lightllm.common.basemodel import InferStateInfo +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class StarcoderTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = StarcoderPreAndPostLayerWeight @@ -22,13 +29,19 @@ class StarcoderTpPartModel(TpPartBaseModel): def __init__(self, kvargs): super().__init__(kvargs) - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _init_config(self): super()._init_config() diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 801897d63d..33f8850c5d 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.starcoder2.layer_weights.pre_and_post_layer_weight import Starcoder2PreAndPostLayerWeight @@ -10,6 +11,12 @@ from lightllm.common.basemodel import TpPartBaseModel +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class Starcoder2TpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = Starcoder2PreAndPostLayerWeight @@ -25,13 +32,19 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_resource_manager(self): - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.basemodel.basemodel_utils import init_model_full_attention - - if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) - init_model_full_attention(self) + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" + from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + + state = create_full_attention_definition( + model_class=cls, model_config=model_config, args=args, + model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [] if state is None else [state] def _init_config(self): super()._init_config() diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index f94899143e..98d26db78f 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -166,6 +166,8 @@ def init_model(self, kvargs): manager: AttResourceManager = self.model.att_resource_manager assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_creation() + manager.load_full_att_calibration(args=self.model.args) manager.initialize_resources( self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/att_resource_config.py index 9135c275f4..c6813eebc3 100644 --- a/lightllm/utils/att_resource_config.py +++ b/lightllm/utils/att_resource_config.py @@ -1,90 +1,14 @@ -"""静态 FP8 标定加载与 CPU cache 服务布局构建。""" +"""CPU cache 服务布局构建；仅依赖 State 配置，不加载静态标定。""" from __future__ import annotations from typing import TYPE_CHECKING -import json - -import torch - from lightllm.common.att_resources import AttResourceManager from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from collections.abc import Mapping from lightllm.common.att_resources.types import StateTransferLayout - from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - - -def load_static_fp8_calibration( - args: StartArgs | EasyDict, - *, - per_head: bool, - storage_layer_count: int, - heads: int, - tp_rank: int = 0, - tp_world_size: int = 1, - storage_layer_start: int = 0, -) -> tuple[dict[str, object], torch.Tensor, torch.Tensor | None]: - """读取模型和 CPU cache 共用的静态标定参数。 - - storage_layer_count 是当前 State 的物理层数；storage_layer_start 是其在标定表中的 - 物理行起点，两者都不是逻辑层号。共享逻辑层使用同一行标定参数，不重复占用行。 - 返回 Tensor 保留在 CPU；当前 TP 分片在每行中选取对应的 head 参数。 - """ - path = args.kv_quant_calibration_config_path - if path is None: - assert not per_head, "fp8kv_sph requires kv_quant_calibration_config_path with q_calibration" - return {"scale": "unit"}, torch.ones((storage_layer_count, 2), dtype=torch.float32), None - with open(path) as f: - cfg = json.load(f) - calibrated_storage_layer_count = cfg["num_layers"] - # 当前模型只读取自己的物理层区间；main/draft 偏移由装配流程明确提供。 - storage_layer_end = storage_layer_start + storage_layer_count - assert 0 <= storage_layer_start < storage_layer_end <= calibrated_storage_layer_count, ( - "calibration does not cover this model's attention layers" - ) - - if not per_head: - from lightllm.utils.config_utils import get_model_architectures - - assert cfg["architectures"] == get_model_architectures(args.model_dir), ( - "calibration architecture does not match the model" - ) - mode = "per_head" if per_head else "per_tensor" - assert cfg.get("quant_type", mode) == mode, "calibration quantization mode does not match the State" - limits = torch.finfo(torch.float8_e4m3fn) - assert cfg.get("qmin", limits.min) == limits.min and cfg.get("qmax", limits.max) == limits.max, ( - "calibration range does not match float8_e4m3fn" - ) - - def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]) -> torch.Tensor: - tensor = torch.tensor(values["scales"], dtype=torch.float32) - assert tuple(values["scales_shape"]) == tensor.shape == shape, ( - f"calibration scales and scales_shape must match {shape}" - ) - assert bool((torch.isfinite(tensor) & (tensor > 0)).all()), "calibration scales must be finite and positive" - return tensor[storage_layer_start:storage_layer_end] - - def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: - assert calibrated_heads > 0 and (heads * tp_world_size) % calibrated_heads == 0, ( - "calibration heads do not cover the model's KV heads" - ) - expanded = tensor.repeat_interleave(heads * tp_world_size // calibrated_heads, dim=-1) - return expanded[:, tp_rank * heads : (tp_rank + 1) * heads].contiguous() - - if per_head: - calibrated_heads = cfg["num_head"] - all_scales = scale_tensor(cfg, (calibrated_storage_layer_count, 2 * calibrated_heads)) - k, v = all_scales.split(calibrated_heads, dim=-1) - scales = torch.cat((shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) - q_cfg = cfg.get("q_calibration") - assert q_cfg, "fp8kv_sph calibration config requires q_calibration" - q_scales = shard(scale_tensor(q_cfg, (calibrated_storage_layer_count, q_cfg["num_head"])), q_cfg["num_head"]) - else: - scales, q_scales = scale_tensor(cfg, (calibrated_storage_layer_count, 2)), None - return cfg, scales, q_scales def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: @@ -92,11 +16,7 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: from transformers.configuration_utils import PretrainedConfig from lightllm.models import get_draft_model_class from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS - from lightllm.common.basemodel.basemodel_utils import ( - create_full_attention_definition, - create_bounded_attention_definition, - get_resource_config, - ) + from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.llm_utils import get_llm_model_class # 阶段一：读取 main 和启用 attention 的 draft 配置，不构造实际模型。 @@ -111,24 +31,23 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: config, _ = PretrainedConfig.get_config_dict(directory) models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) - # 阶段二：复用模型定义入口创建 meta State，并累计逻辑层号和标定物理行偏移。 + # 阶段二：复用模型定义入口创建 meta State，并累计逻辑层号；CPU 页不包含静态标定参数。 manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 - logical_layer_start = calibration_storage_layer_start = 0 + logical_layer_start = 0 for role, model_class, config in models: count = get_resource_config(config, model_class=model_class)["num_hidden_layers"] logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) - full = create_full_attention_definition( - model_class, config, args, model_logical_layer_ids=logical_layer_ids, calibration_storage_layer_start=calibration_storage_layer_start + full_states = model_class._create_full_att_states( + model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, ) - bounded = create_bounded_attention_definition( - model_class, config, args, model_logical_layer_ids=logical_layer_ids + bounded_states = model_class._create_bounded_att_states( + model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, + request_capacity=1, device="meta", ) - for state in (full, bounded): - if state is None: - continue + # tail 不占 CPU 页；无需创建 tail，也不需要实际模型或临时 CUDA 运行态。 + for state in (*full_states, *bounded_states): manager.set_att_state(role, state) logical_layer_start += count - calibration_storage_layer_start += len(full.config.logical_layer_groups) if full is not None else 0 # 阶段三：确认混合模型的页边界一致，由管理器聚合 CPU 联合页布局。 manager.finish_state_creation() diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 3b477e3e83..4ca6c08ad9 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -309,6 +309,8 @@ def main(): manager: AttResourceManager = model.att_resource_manager assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_creation() + manager.load_full_att_calibration(args=model.args) manager.initialize_resources( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 86246f3835..a3ccd92fc0 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -1290,6 +1290,8 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) manager: AttResourceManager = model.att_resource_manager assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_creation() + manager.load_full_att_calibration(args=model.args) manager.initialize_resources( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index defa3e59ca..faae489405 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -17,7 +17,9 @@ FullAttentionConfig, ) from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig -from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import set_env_start_args @@ -56,13 +58,25 @@ def main(): FullAttentionConfig(tuple((layer,) for layer in tuple(range(12))), 1, 256, torch.bfloat16, 0, 8), ) full.initialize_storage(tokens) - resources = Qwen3NextAttResourceManager(cfg, request_capacity=1, mtp_step=0) - resources.bind_storage( - full, - FullAttKVAllocator(tokens), - big_page_capacity=args.pages + 2, - small_page_capacity=0, + resources = AttResourceManager(request_capacity=1) + bounded_state = create_gdn_state( + config=ConvSSMStateConfig( + logical_layer_ids=tuple(layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval), + global_k_heads=cfg.global_linear_k_heads, global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, conv_dtype=cfg.conv_state_dtype, ssm_dtype=cfg.ssm_state_dtype, + tp_rank=0, tp_world_size=cfg.tp_world_size, + ), + request_capacity=1, mtp_step=0, device="cuda", ) + bounded_state.initialize_page_pools(big_page_capacity=args.pages + 2, small_page_capacity=0) + resources.set_att_state("main", full) + resources.set_att_state("main", bounded_state) + resources.full_att_allocator = FullAttKVAllocator(tokens) + resources.bounded_att_big_page_allocator = StatePageAllocator(args.pages + 2, reserved_slots=2) + resources.bounded_att_cpu_load_big_page_slot = args.pages + resources.bounded_att_cpu_offload_big_page_slot = args.pages + 1 + resources.finish_state_creation() pages = list(range(args.pages)) slots = resources.bounded_att_big_page_allocator.alloc(args.pages) cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) @@ -77,8 +91,7 @@ def main(): cpu_page_buffers, pages, layouts=resources.get_cpu_transfer_layouts(), - page_kind="big", - slot_ids=slots, + big_page_slot_ids=slots, ) torch.cuda.synchronize() @@ -89,6 +102,7 @@ def restore(): token_indices, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=slots, + req_idx=0, ) for _ in range(args.warmup): diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 1def48fb89..f759b2f05d 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -16,11 +16,40 @@ if TYPE_CHECKING: + from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig + from lightllm.common.att_resources.allocators import FullAttKVAllocator from collections.abc import Mapping, Sequence from lightllm.common.att_resources.att_resource_manager import AttResourceManager from lightllm.common.basemodel.basemodel import TpPartBaseModel +def bind_test_fp8_calibration( + state: StaticFP8FullAttentionState, + scales: torch.Tensor, + q_scales: torch.Tensor | None = None, +) -> None: + """将测试的本地期望值展开为标定文件格式，再通过生产绑定入口初始化。 + + 每个 TP 分片放置相同测试值；这样重建任意 rank 都能保留原测试的本地 scale。 + 真实的非重复全局 head 切片由 test_model_calibration 单独覆盖。 + """ + config = state.config + calibration = {"num_layers": len(config.logical_layer_groups), "quant_type": config.scale_mode} + if config.scale_mode == "per_head": + heads = config.head_num + k, v = scales.split(heads, dim=-1) + scales = torch.cat((k.repeat(1, config.tp_world_size), v.repeat(1, config.tp_world_size)), dim=-1) + assert q_scales is not None + q_scales = q_scales.repeat(1, config.tp_world_size) + calibration["num_head"] = heads * config.tp_world_size + calibration["q_calibration"] = { + "num_head": q_scales.shape[1], "scales_shape": list(q_scales.shape), "scales": q_scales.tolist(), + } + calibration["scales_shape"] = list(scales.shape) + calibration["scales"] = scales.tolist() + state.bind_calibration(calibration=calibration, storage_layer_start=0) + + def rebuild_full_state( template: FullAttentionState, *, @@ -39,6 +68,9 @@ def rebuild_full_state( if isinstance(template, StaticFP8FullAttentionState): prepared_full_state = type(template)( config, page_size=page_size, device=device, + ) + bind_test_fp8_calibration( + state=prepared_full_state, scales=template.scales if scales is None else scales, q_scales=template.q_scales if q_scales is None else q_scales, ) @@ -151,3 +183,70 @@ def full_state_test_fields( name: TensorField(tensor=tensor, axes=("layer", "token", "head", "group") if name.endswith("_scale") else axes) for name, tensor in buffers.items() } + + +def build_hybrid_resources( + config: Qwen3NextResourceConfig, + request_capacity: int, + mtp_step: int, + full_state: FullAttentionState, + full_allocator: FullAttKVAllocator, + *, + big_page_capacity: int, + small_page_capacity: int, + draft_full_state: FullAttentionState | None = None, + tp_rank: int = 0, +) -> AttResourceManager: + """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" + from lightllm.common.att_resources import AttResourceManager + from lightllm.common.att_resources.allocators import StatePageAllocator + from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state + from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig + + manager = AttResourceManager(request_capacity=request_capacity) + glm = isinstance(config, Glm5NextResourceConfig) + create_state = create_kda_state if glm else create_gdn_state + bounded_state = create_state( + config=ConvSSMStateConfig( + logical_layer_ids=tuple( + layer for layer in range(config.all_layer_num) if (layer + 1) % config.full_attention_interval + ), + global_k_heads=config.global_linear_k_heads, + global_v_heads=config.global_linear_v_heads, + key_dim=config.head_linear_k_dim, + value_dim=config.head_linear_v_dim, + conv_width=config.conv_kernel_size - 1, + conv_dtype=config.conv_state_dtype, + ssm_dtype=config.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=config.tp_world_size, + ), + request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + ) + bounded_state.initialize_page_pools( + big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=True, + ) + manager.set_att_state("main", bounded_state) + for role, state in (("main", full_state), ("draft0", draft_full_state)): + if state is None: + continue + manager.set_att_state(role, state) + if glm: + tail = create_kpool_state( + config=KPoolStateConfig( + logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), + pool_size=config.index_kpool, head_dim=config.index_head_dim, dtype=config.full_att_dtype, + tp_rank=tp_rank, tp_world_size=config.tp_world_size, + ), + request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + ) + manager.set_att_state(role, tail) + manager.full_att_allocator = full_allocator + manager.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) + manager.bounded_att_small_page_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None + manager.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 + manager.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 + manager.cpu_tail_restore = glm + manager.finish_state_creation() + return manager diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index c826fbe1bc..bc5eef4e67 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -126,6 +126,17 @@ def test_backend_allocates_all_roles_before_warmup( checked = [] mapping = SimpleNamespace(bind_hold_page=lambda indices: None) request_allocations = [] + manager_creations = [] + from lightllm.common.att_resources import AttResourceManager + + def create_empty_manager(request_capacity: int): + manager = AttResourceManager(request_capacity=request_capacity) + assert not manager._states_created + assert not any(name.endswith("_state") for name in manager.__dict__) + manager_creations.append(request_capacity) + return manager + + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttResourceManager", create_empty_manager) def create_req_manager(max_request_num: int, max_sequence_length: int): assert len(models) == 1 + draft_count @@ -250,6 +261,8 @@ def start_runtime(): assert all(other.att_resource_manager is models[0].att_resource_manager for other in models[: index + 1]) assert all(not hasattr(other, "att_resource_manager") for other in models[index + 1 :]) main = models[0] + main.att_resource_manager.finish_state_creation() + main.att_resource_manager.load_full_att_calibration(args=args) main.att_resource_manager.initialize_resources( args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction ) @@ -263,6 +276,7 @@ def start_runtime(): start_runtime() assert request_allocations == [(2, 32)] + assert manager_creations == [models[0].max_req_num + 1] assert captured == {id(model) for model in models} assert checked == [models[0]] assert ready == ([True] if via_backend else []) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index 86f1f20534..d032d2fabd 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -1,5 +1,7 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + from unit_tests.common.att_resources.model_helpers import full_state_test_fields from unit_tests.common.att_resources.model_helpers import full_state_test_buffers @@ -47,12 +49,17 @@ def definition(kind: str): ) if kind.startswith("static"): heads = 2 if kind == "static_head" else 1 - return StaticFP8FullAttentionState( + state = StaticFP8FullAttentionState( StaticFP8AttentionConfig( - **vars(config), scale_mode="per_head" if heads == 2 else "per_tensor", calibration_id="test" + **vars(config), scale_mode="per_head" if heads == 2 else "per_tensor" ), - page_size=4, device="meta", scales=torch.ones((2, heads * 2), dtype=torch.float32), q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, + page_size=4, device="meta", + ) + bind_test_fp8_calibration( + state=state, scales=torch.ones((2, heads * 2), dtype=torch.float32), + q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, ) + return state if kind.startswith("packed"): fp8 = kind == "packed_fp8" kv_dim = 512 if kind == "packed_glm" else 576 @@ -123,6 +130,7 @@ def test_budget_matches_token_stores_including_hold_and_padding( resources = AttResourceManager(request_capacity=1) resources.set_att_state("main", rebuild_full_state(template, config=template.config)) monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) + resources.finish_state_creation() if expected is None: with pytest.raises(MemoryError, match="HOLD"): resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 88088abcf0..176bb7592e 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -2,6 +2,8 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + from unit_tests.common.att_resources.model_helpers import rebuild_full_state from typing import TYPE_CHECKING @@ -121,9 +123,12 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu state = StaticFP8FullAttentionState( layout_config=StaticFP8AttentionConfig( logical_layer_groups=((5,),), global_head_num=head_num, head_dim=128, dtype=torch.bfloat16, - scale_mode=scale_mode, calibration_id="split-storage", + scale_mode=scale_mode, ), - page_size=4, device="cuda", scales=scales, + page_size=4, device="cuda", + ) + bind_test_fp8_calibration( + state=state, scales=scales, q_scales=torch.ones((1, head_num), dtype=torch.float32, device="cuda") if scale_mode == "per_head" else None, ) state.initialize_storage(8) diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index 5a816ae05a..b8a985d661 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -85,6 +85,8 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( model.tp_rank_, model.tp_world_size_ = 0, 1 model.max_total_token_num, model.max_req_num, model.mem_fraction = None, 1, 1.0 prepare_model_states(model) + model.att_resource_manager.finish_state_creation() + model.att_resource_manager.load_full_att_calibration(args=args) model.att_resource_manager.initialize_resources( args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index eec81f7195..48b651487b 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -13,16 +13,17 @@ from unit_tests.common.att_resources.test_tail_state import make_tail from unit_tests.common.att_resources.test_capacity import definition from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.basemodel_utils import init_model_full_attention +from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: - model = TpPartBaseModel.__new__(TpPartBaseModel) + model = LlamaTpPartModel.__new__(LlamaTpPartModel) model.args, model.config = SimpleNamespace(), {} model.tp_rank_, model.tp_world_size_ = 0, 1 model.layers_infer = [SimpleNamespace(layer_num_=0, indexer=SimpleNamespace())] - model.att_resource_manager = AttResourceManager(request_capacity=1) + model.max_req_num = 0 + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) state = GQAFullAttentionState( layout_config=FullAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=128, dtype=torch.bfloat16, @@ -32,7 +33,7 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: "lightllm.common.basemodel.basemodel_utils.create_full_attention_definition", lambda *args, **kwargs: state, ) - init_model_full_attention(model) + model.prepare_att_resources() assert model.att_resource_manager.main_gqa_bf16_state is state assert not hasattr(model.layers_infer[0], "full_att_state_name") @@ -40,7 +41,7 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("role", ["main", "draft0"]) -def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str) -> None: +def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, monkeypatch: pytest.MonkeyPatch) -> None: manager = AttResourceManager(request_capacity=1) for model_role in ("main", "draft0"): for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): @@ -55,7 +56,9 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str) -> model.max_req_num = 0 model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) model.main_model = SimpleNamespace(att_resource_manager=manager) - model._init_att_resource_manager = lambda: setattr(model, "att_resource_manager", manager) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttResourceManager", lambda request_capacity: manager) + model.config, model.args = {}, SimpleNamespace(mtp_mode=None) + model.tp_rank_, model.tp_world_size_ = 0, 1 model.layers_infer = [] for logical_layer_id in range(4): layer = TransformerLayerInfer.__new__(TransformerLayerInfer) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index 0fc9f4dd2a..b51d6c59d2 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -1,4 +1,4 @@ -"""Model-owned calibration slices must not depend on State initialization order.""" +"""统一标定阶段不依赖 State 注册顺序，不在布局定义阶段加载参数。""" from __future__ import annotations @@ -13,9 +13,6 @@ import torch from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.full_att import FullAttentionState -from unit_tests.common.att_resources.model_helpers import owned_states -from lightllm.common.basemodel.basemodel_utils import init_model_full_attention from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.models.qwen3next.model import Qwen3NextTpPartModel @@ -23,10 +20,11 @@ if TYPE_CHECKING: from pathlib import Path + from easydict import EasyDict @pytest.mark.parametrize("encoding", ["fp8kv_sph", "fp8kv_spt"]) -def test_calibration_slices_follow_model_configs_before_other_states_exist(tmp_path: Path, encoding: str): +def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encoding: str): config = dict( architectures=["Qwen3NextForCausalLM"], num_hidden_layers=8, @@ -97,13 +95,26 @@ def test_calibration_slices_follow_model_configs_before_other_states_exist(tmp_p models.append(model) layer_start += layer_count - # Prepare KV definitions in reverse: no preceding model has a manager or State yet. + # 反向创建、注册，标定仍按 main 和 draft 的数值顺序绑定。 + manager = AttResourceManager(request_capacity=1) for index in reversed(range(len(models))): model = models[index] - assert all(not hasattr(previous, "att_resource_manager") for previous in models[:index]) - model.att_resource_manager = AttResourceManager(request_capacity=1) - init_model_full_attention(model) - states = owned_states(model.att_resource_manager.__dict__, FullAttentionState) + model.att_resource_manager = manager + states = [] + if not (model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att")): + states = model._create_full_att_states( + model_config=model.config, args=model.args, + 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_, + ) + for state in states: + assert state.resource_id is None + assert state.scales is None and state.q_scales is None + manager.set_att_state(model.att_resource_role, state) + manager.finish_state_creation() + manager.load_full_att_calibration(args=args) + for index, model in enumerate(models): + states = manager.get_full_attention_states(model.att_resource_role) state = states[0] if states else None expected_layers, (start, end) = definitions[index][-2:] if not expected_layers: @@ -114,3 +125,106 @@ def test_calibration_slices_follow_model_configs_before_other_states_exist(tmp_p assert torch.equal(state.scales[:, 0], scales[start:end]) if encoding == "fp8kv_sph": assert torch.equal(state.q_scales[:, 0], scales[start:end] * 2) + + +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, tp_world_size: int, +): + from builtins import open as original_open + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + + # 第 0 行属于普通精度的共享层；其余行分别属于 main、draft2、draft10 的 FP8 State。 + scales = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + q_scales = scales[:, :2] + 20 + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text(json.dumps(dict( + num_layers=4, num_head=2, quant_type="per_head", scales_shape=[4, 4], scales=scales.tolist(), + q_calibration=dict(num_head=2, scales_shape=[4, 2], scales=q_scales.tolist()), + ))) + args = StartArgs(kv_quant_calibration_config_path=str(calibration_path)) + reads = [] + + def count_open(file: str, *args: object, **kwargs: object): + if file == str(calibration_path): + reads.append(file) + return original_open(file, *args, **kwargs) + + monkeypatch.setattr("builtins.open", count_open) + for tp_rank in range(tp_world_size): + manager = AttResourceManager(request_capacity=1) + expected = [] + # 特意打乱角色及同角色 State 的注册顺序。 + for role, groups, row, encoding in ( + ("draft10", ((5,),), 3, "fp8kv_sph"), + ("main", ((2, 3),), 1, "fp8kv_sph"), + ("draft2", ((4,),), 2, "fp8kv_sph"), + ("main", ((0, 1),), 0, "None"), + ): + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=groups, global_head_num=4, head_dim=16, dtype=torch.bfloat16, + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), kv_type=encoding, device="meta", + ) + manager.set_att_state(role, state) + if encoding != "None": + expected.append((state, row)) + with pytest.raises(AssertionError, match="State creation"): + manager.load_full_att_calibration(args=args) + manager.finish_state_creation() + manager.load_full_att_calibration(args=args) + heads = 4 // tp_world_size + local_heads = slice(tp_rank * heads, (tp_rank + 1) * heads) + for state, row in expected: + expected_k = scales[row:row + 1, :2].repeat_interleave(2, dim=-1)[:, local_heads] + expected_v = scales[row:row + 1, 2:].repeat_interleave(2, dim=-1)[:, local_heads] + expected_q = q_scales[row:row + 1].repeat_interleave(2, dim=-1)[:, local_heads] + assert torch.equal(state.scales, torch.cat((expected_k, expected_v), dim=-1)) + assert torch.equal(state.q_scales, expected_q) + assert state.k_storage is None + assert len(reads) == tp_world_size # 每个管理器只读取一次，不随 State 数量增加。 + + +@pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) +def test_cpu_layout_does_not_require_calibration_file( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, cpu_args: EasyDict, encoding: str, +): + from lightllm.utils.att_resource_config import build_cpu_cache_layouts + + config = dict(model_type="llama", num_attention_heads=2, hidden_size=32, num_hidden_layers=2) + (tmp_path / "config.json").write_text(json.dumps(config)) + cpu_args.model_dir = str(tmp_path) + cpu_args.llm_kv_type = encoding + cpu_args.data_type = "bfloat16" + cpu_args.kv_quant_calibration_config_path = str(tmp_path / "does-not-exist.json") + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: LlamaTpPartModel) + layouts = build_cpu_cache_layouts() + assert len(layouts) == 1 + (layout,) = layouts.values() + assert layout.byte_end > layout.byte_start + assert {region.field for region in layout.regions} == {"k", "v"} + + +@pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) +def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=((0, 1), (2,)), global_head_num=2, head_dim=16, dtype=torch.bfloat16, + ), kv_type=encoding, device="cpu", + ) + manager = AttResourceManager(request_capacity=1) + manager.set_att_state("main", state) + manager.finish_state_creation() + with pytest.raises(AssertionError, match="calibration must be bound"): + state.initialize_storage(4) + if encoding == "fp8kv_sph": + with pytest.raises(AssertionError, match="requires a calibration file"): + manager.load_full_att_calibration(args=StartArgs()) + else: + manager.load_full_att_calibration(args=StartArgs()) + assert torch.equal(state.scales, torch.ones((2, 2))) + state.initialize_storage(4) + assert state.k_storage.shape[0] == 2 diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index c78b4f7f53..17783eb6b6 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -131,6 +131,8 @@ def check_runtime_before_profile(device: str | torch.device | int): for item in models: item.bind_req_manager(req_manager=req_manager) assert all(item.req_manager is req_manager for item in models) + model.att_resource_manager.finish_state_creation() + model.att_resource_manager.load_full_att_calibration(args=args) model.att_resource_manager.initialize_resources( args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) @@ -142,8 +144,9 @@ def check_runtime_before_profile(device: str | torch.device | int): assert resources.full_att_allocator.size == full_state.capacity == 32 assert not hasattr(model, "mem_manager") assert not hasattr(model, "full_att_state") + full_interval = 4 if glm else model.config["full_attention_interval"] assert full_state.config.logical_layer_groups == tuple( - (layer,) for layer in model.linear_config.get_full_att_logical_layer_ids() if layer < model.config["num_hidden_layers"] + (layer,) for layer in range(full_interval - 1, model.config["num_hidden_layers"], full_interval) ) assert observed == [main_linear_state(resources).conv_state.data_ptr()] assert resources.bounded_att_big_page_allocator.size == 6 diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 4c86507c08..4f8ce9e3ae 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -1,5 +1,6 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources from unit_tests.common.att_resources.model_helpers import full_state_test_fields from unit_tests.common.att_resources.model_helpers import full_state_test_buffers @@ -21,8 +22,6 @@ from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig -from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager -from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args @@ -58,8 +57,6 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int all_layer_num=4, draft_full_att_kv_layer_num=draft_layers, ) - manager_cls = Glm5NextAttResourceManager if glm else Qwen3NextAttResourceManager - resources = manager_cls(config, 3, mtp, tp_rank=rank) if glm: full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( @@ -77,9 +74,10 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int ) full.initialize_storage(16) main, draft = create_full_attention_roles(full, 4) - resources.bind_storage( - main, - FullAttKVAllocator(16, 4), + resources = build_hybrid_resources( + config=config, request_capacity=3, mtp_step=mtp, tp_rank=rank, + full_state=main, + full_allocator=FullAttKVAllocator(16, 4), draft_full_state=draft, big_page_capacity=4, small_page_capacity=2, @@ -95,22 +93,22 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: i targets = [make_resources(glm, tp, rank, mtp) for rank in range(tp)] expected = [] for rank, (resources, full) in enumerate(sources): - cfg = resources.cache_config + cfg = main_linear_state(resources).config for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.fill_(1 if glm else rank // max(1, tp // 2) + 1) state = main_linear_state(resources) local_start = 0 for local_heads, global_heads in ( - (cfg.num_linear_k_heads, cfg.global_linear_k_heads), - (cfg.num_linear_k_heads, cfg.global_linear_k_heads), - (cfg.num_linear_v_heads, cfg.global_linear_v_heads), + (cfg.num_k_heads, cfg.global_k_heads), + (cfg.num_k_heads, cfg.global_k_heads), + (cfg.num_v_heads, cfg.global_v_heads), ): start = rank * local_heads // max(1, tp * local_heads // global_heads) for h in range(local_heads): state.conv_state[:, 0, local_start + h * 8 : local_start + (h + 1) * 8].fill_(start + h + 3) local_start += local_heads * 8 - for h in range(cfg.num_linear_v_heads): - state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_linear_v_heads + h + 7) + for h in range(cfg.num_v_heads): + state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_v_heads + h + 7) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) torch.cuda.synchronize() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index e8771c686d..c0f958c79e 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -118,6 +118,7 @@ def test_model_and_cpu_service_use_identical_storage_geometry( model.max_total_token_num, model.max_req_num = 19, 1 prepare_model_states(model) resources = model.att_resource_manager + resources.finish_state_creation() resources.initialize_resources(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) (state,) = resources.get_full_attention_states("main") assert state.config.head_num == local_heads and state.config.head_dim == dim @@ -172,7 +173,51 @@ def memory_info(device: str | torch.device | int): template = rebuild_full_state(state, config=state.config, token_capacity=0, device="meta") resources = AttResourceManager(request_capacity=1) resources.set_att_state("main", rebuild_full_state(template, config=template.config)) + resources.finish_state_creation() resources.initialize_resources(args, token_capacity=None, mem_fraction=0.75, device="cpu") assert resources.full_att_allocator.size == 28 assert resources.get_pd_page_bytes() == reference.get_pd_page_bytes() assert observed_devices == [torch.device("cuda", 3)] + + +def test_cpu_layout_uses_all_model_factory_states( + args: StartArgs | EasyDict, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +): + """同一模型返回多种方案和共享层时，CPU 页采用实际 State 列表。""" + from lightllm.common.att_resources.full_att import FullAttentionState, GQAFullAttentionState + + class MixedStorageModel(LlamaTpPartModel): + @classmethod + def _create_full_att_states( + cls, model_config: dict[str, object], args: StartArgs | EasyDict, *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list[FullAttentionState]: + (template,) = super()._create_full_att_states( + model_config=model_config, args=args, model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, tp_world_size=tp_world_size, + ) + return [ + GQAFullAttentionState( + layout_config=dataclasses.replace( + template.config, logical_layer_groups=(model_logical_layer_ids[:2],), dtype=torch.bfloat16, + ), page_size=args.page_size, device="meta", + ), + GQAFullAttentionState( + layout_config=dataclasses.replace( + template.config, logical_layer_groups=(model_logical_layer_ids[2:],), dtype=torch.float16, + ), page_size=args.page_size, device="meta", + ), + ] + + config = dict(model_type="llama", num_attention_heads=8, hidden_size=128, num_hidden_layers=3) + (tmp_path / "config.json").write_text(json.dumps(config)) + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: MixedStorageModel) + cpu_layouts = build_cpu_cache_layouts() + model = MixedStorageModel.__new__(MixedStorageModel) + model.config, model.args = config, args + model.tp_rank_, model.tp_world_size_, model.max_req_num = 0, 1, 1 + prepare_model_states(model) + model.att_resource_manager.finish_state_creation() + runtime_layouts = model.att_resource_manager.get_cpu_transfer_layouts() + assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_state", "main_gqa_fp16_state"} + assert cpu_layouts == runtime_layouts diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 3de2789a81..5be91d3172 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -1,15 +1,15 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + from unit_tests.common.att_resources.model_helpers import full_state_test_fields from unit_tests.common.att_resources.model_helpers import full_state_test_buffers from unit_tests.common.att_resources.model_helpers import read_kv_storage -from unit_tests.common.att_resources.model_helpers import rebuild_full_state from typing import TYPE_CHECKING -from dataclasses import replace from math import prod import pytest @@ -44,10 +44,14 @@ def make_quant_state(kind: str, device: str | torch.device | int): if kind.startswith("fp8_"): per_head = kind == "fp8_head" config = StaticFP8AttentionConfig( - ((1,), (3,)), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor", calibration_id="test" + ((1,), (3,)), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor" ) scales = torch.arange(1, 9 if per_head else 5, dtype=torch.float32, device=device).reshape(2, -1) / 16 - prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device, scales=scales) + prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device) + bind_test_fp8_calibration( + state=prepared_full_state, scales=scales, + q_scales=torch.ones((2, config.head_num), dtype=torch.float32) if per_head else None, + ) prepared_full_state.initialize_storage(8) return prepared_full_state fp8 = kind == "packed_fp8" @@ -208,26 +212,16 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): assert not read_kv_storage(state)[1].any() -def test_fp8_calibration_mismatch_fails(): - src = make_quant_state("fp8_head", "cpu") - dst = rebuild_full_state(src, config=replace(src.config, calibration_id="other")) - with pytest.raises(AssertionError, match="calibration"): - dst.copy_tokens_from_other_dp( - src_token_indices=torch.tensor([0], dtype=torch.long), - dst_token_indices=torch.tensor([1], dtype=torch.long), - remote_dp_full_att_state=src, - ) - - @pytest.mark.parametrize("kind", ["fp8_tensor", "fp8_head"]) def test_deferred_storage_preserves_real_fp8_calibration(kind: str): source = make_quant_state(kind, "cpu") state = StaticFP8FullAttentionState( source.config, page_size=source.page_size, device="meta", - scales=source.scales, q_scales=source.q_scales, ) + assert state.scales is None and state.q_scales is None + bind_test_fp8_calibration(state=state, scales=source.scales, q_scales=source.q_scales) assert read_kv_storage(state) is None and not state.scales.is_meta - assert state.scales is source.scales + assert torch.equal(state.scales, source.scales) state.device = torch.device("cpu") state.initialize_storage(source.capacity) assert torch.equal(state.scales, source.scales) diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index c1980b3339..78a8333562 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -1,5 +1,7 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + from unit_tests.common.att_resources.model_helpers import read_kv_storage from collections.abc import Callable @@ -38,7 +40,6 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( q_scales = torch.ones((2, 2), dtype=torch.float32) if kv_type == "fp8kv_sph" else None state = create_gqa_state( config=config, kv_type=kv_type, page_size=4, device="meta", group_size=16, - scales=scales, q_scales=q_scales, calibration_id="test-calibration", ) assert state.get_state_type() == expected_type assert state.resource_id is None @@ -48,8 +49,11 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( assert state.config.head_dim == config.head_dim and config.dtype == torch.bfloat16 assert state.capacity == 0 and read_kv_storage(state) is None if kv_type.startswith("fp8"): - assert state.scales is scales and not state.scales.is_meta - assert state.q_scales is q_scales + assert state.scales is None and state.q_scales is None + bind_test_fp8_calibration(state=state, scales=scales, q_scales=q_scales) + assert torch.equal(state.scales, scales) + if q_scales is not None: + assert torch.equal(state.q_scales, q_scales) state.device = torch.device("cpu") state.initialize_storage(8) assert state.k_storage.shape == state.v_storage.shape == (2, 12, 2, width) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 4ed5aa0326..041c2c2411 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -52,6 +52,7 @@ def check_created_before_allocation(manager: AttResourceManager, *args: object, monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) + manager.finish_state_creation() manager.initialize_resources(args, token_capacity=8, mem_fraction=1.0, device="cpu") assert manager.get_full_attention_states() is cached[0] for role, meta in zip(("main", "draft0"), definitions): diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 6d467fbc83..8f372322b7 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -1,3 +1,4 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses from types import SimpleNamespace @@ -11,7 +12,6 @@ PackedMLAAttentionConfig, ) from lightllm.common.req_manager import ReqManager -from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name @@ -96,13 +96,13 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t ) full.initialize_storage(16) req = ReqManager(3, 16) - resources = Glm5NextAttResourceManager(config, 4, mtp_step) from unit_tests.common.att_resources.model_helpers import create_full_attention_roles main_full, draft_full = create_full_attention_roles(full, config.all_layer_num) - resources.bind_storage( - main_full, - FullAttKVAllocator(full.capacity, full.page_size), + resources = build_hybrid_resources( + config=config, request_capacity=4, mtp_step=mtp_step, + full_state=main_full, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft_full, big_page_capacity=4, small_page_capacity=2, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 7881e58519..0dbbb4e749 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,3 +1,4 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources from unit_tests.common.att_resources.model_helpers import full_state_test_fields import dataclasses @@ -14,7 +15,6 @@ PackedMLAAttentionConfig, ) from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page -from lightllm.models.glm5_next.att_resource_manager import Glm5NextAttResourceManager from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs @@ -84,11 +84,11 @@ def create(tp, target_mtp_step=None): page_size=args.page_size, ) full.initialize_storage(32) - resources = Glm5NextAttResourceManager(config, 4, local_mtp_step, tp_rank=rank) main, draft = create_full_attention_roles(full, config.all_layer_num) - resources.bind_storage( - main, - FullAttKVAllocator(full.capacity, full.page_size), + resources = build_hybrid_resources( + config=config, request_capacity=4, mtp_step=local_mtp_step, tp_rank=rank, + full_state=main, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft, big_page_capacity=6, small_page_capacity=2, diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index 6acb945931..a6cd2ce5e6 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -175,6 +175,8 @@ def make_runtime( draft.tp_world_size_, draft.tp_rank_ = tp, rank prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) drafts.append(draft) + model.att_resource_manager.finish_state_creation() + model.att_resource_manager.load_full_att_calibration(args=args) model.att_resource_manager.initialize_resources( args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction ) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index a988f3fa28..dff7de34c9 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -1,3 +1,4 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses import json from types import SimpleNamespace @@ -12,7 +13,6 @@ FullAttentionConfig, ) from lightllm.common.req_manager import ReqManager -from lightllm.models.qwen3next.att_resource_manager import Qwen3NextAttResourceManager from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq @@ -72,10 +72,10 @@ def linear_cache(monkeypatch, tmp_path, request): ) full.initialize_storage(384) requests = ReqManager(2, 512) - resources = Qwen3NextAttResourceManager(config, 3, 0) - resources.bind_storage( - full, - FullAttKVAllocator(full.capacity, full.page_size), + resources = build_hybrid_resources( + config=config, request_capacity=3, mtp_step=0, + full_state=full, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), big_page_capacity=5, small_page_capacity=2, ) From f77b3bab98830d7945ca91064e29169d0a197fa3 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 07:35:10 +0000 Subject: [PATCH 08/56] refactor: move attention state layout decisions into models --- docs/cache_manager_design.md | 8 +- lightllm/common/basemodel/basemodel_utils.py | 86 ++----------------- lightllm/models/bloom/model.py | 24 ++++-- lightllm/models/deepseek2/model.py | 37 ++++++-- lightllm/models/deepseek3_2/model.py | 35 +++++++- lightllm/models/gemma3/model.py | 36 ++++++-- .../layer_infer/transformer_layer_infer.py | 2 +- lightllm/models/gemma4/model.py | 78 ++++++++++------- lightllm/models/gemma_2b/model.py | 36 ++++++-- .../models/glm5_next/att_resource_config.py | 32 ------- lightllm/models/glm5_next/model.py | 57 ++++++++---- lightllm/models/glm5_next_mtp/model.py | 55 ++++++++---- lightllm/models/llama/model.py | 26 ++++-- lightllm/models/mistral/model.py | 24 ++++-- lightllm/models/mixtral/model.py | 26 ++++-- lightllm/models/qwen3next/model.py | 33 ++++--- lightllm/models/starcoder/model.py | 33 ++++--- lightllm/models/starcoder2/model.py | 26 ++++-- .../common/att_resources/test_layer_groups.py | 5 +- .../test_layer_state_selection.py | 2 +- 20 files changed, 389 insertions(+), 272 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index e2e0ac176c..6997e7eb19 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1634,7 +1634,7 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -单一 KV 方案的 `_create_full_att_states()` 可复用 `basemodel_utils.create_full_attention_definition(...)`，将返回的单个定义包装为列表；不存在 full 层时返回空列表。该 helper 不访问管理器，也不构成独立启动阶段。bounded 模型工厂直接调用对应算法 create，不再先生成 meta State 再重建实际 State。GLM main/draft 的 tail 工厂仅从自身配置读取 K-pool 参数，不依赖 main 的 `linear_config`。 +`_create_full_att_states()` 由具体模型直接构造 `FullAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 `AttResourceManager.initialize_resources()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 @@ -1661,9 +1661,11 @@ State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_si full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例如复制的 MLA latent；否则按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 -`basemodel_utils.py` 提供 `get_resource_config(model_config, *, model_class=None)` 和 `get_full_attention_options(config, args, *, model_class=None)`，统一处理字段别名、text config 和默认 KV 维度。它们是普通工具函数，BaseModel 不再声明对应的 `_resource_config`、`_full_attention_options` 方法；具体模型仅在存在专有规则时提供可选 hook，并调用工具函数取得默认配置。传入 `model_class` 时，工具函数选择该模型的专有实现，否则使用通用逻辑。 +`basemodel_utils.py` 仅保留 `get_resource_config(model_config, *, model_class=None)`，处理字段别名、text config 和模型专有的配置转换。原生 MTP、DFlash 等模型可通过 `_resource_config` 将原始配置转为自己的有效层配置；该工具不选择 attention 算法、不构造 State，也不通过 options 字典分派布局。 -同一文件中的 `create_full_attention_definition(model_class, ...)` 根据配置直接调用算法级 `create_gqa_state` / `create_mla_state`，供模型 full 工厂复用。CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例或运行推理初始化。返回的 State 列表统一注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可包含同一模型的多个 full/bounded State；静态标定不参与布局计算，CPU 服务无需打开标定文件。原生 MTP、DFlash 等配置转换由对应模型实现。`att_resource_config.py` 仅保留 CPU 服务布局构建；标定加载由管理器集中组织，具体格式解释由静态 FP8 State 完成。 +GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `_create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 + +CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`att_resource_config.py` 仅保留 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index 46af0f0822..c29f95e64a 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -1,4 +1,4 @@ -"""Shared helpers for model resource definitions used by startup and CPU caches.""" +"""模型启动与 CPU cache 共用的配置展开及字段别名处理。""" from __future__ import annotations @@ -6,15 +6,13 @@ from lightllm.common.build_utils import repair_config if TYPE_CHECKING: - from collections.abc import Sequence - from lightllm.common.att_resources.full_att import FullAttentionState - from easydict import EasyDict from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.server.core.objs.start_args_type import StartArgs -def get_resource_config(model_config: dict[str, object], *, model_class: type[TpPartBaseModel] | None = None): - """Normalize resource config, applying a model-specific override when present.""" +def get_resource_config( + model_config: dict[str, object], *, model_class: type[TpPartBaseModel] | None = None, +) -> dict[str, object]: + """优先采用模型专有配置转换；默认展开 text_config 并补齐常见字段别名。""" build_config = getattr(model_class, "_resource_config", None) if build_config is not None: return build_config(model_config) @@ -23,77 +21,3 @@ def get_resource_config(model_config: dict[str, object], *, model_class: type[Tp repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) repair_config(config, same_names=["num_hidden_layers", "n_layer"]) return config - - -def get_full_attention_options( - config: dict[str, object], - args: StartArgs | EasyDict, - *, - model_class: type[TpPartBaseModel] | None = None, - model_logical_layer_ids: Sequence[int] | None = None, -): - """Return default KV geometry or the model's specialized storage options.""" - build_options = getattr(model_class, "_full_attention_options", None) - if build_options is not None: - options = build_options(config, args) - else: - heads = config.get("num_key_value_heads", config["num_attention_heads"]) - dim = config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]) - logical_layer_count = config["num_hidden_layers"] - options = dict( - heads=heads, dim=dim, - logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in range(logical_layer_count)), - ) - if model_logical_layer_ids is not None: - # 模型选项使用本模型内的层号；转换为实际逻辑层号时保留共享分组及物理层顺序。 - options["logical_layer_groups"] = tuple( - tuple( - model_logical_layer_ids[local_logical_layer_id] - for local_logical_layer_id in logical_layer_group - if local_logical_layer_id < len(model_logical_layer_ids) - ) - for logical_layer_group in options["logical_layer_groups"] - if any(local_logical_layer_id < len(model_logical_layer_ids) for local_logical_layer_id in logical_layer_group) - ) - return options - - -def create_full_attention_definition( - model_class: type[TpPartBaseModel], - model_config: dict[str, object], - args: StartArgs | EasyDict, - *, - tp_rank: int = 0, - tp_world_size: int = 1, - model_logical_layer_ids: Sequence[int] | None = None, -) -> FullAttentionState | None: - """按模型配置创建算法 State；静态 FP8 标定由全部 State 就绪后的统一阶段绑定。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state, create_mla_state - from lightllm.utils.torch_dtype_utils import get_torch_dtype - - if model_class.is_mtp_draft_model and args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): - return None - config = get_resource_config(model_config, model_class=model_class) - options = get_full_attention_options(config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids) - if not options["logical_layer_groups"]: - return None - mla = options.get("mla", False) - packed = options.get("packed", False) - dtype = options.get("dtype") - layout_config = FullAttentionConfig( - logical_layer_groups=options["logical_layer_groups"], - global_head_num=options["heads"], - head_dim=options["dim"], - dtype=get_torch_dtype(args.data_type) if dtype is None else dtype, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - replicated=options.get("replicated", False) or mla or packed, - ) - if mla or packed: - return create_mla_state( - config=layout_config, kv_type=args.llm_kv_type, packed=packed, page_size=args.page_size, device="meta", - ) - - return create_gqa_state( - config=layout_config, kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", - ) diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 0685e0f9b5..b0f52c182b 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -36,14 +36,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _init_config(self): super()._init_config() diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index a70f21b6b8..39e763bb56 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch from typing import final from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer @@ -12,6 +13,12 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class Deepseek2TpPartModel(LlamaTpPartModel): # weight class transformer_weight_class = Deepseek2TransformerLayerWeight @@ -57,12 +64,30 @@ def _verify_params(self): return super()._verify_params() @classmethod - def _full_attention_options(cls, config, args): - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options - - options = get_full_attention_options(config, args) - options.update(heads=1, dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], mla=True) - return options + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + ) + return [state] def _init_to_get_yarn_rotary(self): from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask, get_deepseek_mscale diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 197b25fafe..e9baff5b26 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -1,9 +1,16 @@ +from typing import TYPE_CHECKING from lightllm.models.deepseek2.model import Deepseek2TpPartModel from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class Deepseek3_2TpPartModel(Deepseek2TpPartModel): # weight class @@ -13,10 +20,30 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): transformer_layer_infer_class = Deepseek3_2TransformerLayerInfer @classmethod - def _full_attention_options(cls, config, args): - options = super()._full_attention_options(config, args) - options["packed"] = True - return options + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", + ) + return [state] def _init_att_backend(self): self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0)(model=self) diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index 5c769842b9..fb013cd29b 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import os import json import torch @@ -15,6 +16,12 @@ from transformers import AutoConfig from lightllm.utils.log_utils import init_logger +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + logger = init_logger(__name__) @@ -137,12 +144,29 @@ def _init_custom(self): self._init_to_get_rotary() @classmethod - def _full_attention_options(cls, config, args): - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options - - options = get_full_attention_options(config, args) - options.update(dim=256, dtype=torch.bfloat16) - return options + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=256, + dtype=torch.bfloat16, + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + ) + return [state] def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index bd7720c81e..b8a27cc459 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -22,7 +22,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): time. The KV cache layout is uniform (sliding shape: num_kv_heads=16, head_dim=256); full-attention layers pack their (4, 512) tensor into the first 8 heads of the 16-head slot at cache-write time, then reshape on - read. See Gemma4TpPartModel._full_attention_options for context. + read. See Gemma4TpPartModel._create_full_att_states for context. """ def __init__(self, layer_num, network_config): diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 003840ab77..056a63f3d2 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import os import json import torch @@ -17,6 +18,12 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class Gemma4TpPartModel(LlamaTpPartModel): pre_and_post_weight_class = Gemma4PreAndPostLayerWeight transformer_weight_class = Gemma4TransformerLayerWeight @@ -80,43 +87,50 @@ def _verify_params(self): return @classmethod - def _full_attention_options(cls, config, args): - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options - - # Uniform per-layer KV cache layout. The per-layer cache slot must fit - # whichever layer type has the largest per-token K/V width: sliding - # (num_key_value_heads * head_dim) or full - # (num_global_kv * global_head_dim). Keep cache_slot_dim = head_dim - # and pick cache_slot_num = max-width / head_dim. For 31B this - # collapses to num_key_value_heads; for E4B the full-attn shape wins - # (2*512 > 2*256), so it uses 4 storage slots of 256 dims. - # Gemma4TransformerLayerInfer.__init__ computes the same value and - # uses it to pack/unpack K/V at write/read time. + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + # 每层的 KV 槽位须容纳 sliding/full 两类 attention 的最大宽度。 head_dim = config["head_dim"] num_global_kv = config.get("num_global_key_value_heads") or config["num_key_value_heads"] sliding_total = config["num_key_value_heads"] * head_dim full_total = num_global_kv * config["global_head_dim"] per_token_k_width = max(sliding_total, full_total) - if per_token_k_width % head_dim: - raise ValueError("Gemma4 KV width must contain complete storage heads") - options = get_full_attention_options(config, args) - options.update(heads=per_token_k_width // head_dim, dim=head_dim) - shared_logical_layer_count = config.get("num_kv_shared_layers") or 0 - if shared_logical_layer_count: - # 后部共享层读取同类型的最后一个非共享层；仅生产层占一份物理 KV。 - # 由模型配置生成分组，GPU State 和 CPU cache 布局使用完全相同的层顺序。 - storage_layer_count = config["num_hidden_layers"] - shared_logical_layer_count - assert 0 < storage_layer_count < config["num_hidden_layers"], "invalid KV-shared layer count" - logical_layer_groups = [[logical_layer_id] for logical_layer_id in range(storage_layer_count)] - for logical_layer_id in range(storage_layer_count, config["num_hidden_layers"]): - for source_logical_layer_id in range(storage_layer_count - 1, -1, -1): - if config["layer_types"][source_logical_layer_id] == config["layer_types"][logical_layer_id]: - logical_layer_groups[source_logical_layer_id].append(logical_layer_id) - break - else: - raise AssertionError("KV-shared layer requires an earlier producer of the same attention type") - options["logical_layer_groups"] = tuple(tuple(logical_layer_group) for logical_layer_group in logical_layer_groups) - return options + assert per_token_k_width % head_dim == 0, "Gemma4 KV width must contain complete storage heads" + + shared_layer_count = config.get("num_kv_shared_layers") or 0 + storage_layer_count = len(model_logical_layer_ids) - shared_layer_count + assert 0 <= shared_layer_count < len(model_logical_layer_ids), "invalid KV-shared layer count" + layer_groups = [[layer_id] for layer_id in model_logical_layer_ids[:storage_layer_count]] + # 后部共享层读取同类型的最后一个生产层，只增加逻辑映射，不增加物理存储行。 + for local_layer in range(storage_layer_count, len(model_logical_layer_ids)): + for source_layer in range(storage_layer_count - 1, -1, -1): + if config["layer_types"][source_layer] == config["layer_types"][local_layer]: + layer_groups[source_layer].append(model_logical_layer_ids[local_layer]) + break + else: + raise AssertionError("KV-shared layer requires an earlier producer of the same attention type") + logical_layer_groups = tuple(tuple(group) for group in layer_groups) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=per_token_k_width // head_dim, + head_dim=head_dim, + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + ) + return [state] def _init_att_backend(self): # Gemma-4 has per-layer heterogeneous attention: sliding layers use diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index 2e65a13d79..c417982e2f 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from lightllm.models.gemma_2b.layer_weights.transformer_layer_weight import Gemma_2bTransformerLayerWeight from lightllm.models.gemma_2b.layer_weights.pre_and_post_layer_weight import Gemma_2bPreAndPostLayerWeight from lightllm.models.gemma_2b.layer_infer.pre_layer_infer import Gemma_2bPreLayerInfer @@ -7,6 +8,12 @@ from lightllm.models.llama.model import LlamaTpPartModel +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.att_resources.full_att import FullAttentionState + + class Gemma_2bTpPartModel(LlamaTpPartModel): # weight class pre_and_post_weight_class = Gemma_2bPreAndPostLayerWeight @@ -35,9 +42,28 @@ def _verify_params(self): return @classmethod - def _full_attention_options(cls, config, args): - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options + def _create_full_att_states( + cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, + model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - options = get_full_attention_options(config, args) - options.update(dim=config["hidden_size"] // config["num_attention_heads"], replicated=True) - return options + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config["hidden_size"] // config["num_attention_heads"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + replicated=True, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + ) + return [state] diff --git a/lightllm/models/glm5_next/att_resource_config.py b/lightllm/models/glm5_next/att_resource_config.py index eec434053b..571ccdd084 100644 --- a/lightllm/models/glm5_next/att_resource_config.py +++ b/lightllm/models/glm5_next/att_resource_config.py @@ -9,8 +9,6 @@ from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.att_resources.att_tail import AttentionTailState from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs @@ -65,33 +63,3 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) - - -def create_model_att_tail_states( - model_class: type[TpPartBaseModel], model_config: dict[str, object], args: StartArgs | EasyDict, *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int, - tp_rank: int, tp_world_size: int, device: str, -) -> list[AttentionTailState]: - """由模型配置创建 K-pool tail 列表，不读取模型实例或注册管理器。""" - from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options, get_resource_config - - config = get_resource_config(model_config, model_class=model_class) - options = get_full_attention_options( - config, args, model_class=model_class, model_logical_layer_ids=model_logical_layer_ids, - ) - if not options["logical_layer_groups"]: - return [] - assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" - state = create_kpool_state( - config=KPoolStateConfig( - logical_layer_ids=tuple( - logical_layer_id for logical_layer_group in options["logical_layer_groups"] - for logical_layer_id in logical_layer_group - ), - pool_size=config["index_kpool"], head_dim=config["index_head_dim"], - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, - ) - return [state] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index a00f450304..eec783e0b5 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -104,14 +104,28 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + config = get_resource_config(model_config, model_class=cls) + assert args.llm_kv_type == "None", "GLM K-pool requires unquantized KV" + cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_logical_layer_ids()) + if not logical_layer_groups: + return [] + state = create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["kv_lora_rank"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] @classmethod def _create_att_tail_states( @@ -119,21 +133,26 @@ def _create_att_tail_states( model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", ) -> list["AttentionTailState"]: - """按本模型的 full-attention 层创建 K-pool tail，不依赖其他创建接口的副作用。""" - from lightllm.models.glm5_next.att_resource_config import create_model_att_tail_states - - return create_model_att_tail_states( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, request_capacity=request_capacity, - tp_rank=tp_rank, tp_world_size=tp_world_size, device=device, - ) + """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" + from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - @classmethod - def _full_attention_options(cls, config, args): - if args.llm_kv_type != "None": - raise ValueError(f"GLM K-pool does not support llm_kv_type={args.llm_kv_type}") + config = get_resource_config(model_config, model_class=cls) cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) - return dict(heads=1, dim=config["kv_lora_rank"], logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in cfg.get_full_att_logical_layer_ids()), packed=True) + logical_layer_ids = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_logical_layer_ids()) + if not logical_layer_ids: + return [] + assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" + state = create_kpool_state( + config=KPoolStateConfig( + logical_layer_ids=logical_layer_ids, + pool_size=config["index_kpool"], head_dim=config["index_head_dim"], + dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ) + return [state] def _init_att_backend(self): self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index c07511ee21..fb2cdc2a11 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -72,14 +72,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["kv_lora_rank"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] @classmethod def _create_att_tail_states( @@ -87,14 +99,25 @@ def _create_att_tail_states( model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", ) -> list["AttentionTailState"]: - """按本模型的 full-attention 层创建 K-pool tail，不依赖其他创建接口的副作用。""" - from lightllm.models.glm5_next.att_resource_config import create_model_att_tail_states - - return create_model_att_tail_states( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, request_capacity=request_capacity, - tp_rank=tp_rank, tp_world_size=tp_world_size, device=device, + """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" + from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_ids = model_logical_layer_ids + if not logical_layer_ids: + return [] + assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" + state = create_kpool_state( + config=KPoolStateConfig( + logical_layer_ids=logical_layer_ids, + pool_size=config["index_kpool"], head_dim=config["index_head_dim"], + dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, ) + return [state] @classmethod def _resource_config(cls, model_config): @@ -104,10 +127,6 @@ def _resource_config(cls, model_config): config.update(num_hidden_layers=1, n_layer=1) return config - @classmethod - def _full_attention_options(cls, config, args): - return dict(heads=1, dim=config["kv_lora_rank"], logical_layer_groups=((0,),), packed=True) - def _init_att_backend(self): backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 5e494d0772..d80873e196 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -44,14 +44,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _init_config(self): super()._init_config() diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 40669e9fea..a0ca1c6faf 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -38,14 +38,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _verify_params(self): assert self.load_way in ["HF"], "mistral only supports HF format to load Now!" diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 9e79f99887..0b86d866d2 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -46,14 +46,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _verify_params(self): assert self.load_way in ["HF"], "mixtral only supports HF format to load Now!" diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 6937660d44..4cb7d323ee 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -95,23 +95,28 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, - ) - return [] if state is None else [state] + """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - @classmethod - def _full_attention_options(cls, config, args): - if args.llm_kv_type not in ("None", "fp8kv_sph", "fp8kv_spt"): - raise ValueError(f"Qwen3Next does not support llm_kv_type={args.llm_kv_type}") + config = get_resource_config(model_config, model_class=cls) + assert args.llm_kv_type in ("None", "fp8kv_sph", "fp8kv_spt"), "unsupported Qwen3Next KV encoding" cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) - return dict( - heads=cfg.full_att_all_num_kv_heads, dim=cfg.full_att_head_dim, logical_layer_groups=tuple((logical_layer_id,) for logical_layer_id in cfg.get_full_att_logical_layer_ids()) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_logical_layer_ids()) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=cfg.full_att_all_num_kv_heads, + head_dim=cfg.full_att_head_dim, + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) + return [state] def _init_att_backend1(self): role = self.att_resource_role diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index 4bb9ec1888..f78760b858 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -34,14 +34,27 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition + """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["hidden_size"] // config["num_attention_heads"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + replicated=True, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _init_config(self): super()._init_config() @@ -58,14 +71,6 @@ def _reset_num_key_value_heads(self): def _verify_params(self): assert self.load_way == "HF", "StarCoder only support HF format to load Now!" - @classmethod - def _full_attention_options(cls, config, args): - from lightllm.common.basemodel.basemodel_utils import get_full_attention_options - - options = get_full_attention_options(config, args) - options.update(heads=1, dim=config["hidden_size"] // config["num_attention_heads"], replicated=True) - return options - def _init_some_value(self): super()._init_some_value() self.tp_k_head_num_ = self.config["num_key_value_heads"] diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 33f8850c5d..9044c78c2b 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -37,14 +37,26 @@ def _create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: - """按当前模型的方案创建 full-attention 定义，不分配 KV 容量或注册管理器。""" - from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition - - state = create_full_attention_definition( - model_class=cls, model_config=model_config, args=args, - model_logical_layer_ids=model_logical_layer_ids, tp_rank=tp_rank, tp_world_size=tp_world_size, + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.basemodel.basemodel_utils import get_resource_config + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = get_resource_config(model_config, model_class=cls) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=FullAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", ) - return [] if state is None else [state] + return [state] def _init_config(self): super()._init_config() diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 9f8acdaf4f..876267519a 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -10,7 +10,6 @@ from lightllm.common.att_resources import AttResourceManager from lightllm.common.att_resources.full_att import FullAttentionConfig from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.basemodel_utils import create_full_attention_definition from lightllm.models.gemma4.model import Gemma4TpPartModel from unit_tests.common.att_resources.model_helpers import ( full_state_test_buffers, @@ -69,8 +68,8 @@ def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "E head_dim=128, global_head_dim=128, num_kv_shared_layers=2, layer_types=["sliding_attention", "full_attention"] * 3, ) - state = create_full_attention_definition( - model_class=Gemma4TpPartModel, model_config=config, args=pd_args, + (state,) = Gemma4TpPartModel._create_full_att_states( + model_config=config, args=pd_args, model_logical_layer_ids=(10, 11, 12, 13, 14, 15), ) # 共享消费者 14/15 分别使用最后一个同类型生产层 12/13，保持四份物理层存储。 diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 48b651487b..87b6dd68c3 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -30,7 +30,7 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: ), device="meta", ) monkeypatch.setattr( - "lightllm.common.basemodel.basemodel_utils.create_full_attention_definition", lambda *args, **kwargs: state, + LlamaTpPartModel, "_create_full_att_states", classmethod(lambda cls, **kwargs: [state]), ) model.prepare_att_resources() From 53b7bf39143f0baeafea204dad442a72fe6c2580 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 07:49:40 +0000 Subject: [PATCH 09/56] refactor: rename CPU cache layout module --- docs/cache_manager_design.md | 7 +++++-- .../utils/{att_resource_config.py => cpu_cache_layout.py} | 0 lightllm/utils/kv_cache_utils.py | 2 +- unit_tests/common/att_resources/test_backend_startup.py | 2 +- unit_tests/common/att_resources/test_model_calibration.py | 2 +- .../common/att_resources/test_model_storage_definition.py | 2 +- 6 files changed, 9 insertions(+), 6 deletions(-) rename lightllm/utils/{att_resource_config.py => cpu_cache_layout.py} (100%) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 6997e7eb19..28005dd0b1 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -327,6 +327,9 @@ lightllm/common/ lightllm/models// ├── model.py # 三类 State 创建接口，决定本模型的资源组合 └── att_resource_config.py # 模型配置与可复用的 State 创建工具 + +lightllm/utils/ +└── cpu_cache_layout.py # build_cpu_cache_layouts：构建 CPU cache 联合页布局 ``` `allocators/` 与三类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 @@ -363,7 +366,7 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 -`lightllm/utils/att_resource_config.py` 负责把启动参数和标定文件转换为这些显式参数，并调用对应算法入口。full-attention 创建时仍为零普通 token 容量，后续由管理器统一分配；bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 +各模型的三类 State 工厂负责构造配置并调用对应算法入口；`lightllm/utils/cpu_cache_layout.py` 复用模型工厂创建 meta State，生成 CPU cache 联合页布局，不加载标定。full-attention 的实际存储由管理器后续统一分配；GPU 启动时 bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 ```python state = create_gqa_state( @@ -1665,7 +1668,7 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `_create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 -CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`att_resource_config.py` 仅保留 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 +CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 diff --git a/lightllm/utils/att_resource_config.py b/lightllm/utils/cpu_cache_layout.py similarity index 100% rename from lightllm/utils/att_resource_config.py rename to lightllm/utils/cpu_cache_layout.py diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 9f55465c7c..ec89a652ba 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -51,7 +51,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - from lightllm.utils.att_resource_config import build_cpu_cache_layouts + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts layouts = build_cpu_cache_layouts() page_bytes = 0 diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index bc5eef4e67..68433786d0 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -20,7 +20,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.server.router.model_infer.mode_backend import base_backend from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.utils.att_resource_config import build_cpu_cache_layouts +from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args if TYPE_CHECKING: diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index b51d6c59d2..1472454222 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -190,7 +190,7 @@ def count_open(file: str, *args: object, **kwargs: object): def test_cpu_layout_does_not_require_calibration_file( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, cpu_args: EasyDict, encoding: str, ): - from lightllm.utils.att_resource_config import build_cpu_cache_layouts + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts config = dict(model_type="llama", num_attention_heads=2, hidden_size=32, num_hidden_layers=2) (tmp_path / "config.json").write_text(json.dumps(config)) diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index c0f958c79e..aad1f5d245 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -27,7 +27,7 @@ from lightllm.models.gemma3.model import Gemma3TpPartModel from lightllm.models.gemma4.model import Gemma4TpPartModel from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.att_resource_config import build_cpu_cache_layouts +from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.att_resources import AttResourceManager from lightllm.common.att_resources.allocators import FullAttKVAllocator From 92e6450ec7c9e0489e60b5c27946217eb239c685 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 07:51:48 +0000 Subject: [PATCH 10/56] refactor: rename model resource configuration modules --- docs/cache_manager_design.md | 4 ++-- .../models/glm5_next/{att_resource_config.py => config.py} | 2 +- lightllm/models/glm5_next/model.py | 2 +- .../models/qwen3next/{att_resource_config.py => config.py} | 0 lightllm/models/qwen3next/model.py | 2 +- lightllm/utils/backend_validator.py | 2 +- test/cpu_cache_kernel/test_speed.py | 2 +- unit_tests/common/att_resources/model_helpers.py | 4 ++-- unit_tests/common/att_resources/test_model_resources.py | 4 ++-- unit_tests/common/basemodel/attention/linear/test_gdn.py | 2 +- unit_tests/models/glm5_next/test_cache.py | 2 +- unit_tests/models/glm5_next/test_pd_cache.py | 2 +- .../server/router/model_infer/test_hybrid_cpu_state_cache.py | 2 +- 13 files changed, 15 insertions(+), 15 deletions(-) rename lightllm/models/glm5_next/{att_resource_config.py => config.py} (97%) rename lightllm/models/qwen3next/{att_resource_config.py => config.py} (100%) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 28005dd0b1..d2cdcb1702 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -326,7 +326,7 @@ lightllm/common/ lightllm/models// ├── model.py # 三类 State 创建接口，决定本模型的资源组合 -└── att_resource_config.py # 模型配置与可复用的 State 创建工具 +└── config.py # 模型专有配置与派生参数 lightllm/utils/ └── cpu_cache_layout.py # build_cpu_cache_layouts：构建 CPU cache 联合页布局 @@ -346,7 +346,7 @@ CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式 无需模型特有初始化的场景直接使用 `AttResourceManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `set_att_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 -模型特有的 State 组合由 `model.py` 中三类创建接口确定，复用工具放在现有的 `att_resource_config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttResourceManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 +模型特有的 State 组合由 `model.py` 中三类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttResourceManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttResourceManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttResourceManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 resource_id 匹配时可构建临时映射，它不决定模型的资源组织方式。 diff --git a/lightllm/models/glm5_next/att_resource_config.py b/lightllm/models/glm5_next/config.py similarity index 97% rename from lightllm/models/glm5_next/att_resource_config.py rename to lightllm/models/glm5_next/config.py index 571ccdd084..3a66cfc784 100644 --- a/lightllm/models/glm5_next/att_resource_config.py +++ b/lightllm/models/glm5_next/config.py @@ -6,7 +6,7 @@ from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.torch_dtype_utils import get_torch_dtype -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig if TYPE_CHECKING: from easydict import EasyDict diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index eec783e0b5..cbbcc22f6b 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -9,7 +9,7 @@ from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer diff --git a/lightllm/models/qwen3next/att_resource_config.py b/lightllm/models/qwen3next/config.py similarity index 100% rename from lightllm/models/qwen3next/att_resource_config.py rename to lightllm/models/qwen3next/config.py diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 4cb7d323ee..18340e6239 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -11,7 +11,7 @@ ) from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig if TYPE_CHECKING: from easydict import EasyDict diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 8927374915..0d68854ebc 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -99,7 +99,7 @@ def _validate_flashqla(): from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) - from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig linear_config = Qwen3NextResourceConfig.load_from_args() num_k_heads = linear_config.num_linear_k_heads diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index faae489405..4245367ead 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -16,7 +16,7 @@ GQAFullAttentionState, FullAttentionConfig, ) -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.common.att_resources import AttResourceManager from lightllm.common.att_resources.allocators import StatePageAllocator from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index f759b2f05d..463da5008c 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: - from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.common.att_resources.allocators import FullAttKVAllocator from collections.abc import Mapping, Sequence from lightllm.common.att_resources.att_resource_manager import AttResourceManager @@ -202,7 +202,7 @@ def build_hybrid_resources( from lightllm.common.att_resources.allocators import StatePageAllocator from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state - from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig + from lightllm.models.glm5_next.config import Glm5NextResourceConfig manager = AttResourceManager(request_capacity=request_capacity) glm = isinstance(config, Glm5NextResourceConfig) diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 4f8ce9e3ae..af505e8d34 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -20,8 +20,8 @@ FullAttentionConfig, ) from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig -from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 38a5113848..cf448e0704 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -12,7 +12,7 @@ import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.api_cli import make_argument_parser import lightllm.utils.backend_validator as backend_validator diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 8f372322b7..2088c76ad8 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -12,7 +12,7 @@ PackedMLAAttentionConfig, ) from lightllm.common.req_manager import ReqManager -from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 0dbbb4e749..cd5f7360a3 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -15,7 +15,7 @@ PackedMLAAttentionConfig, ) from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page -from lightllm.models.glm5_next.att_resource_config import Glm5NextResourceConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index dff7de34c9..34f5203713 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -16,7 +16,7 @@ from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.models.qwen3next.att_resource_config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name From 20e62d653a8efb5b037330596e53211d2e547bca Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 08:09:24 +0000 Subject: [PATCH 11/56] refactor: resolve attention states and calibration by logical layer --- docs/cache_manager_design.md | 6 + .../att_resources/full_att/base_state.py | 5 + .../att_resources/full_att/gqa_state.py | 8 +- .../full_att/group_quantized_state.py | 8 +- .../att_resources/full_att/mla_state.py | 8 +- .../full_att/packed_mla_state.py | 8 +- .../full_att/static_fp8_state.py | 8 +- .../common/basemodel/attention/base_att.py | 13 +- .../common/basemodel/attention/fa3/fp8.py | 11 +- .../basemodel/attention/flashinfer/fp8.py | 11 +- .../common/basemodel/attention/linear/gdn.py | 41 +++-- .../common/basemodel/attention/linear/kda.py | 41 +++-- .../kv_copy/destindex_copy_kv_flashmla_fp8.py | 104 ++++++------- .../prefill_compact_kv_flashmla_fp8.py | 142 +++++++++--------- .../layer_infer/transformer_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 6 +- .../layer_infer/transformer_layer_infer.py | 1 + .../att_resources/test_compute_binding.py | 3 +- .../common/att_resources/test_layer_groups.py | 5 +- .../basemodel/attention/linear/test_gdn.py | 27 +++- .../attention/test_fp8_logical_layer.py | 104 +++++++++++++ unit_tests/models/glm5_next/test_mtp.py | 21 ++- 22 files changed, 390 insertions(+), 195 deletions(-) create mode 100644 unit_tests/common/basemodel/attention/test_fp8_logical_layer.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index d2cdcb1702..01a5013a07 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -450,6 +450,8 @@ logical_layer_groups = ( 层命名统一使用 `logical_` 和 `storage_` 区分含义：`logical_layer_groups` 描述共享分组，`logical_layer_id` 是推理层传入的逻辑层号，`storage_layer_index` 是 Tensor layer 轴上的物理下标。State 的 `_logical_to_storage_layer` 保存两者的对应关系。KDA/GDN 和 K-pool 当前仍是一对一存储，配置使用 `logical_layer_ids`；三类 State 的 `get_layer_view(logical_layer_id=...)` 使用相同参数语义。FP8 标定的 `storage_layer_count`、`storage_layer_start` 表示物理行数和偏移，不能传逻辑层号。Tensor 轴标签 `"layer"` 保留，表示存储层轴。 +静态 FP8 的 FA3/FlashInfer 后端通过 `AttControl.logical_layer_id` 接收当前逻辑层号，再由 State 的 `get_storage_layer_index()` 查询标定行，不通过 KV 指针反查层号。该接口属于 `FullAttentionState` 的公共契约，所有具体实现按自己的层映射返回物理行；层视图和 KV 写入也使用此接口。 + State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get_layer_view(logical_layer_id=0)` 与 `get_layer_view(logical_layer_id=4)` 返回同一份存储视图；`logical_layer_id=8` 返回另一份。分配、单 token 字节数、CPU/PD 布局和静态量化标定行数均按 `len(logical_layer_groups)` 计算，不按组内逻辑层总数计算，不重复保存物理层数或另一份配置映射。 同组逻辑层共享该 State 的全部配套字段，必须使用兼容的布局和参数。分组只表达数据共享关系，不指定写入者或执行顺序；模型负责跳过共享消费者的 KV 写入，避免覆盖生产层数据。普通模型每个逻辑层单独一组；Gemma4 将后部共享层加入同类型的最后一个非共享层所在组。推理层按自己的逻辑层号访问，CPU 服务与模型使用同一套配置构建物理分组。DP/CPU/PD 传输按相同的物理层顺序搬运，main/draft 仍各自拥有独立 State。 @@ -512,6 +514,8 @@ class FullAttentionState: # ======================================================================== def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_storage_layer_index(self, logical_layer_id: int) -> int: ... + def get_layer_view(self, logical_layer_id: int) -> ViewT: ... @@ -1442,6 +1446,8 @@ CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_pag `accepted_steps` 为逐请求接受的、已经执行过 forward 的输入步数，包含本轮基准输入步，不是“接受的 draft token 数”或新采样输出数。取值范围为 `[1, 本轮实际验证宽度]`；KDA/GDN 的原生索引为 `accepted_steps - 1`。实现不应把不存在的零步回滚当作有效操作。未启用 MTP 时，合法的一步提交可为无数据移动的操作。 +GDN/KDA Backend 不保存 State 引用。prefill/decode 按角色、算法和逻辑层号调用管理器的 `get_att_state()`；MTP decode 的 `init_state()` 在需要游标时局部获取当前角色唯一的对应算法 State，将 `mtp_state_index` 显式传给 fixed/dynamic 初始化函数。此初始化仍要求每个角色只有一个对应算法 State，不假定多个 State 共用游标。 + KDA 的提交可以只更新 `mtp_state_index`，下一轮 kernel 根据该游标选择输入状态。当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0，直接读取 Conv 的 `[:W]` 窗口和 SSM 的 `req_idx*(S+1)` 基准槽位。导入时写到目标的基准位置，并将游标设为 0。实现中以 TODO 标明：若以后支持 MTP 提交后传输，需要按已接受索引 j 选择 Conv 的 `[j:j+W]` 窗口和 SSM 的 `req_idx*(S+1)+j` 槽位。 prefill 的最终状态、普通 decode 的结果和已提交验证状态都可以成为“当前已接受状态”；执行流程必须保证所有相关层已完成且选择信息与数据一致。验证结果尚未提交时，不允许将临时候选作为当前状态发布。大小页保存仅发生在 prefill 阶段，此时 `mtp_state_index` 为 0，直接复制 Conv 的 `[:W]` 窗口及 SSM 的 `req_idx*(S+1)` 基准槽位，不需要先捕获候选状态。实现中以 TODO 标明：若以后支持 decode/MTP 后保存，需要改为按已接受游标选取数据。 diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index c693e53b3a..e15a4232c1 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -134,6 +134,11 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" ... + @abstractmethod + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + ... + @abstractmethod def get_layer_view(self, logical_layer_id: int) -> Any: """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index 6611e6b4a7..1253efc78d 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -88,15 +88,19 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index 84186c6af8..3dcc961e67 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -105,15 +105,19 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """返回该层 K/V 编码数据及各自的分组 scale 视图。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) return (self.k_storage[storage_layer_index], self.k_scales[storage_layer_index]), (self.v_storage[storage_layer_index], self.v_scales[storage_layer_index]) def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """量化输入浮点 K/V，并将编码数据和分组 scale 写入相同 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index 162dbb378a..658692656b 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -76,14 +76,18 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" - return self.storage[self._logical_to_storage_layer[logical_layer_id]] + return self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """通过 CUDA 算子将形状匹配的 MLA latent 数据写入指定物理 token 位置。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index 2c61f4fe49..f8c94527a9 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -130,9 +130,13 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" - row = self.storage[self._logical_to_storage_layer[logical_layer_id]] + row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] raw = row.view(torch.uint8) kv = raw[..., : self.kv_nbytes] if self.config.fp8_nope_dim else row[..., : self.config.kv_dim] indexer = raw[..., self.indexer_offset : self.indexer_offset + self.config.indexer_dim + 4] @@ -141,7 +145,7 @@ def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """写入 MLA KV 及其量化参数，保留同一 token 行中的 indexer 区域。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index e9855c9a14..ad2e73cdb9 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -169,15 +169,19 @@ def has_logical_layer(self, logical_layer_id: int) -> bool: """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询 KV 和标定参数共用的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """按固定 per-tensor 或 per-head scale 量化输入 K/V，写入 FP8 编码。""" assert token_indices.is_cuda, "token indices must reside on CUDA" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 095ef834e2..1122df3551 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -10,7 +10,6 @@ from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: - from lightllm.common.att_resources.full_att import FullAttentionState from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -76,16 +75,6 @@ def uses_causal_attention(self) -> bool: is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") return not is_parallel_block_draft - def _find_storage_layer_index( - self, k: torch.Tensor, v: torch.Tensor, state: "FullAttentionState" - ) -> int: - # 标定参数按物理层排列；每组只检查一个代表逻辑层，避免共享层重复计数。 - for storage_layer_index, logical_layer_group in enumerate(state.config.logical_layer_groups): - layer_k, layer_v = state.get_layer_view(logical_layer_id=logical_layer_group[0]) - if k.data_ptr() == layer_k.data_ptr() and v.data_ptr() == layer_v.data_ptr(): - return storage_layer_index - raise AssertionError("K/V views do not belong to the same State layer") - @dataclass class AttControl: @@ -114,6 +103,8 @@ class AttControl: linear_att_prefill_dict: Dict = None linear_att_decode: bool = False linear_att_decode_dict: Dict = None + # 静态 FP8 标定按此逻辑层号查询 State 的物理行；其他后端可不使用。 + logical_layer_id: Optional[int] = None @dataclass diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index d88b258d05..4378bcce91 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -72,10 +72,12 @@ def prefill_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, + *, logical_layer_id: int, ) -> torch.Tensor: self.backend: Fp8Fa3AttBackend = self.backend # for typing @@ -89,7 +91,7 @@ def _fp8_prefill_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) - storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) static_q_scales = state.q_scales[storage_layer_index] q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales) q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) @@ -162,6 +164,7 @@ def decode_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_decode_att( @@ -170,6 +173,8 @@ def _fp8_decode_att( k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, + *, + logical_layer_id: int, ): k_head_num = k.shape[1] k_head_dim = k.shape[2] @@ -181,7 +186,7 @@ def _fp8_decode_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) - storage_layer_index = self.backend._find_storage_layer_index(k=cache_k, v=cache_v, state=state) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) q_head_num = q.shape[1] att_batch_size = self.b_att_seq_len.shape[0] diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 174387aa33..41e8e590fa 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -58,10 +58,12 @@ def prefill_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, + *, logical_layer_id: int, ) -> torch.Tensor: o_tensor = alloc_func(q.shape, q.dtype, device="cuda") k = k.unsqueeze(1).view(torch.float8_e4m3fn) @@ -70,7 +72,7 @@ def _fp8_prefill_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", ) - storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) k_descale = state.cpu_scales[storage_layer_index][0] v_descale = state.cpu_scales[storage_layer_index][1] self.prefill_wrapper.run( @@ -109,6 +111,7 @@ def decode_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_decode_att( @@ -117,6 +120,8 @@ def _fp8_decode_att( k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, + *, + logical_layer_id: int, ): o_tensor = alloc_func(q.shape, q.dtype, device="cuda") @@ -126,7 +131,7 @@ def _fp8_decode_att( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", ) - storage_layer_index = self.backend._find_storage_layer_index(k=k, v=v, state=state) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) k_descale = state.cpu_scales[storage_layer_index][0] v_descale = state.cpu_scales[storage_layer_index][1] diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 9b8f81ef27..7e73bf58fe 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -29,11 +29,6 @@ class LinearAttBackend(BaseAttBackend, ABC): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) - role = model.att_resource_role - (self.linear_state,) = ( - state for state in model.att_resource_manager.get_bounded_attention_states(role) - if isinstance(state, GDNState) - ) self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) self.prefill_kernel = self.get_prefill_kernel() @@ -145,7 +140,11 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: LinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) + state = self.infer_state.att_resource_manager.get_att_state( + model_name=self.infer_state.att_resource_role, att_type="gdn", + logical_layer_id=layer_num, state_class=GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -216,16 +215,26 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step == 0: self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1) + return + + # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 + # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 + (state,) = ( + state for state in self.infer_state.att_resource_manager.get_bounded_attention_states( + self.infer_state.att_resource_role + ) + if isinstance(state, GDNState) + ) + if self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: - self._init_fixed_mtp_decode_state(draft_step) + self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int): + def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: ( self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx, @@ -233,12 +242,12 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.backend.linear_state.mtp_state_index, + req_to_mtp_state_index=req_to_mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) - def _init_fixed_mtp_decode_state(self, draft_step: int): + def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = draft_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -255,7 +264,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.backend.linear_state.mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -287,7 +296,11 @@ def decode_att( backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) + state = self.infer_state.att_resource_manager.get_att_state( + model_name=self.infer_state.att_resource_role, att_type="gdn", + logical_layer_id=layer_num, state_class=GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index fc4a59fdee..60cf21e96f 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -38,11 +38,6 @@ class KDALinearAttBackend(BaseAttBackend): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) - role = model.att_resource_role - (self.linear_state,) = ( - state for state in model.att_resource_manager.get_bounded_attention_states(role) - if isinstance(state, KDAState) - ) config = model.config["linear_attn_config"] self.num_heads = config["num_heads"] self.head_dim = config["head_dim"] @@ -89,7 +84,11 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) + state = self.infer_state.att_resource_manager.get_att_state( + model_name=self.infer_state.att_resource_role, att_type="kda", + logical_layer_id=layer_num, state_class=KDAState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -140,16 +139,26 @@ def init_state(self): mtp_step = self.backend.mtp_step if mtp_step == 0: self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1) + return + + # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 + # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 + (state,) = ( + state for state in self.infer_state.att_resource_manager.get_bounded_attention_states( + self.infer_state.att_resource_role + ) + if isinstance(state, KDAState) + ) + if self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: - self._init_fixed_mtp_decode_state(mtp_step) + self._init_fixed_mtp_decode_state(mtp_step, req_to_mtp_state_index=state.mtp_state_index) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int): + def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: ( self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx, @@ -157,12 +166,12 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.backend.linear_state.mtp_state_index, + req_to_mtp_state_index=req_to_mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) - def _init_fixed_mtp_decode_state(self, mtp_step: int): + def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = mtp_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -179,7 +188,7 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.backend.linear_state.mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -210,7 +219,11 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.backend.linear_state.get_layer_view(logical_layer_id=layer_num) + state = self.infer_state.att_resource_manager.get_att_state( + model_name=self.infer_state.att_resource_role, att_type="kda", + logical_layer_id=layer_num, state_class=KDAState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) conv_update = causal_conv1d_update if backend.mtp_step > 0: diff --git a/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py index eca30ec090..c7e9da2bdf 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py @@ -4,58 +4,6 @@ import triton.language as tl -@torch.no_grad() -def destindex_copy_kv_flashmla_fp8( - KV_nope: torch.Tensor, - KV_rope: torch.Tensor, - DestLoc: torch.Tensor, - O_nope: torch.Tensor, - O_scale: torch.Tensor, - O_rope: torch.Tensor, -): - seq_len = DestLoc.shape[0] - kv_nope_head_dim = KV_nope.shape[2] - kv_rope_head_dim = KV_rope.shape[2] - - assert kv_nope_head_dim == 512, f"Expected kv_nope_head_dim=512, got {kv_nope_head_dim}" - assert kv_rope_head_dim == 64, f"Expected kv_rope_head_dim=64, got {kv_rope_head_dim}" - assert O_nope.shape[2] == 512 - assert O_scale.shape[2] == 4 - assert O_rope.shape[2] == 64 - - _fwd_kernel_destindex_copy_kv_flashmla_fp8[(seq_len,)]( - KV_nope, - KV_rope, - DestLoc, - O_nope, - O_scale, - O_rope, - KV_nope.stride(0), - KV_nope.stride(1), - KV_nope.stride(2), - KV_rope.stride(0), - KV_rope.stride(1), - KV_rope.stride(2), - O_nope.stride(0), - O_nope.stride(1), - O_nope.stride(2), - O_scale.stride(0), - O_scale.stride(1), - O_scale.stride(2), - O_rope.stride(0), - O_rope.stride(1), - O_rope.stride(2), - FP8_MIN=torch.finfo(torch.float8_e4m3fn).min, - FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, - BLOCK_DMODEL_NOPE=512, - BLOCK_DMODEL_ROPE=64, - GROUP_SIZE=128, - num_warps=4, - num_stages=1, - ) - return - - @triton.jit def _quant_scale(max_nope: tl.tensor, fp8_max: float): return tl.exp2(tl.ceil(tl.log2(tl.maximum(max_nope / fp8_max, 1e-4)))) @@ -125,3 +73,55 @@ def _fwd_kernel_destindex_copy_kv_flashmla_fp8( o_scale_ptrs = O_scale + dest_index * stride_o_scale_bs + tile_idx * stride_o_scale_d tl.store(o_scale_ptrs, kv_scale.to(tl.float32)) return + + +@torch.no_grad() +def destindex_copy_kv_flashmla_fp8( + KV_nope: torch.Tensor, + KV_rope: torch.Tensor, + DestLoc: torch.Tensor, + O_nope: torch.Tensor, + O_scale: torch.Tensor, + O_rope: torch.Tensor, +): + seq_len = DestLoc.shape[0] + kv_nope_head_dim = KV_nope.shape[2] + kv_rope_head_dim = KV_rope.shape[2] + + assert kv_nope_head_dim == 512, f"Expected kv_nope_head_dim=512, got {kv_nope_head_dim}" + assert kv_rope_head_dim == 64, f"Expected kv_rope_head_dim=64, got {kv_rope_head_dim}" + assert O_nope.shape[2] == 512 + assert O_scale.shape[2] == 4 + assert O_rope.shape[2] == 64 + + _fwd_kernel_destindex_copy_kv_flashmla_fp8[(seq_len,)]( + KV_nope, + KV_rope, + DestLoc, + O_nope, + O_scale, + O_rope, + KV_nope.stride(0), + KV_nope.stride(1), + KV_nope.stride(2), + KV_rope.stride(0), + KV_rope.stride(1), + KV_rope.stride(2), + O_nope.stride(0), + O_nope.stride(1), + O_nope.stride(2), + O_scale.stride(0), + O_scale.stride(1), + O_scale.stride(2), + O_rope.stride(0), + O_rope.stride(1), + O_rope.stride(2), + FP8_MIN=torch.finfo(torch.float8_e4m3fn).min, + FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, + BLOCK_DMODEL_NOPE=512, + BLOCK_DMODEL_ROPE=64, + GROUP_SIZE=128, + num_warps=4, + num_stages=1, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py index bef84c054b..7d5995db40 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py @@ -5,77 +5,6 @@ import triton.language as tl -@torch.no_grad() -def get_prefill_kv_cache_and_remap_indices_triton( - packed_kv: torch.Tensor, - topk_mem_indices: torch.Tensor, - prefill_mem_index: torch.Tensor, - prefill_cache_kv: torch.Tensor, - prefill_dtype: torch.dtype, -): - squeeze_h_kv = topk_mem_indices.ndim == 2 - if squeeze_h_kv: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - - original_shape = topk_mem_indices.shape - flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) - valid_mask = flat_topk != -1 - valid_topk = flat_topk[valid_mask] - table_size = packed_kv.shape[0] - - prefill_row_table = torch.full((table_size,), -1, dtype=torch.int32, device=packed_kv.device) - _build_prefill_row_table_kernel[(prefill_mem_index.numel(),)]( - prefill_mem_index_ptr=prefill_mem_index.to(torch.int32).contiguous(), - row_table_ptr=prefill_row_table, - prefill_token_num=prefill_mem_index.numel(), - num_warps=4, - ) - - unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) - unique_mem_index = unique_mem_index.to(torch.int32) - unique_count = unique_mem_index.numel() - remapped_flat = torch.full_like(flat_topk, -1) - remapped_flat[valid_mask] = inverse.to(torch.int32) - - compact_kv = torch.empty((unique_count, 1, 576), dtype=prefill_dtype, device=packed_kv.device) - packed_nope = packed_kv[:, :, :512].view(torch.float8_e4m3fn).view(-1, 512) - packed_scale = packed_kv[:, :, 512:528].view(torch.float32).view(-1, 4) - packed_rope = packed_kv[:, :, 528:].view(torch.bfloat16).view(-1, 64) - prefill_kv_2d = prefill_cache_kv.view(-1, 576) - compact_kv_2d = compact_kv.view(-1, 576) - - _fill_compact_kv_kernel[(unique_count, 5)]( - packed_nope_ptr=packed_nope, - packed_scale_ptr=packed_scale, - packed_rope_ptr=packed_rope, - unique_mem_index_ptr=unique_mem_index, - prefill_row_table_ptr=prefill_row_table, - prefill_kv_ptr=prefill_kv_2d, - compact_kv_ptr=compact_kv_2d, - packed_nope_stride_s=packed_nope.stride(0), - packed_nope_stride_d=packed_nope.stride(1), - packed_scale_stride_s=packed_scale.stride(0), - packed_scale_stride_d=packed_scale.stride(1), - packed_rope_stride_s=packed_rope.stride(0), - packed_rope_stride_d=packed_rope.stride(1), - prefill_kv_stride_s=prefill_kv_2d.stride(0), - prefill_kv_stride_d=prefill_kv_2d.stride(1), - compact_kv_stride_s=compact_kv_2d.stride(0), - compact_kv_stride_d=compact_kv_2d.stride(1), - unique_num=unique_count, - KV_NOPE_DIM=512, - KV_ROPE_DIM=64, - GROUP_SIZE=128, - BLOCK_D=128, - num_warps=4, - ) - - remapped = remapped_flat.view(original_shape) - if squeeze_h_kv: - remapped = remapped.squeeze(1) - return compact_kv, remapped - - @triton.jit def _build_prefill_row_table_kernel( prefill_mem_index_ptr: tl.tensor, @@ -175,3 +104,74 @@ def _fill_compact_kv_kernel( value, mask=mask, ) + + +@torch.no_grad() +def get_prefill_kv_cache_and_remap_indices_triton( + packed_kv: torch.Tensor, + topk_mem_indices: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_cache_kv: torch.Tensor, + prefill_dtype: torch.dtype, +): + squeeze_h_kv = topk_mem_indices.ndim == 2 + if squeeze_h_kv: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + + original_shape = topk_mem_indices.shape + flat_topk = topk_mem_indices.reshape(-1).contiguous().to(torch.int32) + valid_mask = flat_topk != -1 + valid_topk = flat_topk[valid_mask] + table_size = packed_kv.shape[0] + + prefill_row_table = torch.full((table_size,), -1, dtype=torch.int32, device=packed_kv.device) + _build_prefill_row_table_kernel[(prefill_mem_index.numel(),)]( + prefill_mem_index_ptr=prefill_mem_index.to(torch.int32).contiguous(), + row_table_ptr=prefill_row_table, + prefill_token_num=prefill_mem_index.numel(), + num_warps=4, + ) + + unique_mem_index, inverse = torch.unique(valid_topk, sorted=False, return_inverse=True) + unique_mem_index = unique_mem_index.to(torch.int32) + unique_count = unique_mem_index.numel() + remapped_flat = torch.full_like(flat_topk, -1) + remapped_flat[valid_mask] = inverse.to(torch.int32) + + compact_kv = torch.empty((unique_count, 1, 576), dtype=prefill_dtype, device=packed_kv.device) + packed_nope = packed_kv[:, :, :512].view(torch.float8_e4m3fn).view(-1, 512) + packed_scale = packed_kv[:, :, 512:528].view(torch.float32).view(-1, 4) + packed_rope = packed_kv[:, :, 528:].view(torch.bfloat16).view(-1, 64) + prefill_kv_2d = prefill_cache_kv.view(-1, 576) + compact_kv_2d = compact_kv.view(-1, 576) + + _fill_compact_kv_kernel[(unique_count, 5)]( + packed_nope_ptr=packed_nope, + packed_scale_ptr=packed_scale, + packed_rope_ptr=packed_rope, + unique_mem_index_ptr=unique_mem_index, + prefill_row_table_ptr=prefill_row_table, + prefill_kv_ptr=prefill_kv_2d, + compact_kv_ptr=compact_kv_2d, + packed_nope_stride_s=packed_nope.stride(0), + packed_nope_stride_d=packed_nope.stride(1), + packed_scale_stride_s=packed_scale.stride(0), + packed_scale_stride_d=packed_scale.stride(1), + packed_rope_stride_s=packed_rope.stride(0), + packed_rope_stride_d=packed_rope.stride(1), + prefill_kv_stride_s=prefill_kv_2d.stride(0), + prefill_kv_stride_d=prefill_kv_2d.stride(1), + compact_kv_stride_s=compact_kv_2d.stride(0), + compact_kv_stride_d=compact_kv_2d.stride(1), + unique_num=unique_count, + KV_NOPE_DIM=512, + KV_ROPE_DIM=64, + GROUP_SIZE=128, + BLOCK_D=128, + num_warps=4, + ) + + remapped = remapped_flat.view(original_shape) + if squeeze_h_kv: + remapped = remapped.squeeze(1) + return compact_kv, remapped diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index b8a27cc459..2c73be470c 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -179,8 +179,8 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): def _att_control(self): if self.is_sliding and self.sliding_window_ > 0: w = self.sliding_window_ - 1 - return AttControl(use_sliding_window=True, sliding_window=(w, 0)) - return AttControl(use_sliding_window=False, sliding_window=(-1, -1)) + return AttControl(use_sliding_window=True, sliding_window=(w, 0), logical_layer_id=self.layer_num_) + return AttControl(use_sliding_window=False, sliding_window=(-1, -1), logical_layer_id=self.layer_num_) def _get_layer_kv(self, infer_state: InferStateInfo): # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 19a6a8c688..d9fca23aff 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,4 +1,5 @@ from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.basemodel.attention import AttControl import torch import triton import torch.distributed as dist @@ -57,6 +58,7 @@ def _context_attention_kernel( q=_q, k=_k, v=_v, + att_control=AttControl(logical_layer_id=self.layer_num_), alloc_func=self.alloc_tensor, ) o_tensor = o_tensor.view(q.shape) @@ -76,7 +78,9 @@ def _token_attention_kernel( ) _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor) + o_tensor = infer_state.decode_att_state.decode_att( + q=_q, k=_k, v=_v, att_control=AttControl(logical_layer_id=self.layer_num_), alloc_func=self.alloc_tensor, + ) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index e1873b308c..b4001b4044 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -33,6 +33,7 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): k=k, v=v, att_control=AttControl( + logical_layer_id=self.layer_num_, use_sliding_window=use_sliding_window, sliding_window=(window, window) if use_sliding_window else (-1, -1), ), diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index 9817e23265..f280477a92 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -14,7 +14,6 @@ from lightllm.common.att_resources import AttResourceManager from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.basemodel.attention.base_att import BaseAttBackend from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( @@ -70,7 +69,7 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, actual = actual.view(torch.float8_e4m3fn).float() * state.scales[1][None, :, None] # Backends index fixed scales by physical row, even when model layer # IDs are sparse or include appended draft layers. - assert BaseAttBackend._find_storage_layer_index(None, k, v, state) == 1 + assert state.get_storage_layer_index(logical_layer_id=8) == 1 assert state.cpu_scales.shape == (2, 4) assert torch.equal(actual, kv) diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 876267519a..bffbe941d6 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -9,7 +9,6 @@ from lightllm.common.att_resources import AttResourceManager from lightllm.common.att_resources.full_att import FullAttentionConfig -from lightllm.common.basemodel.attention.base_att import BaseAttBackend from lightllm.models.gemma4.model import Gemma4TpPartModel from unit_tests.common.att_resources.model_helpers import ( full_state_test_buffers, @@ -50,9 +49,7 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - assert views[2] == views[3] assert views[0] != views[2] - if kind in ("gqa", "static_head", "static_tensor"): - k, v = state.get_layer_view(logical_layer_id=11) - assert BaseAttBackend._find_storage_layer_index(None, k, v, state) == 1 + assert [state.get_storage_layer_index(logical_layer_id=layer) for layer in (0, 4, 8, 11)] == [0, 0, 1, 1] @pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index cf448e0704..4fe82759ab 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -6,6 +6,8 @@ import pytest import torch +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.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 @@ -74,7 +76,8 @@ def test_decode_state_initializes_normal_layout(): assert state.b_num_accepted_tokens is None -def test_decode_state_initializes_fixed_mtp_layout(): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_decode_state_initializes_fixed_mtp_layout(role: str): b_req_idx = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) state, infer_state = _create_decode_state( @@ -82,6 +85,7 @@ def test_decode_state_initializes_fixed_mtp_layout(): dynamic_layout=False, b_req_idx=b_req_idx, req_to_mtp_state_index=req_to_mtp_state_index, + role=role, ) state.init_state() @@ -95,7 +99,8 @@ def test_decode_state_initializes_fixed_mtp_layout(): ) -def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): b_req_idx = torch.tensor([2, 2, 4, 4, 4], dtype=torch.int32) req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) expected_cu_q_seq_len = torch.tensor([0, 2, 5, 5, 5, 5], dtype=torch.int32) @@ -113,6 +118,7 @@ def build_params(**kwargs): dynamic_layout=True, b_req_idx=b_req_idx, req_to_mtp_state_index=req_to_mtp_state_index, + role=role, ) state.init_state() @@ -338,7 +344,7 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ assert validate_calls == ["flashqla2", "flashqla3"] -def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None): +def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None, role="main"): model = SimpleNamespace( is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), @@ -350,10 +356,23 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st infer_state.req_manager = SimpleNamespace( HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, ) + infer_state.att_resource_role = role + if req_to_mtp_state_index is not None: + request_capacity = req_to_mtp_state_index.numel() + resources = AttResourceManager(request_capacity=request_capacity) + config = ConvSSMStateConfig( + logical_layer_ids=(3,), global_k_heads=1, global_v_heads=1, + key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, + ) + for model_role in ("main", "draft0"): + linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") + linear.mtp_state_index = req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) + resources.set_att_state(model_role, linear) + resources.finish_state_creation() + infer_state.att_resource_manager = resources state = gdn.LinearAttDecodeAttState( backend=SimpleNamespace( model=model, - linear_state=SimpleNamespace(mtp_state_index=req_to_mtp_state_index), uses_dynamic_spec_verify_layout=lambda: dynamic_layout, ), _infer_state_weak_ref=weakref.ref(infer_state), diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py new file mode 100644 index 0000000000..a87b88e503 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -0,0 +1,104 @@ +"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" + +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 +from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.test_capacity import definition + + +@pytest.mark.parametrize("backend_name", ["fa3", "flashinfer"]) +@pytest.mark.parametrize("operation", ["prefill", "decode"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("logical_layer_id,storage_layer_index", [(3, 0), (7, 0), (8, 1), (11, 1)]) +def test_fp8_backend_selects_calibration_from_logical_layer( + monkeypatch: pytest.MonkeyPatch, + backend_name: str, + operation: str, + role: str, + logical_layer_id: int, + storage_layer_index: int, +) -> None: + template = definition("static_head" if backend_name == "fa3" else "static_tensor") + state = rebuild_full_state( + template, config=replace(template.config, logical_layer_groups=((3, 7), (8, 11))), + token_capacity=8, device="cpu", + ) + state.scales[0].fill_(2) + state.scales[1].fill_(5) + expected_scale = 2 if storage_layer_index == 0 else 5 + manager = AttResourceManager(request_capacity=1) + manager.set_att_state(role, state) + manager.finish_state_creation() + infer = LlamaInferStateInfo() + infer.att_resource_manager = manager + infer.att_resource_role = role + infer.b_seq_len = torch.tensor([1], dtype=torch.int32) + infer.max_q_seq_len = 1 + calls = [] + + # 仅替换外部计算内核，保留推理层、公开 attention 接口及 State 标定查询链路。 + if backend_name == "fa3": + state.q_scales[0].fill_(3) + state.q_scales[1].fill_(7) + expected_q_scale = 3 if storage_layer_index == 0 else 7 + + def quantize(q: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + assert torch.all(scales == expected_q_scale) + return q + + def attend(**kwargs: object) -> torch.Tensor: + assert torch.all(kwargs["k_descale"] == expected_scale) + assert torch.all(kwargs["v_descale"] == expected_scale) + assert torch.all(kwargs["q_descale"] == expected_q_scale) + calls.append(operation) + return kwargs["q"] + + monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", quantize) + monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) + state_class = fa3_fp8.Fp8Fa3PrefillAttState if operation == "prefill" else fa3_fp8.Fp8Fa3DecodeAttState + else: + state_class = ( + flashinfer_fp8.Fp8FlashInferPrefillAttState if operation == "prefill" + else flashinfer_fp8.Fp8FlashInferDecodeAttState + ) + + attention = state_class(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer)) + if backend_name == "fa3": + attention.k_descale = state.scales[:, :2].unsqueeze(1) + attention.v_descale = state.scales[:, 2:].unsqueeze(1) + attention.page_table = attention.cu_seqlens_q = attention.cu_seqlens_k = None + attention.b_att_seq_len = infer.b_seq_len + attention.decode_max_q_seq_len = 1 + attention.causal = True + else: + def run(q: torch.Tensor, kv: tuple[torch.Tensor, torch.Tensor], *, + k_scale: float, v_scale: float, out: torch.Tensor) -> None: + assert k_scale == v_scale == expected_scale + calls.append(operation) + out.copy_(q) + + attention.prefill_wrapper = attention.decode_wrapper = SimpleNamespace(run=run) + + infer.prefill_att_state = infer.decode_att_state = attention + layer = LlamaTransformerLayerInfer.__new__(LlamaTransformerLayerInfer) + layer.layer_num_ = logical_layer_id + layer.tp_q_head_num_ = layer.tp_k_head_num_ = state.config.head_num + layer.head_dim_ = state.config.head_dim + layer.alloc_tensor = lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu") + q = torch.ones((1, state.config.head_num, state.config.head_dim)) + if operation == "prefill": + output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) + else: + output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) + assert torch.equal(output, q) + assert calls == [operation] diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index ef3879a991..d2466bc6d1 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -2,6 +2,8 @@ import pytest import torch +from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -124,7 +126,8 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl @@ -191,16 +194,26 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): uses_dynamic_spec_verify_layout=lambda: dynamic, split_qkv=lambda x: x.split(hidden, -1), ) - backend.linear_state = SimpleNamespace( - mtp_state_index=accepted_offsets, - get_layer_view=lambda logical_layer_id: (conv, ssm), + resources = AttResourceManager(request_capacity=4) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), global_k_heads=heads, global_v_heads=heads, + key_dim=dim, value_dim=dim, conv_width=3, conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, ) + for model_role in ("main", "draft0"): + linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") + linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) + if model_role == role: + linear.conv_state, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) + resources.set_att_state(model_role, linear) + resources.finish_state_creation() manager = SimpleNamespace(HOLD_REQUEST_ID=3) infer = InferStateInfo() infer.batch_size = tokens infer.b_req_idx = torch.tensor(request_rows, device="cuda", dtype=torch.int32) infer.b_mtp_index = torch.tensor(offsets, device="cuda", dtype=torch.int32) infer.req_manager = manager + infer.att_resource_manager = resources + infer.att_resource_role = role weight = SimpleNamespace( get_merged_kda_conv_weight=lambda: conv_weight, linear_A_log=SimpleNamespace(weight=a), From 41d676c0563548e4192cee0e07f7133d446f9053 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 08:31:18 +0000 Subject: [PATCH 12/56] style: standardize attention resource signatures and keyword calls --- docs/cache_manager_design.md | 2 + .../allocators/full_att_kv_allocator.py | 16 +- .../allocators/state_page_allocator.py | 2 +- .../att_resources/att_resource_manager.py | 80 ++-- .../att_resources/att_tail/base_state.py | 16 +- .../att_tail/create_kpool_state.py | 8 +- .../att_resources/att_tail/kpool_state.py | 58 ++- .../att_resources/att_tail/transfer_utils.py | 1 - .../att_resources/bounded_att/base_state.py | 21 +- .../bounded_att/create_gdn_state.py | 8 +- .../bounded_att/create_kda_state.py | 8 +- .../bounded_att/gdn_pages_pool.py | 1 - .../att_resources/bounded_att/gdn_state.py | 121 ++++- .../bounded_att/kda_pages_pool.py | 1 - .../att_resources/bounded_att/kda_state.py | 121 ++++- .../bounded_att/transfer_utils.py | 1 - .../common/att_resources/bounded_att/utils.py | 15 +- .../att_resources/full_att/base_state.py | 5 - .../full_att/create_gqa_state.py | 1 - .../full_att/create_mla_state.py | 1 - .../att_resources/full_att/gqa_state.py | 25 +- .../att_resources/full_att/gqa_transfer.py | 1 - .../full_att/group_quantized_state.py | 25 +- .../full_att/group_quantized_transfer.py | 1 - .../att_resources/full_att/mla_state.py | 25 +- .../att_resources/full_att/mla_transfer.py | 1 - .../full_att/packed_mla_state.py | 25 +- .../full_att/packed_mla_transfer.py | 1 - .../full_att/static_fp8_state.py | 26 +- .../full_att/static_fp8_transfer.py | 1 - .../common/basemodel/attention/fa3/fp8.py | 27 +- .../basemodel/attention/flashinfer/fp8.py | 19 +- lightllm/common/basemodel/basemodel_utils.py | 2 +- lightllm/models/gemma4/model.py | 36 +- lightllm/models/glm5_next/config.py | 2 +- lightllm/models/glm5_next/model.py | 53 ++- lightllm/models/glm5_next_mtp/model.py | 45 +- lightllm/models/qwen3next/config.py | 1 - lightllm/models/qwen3next/model.py | 30 +- .../dp_backend/dp_shared_kv_trans.py | 5 +- .../mode_backend/multi_level_kv_cache.py | 7 +- .../mode_backend/pd/resource_io.py | 5 +- .../model_infer/shared_att_resources_utils.py | 3 +- lightllm/utils/cpu_cache_layout.py | 7 +- test/att_resources/validate_model_cache.py | 60 ++- test/cpu_cache_kernel/test_speed.py | 38 +- .../common/att_resources/manager_fixtures.py | 14 +- .../common/att_resources/model_helpers.py | 23 +- .../att_resources/test_backend_startup.py | 8 +- .../att_resources/test_bounded_state.py | 138 +++++- .../common/att_resources/test_capacity.py | 431 +++++++++++++++--- .../att_resources/test_compute_binding.py | 62 ++- .../att_resources/test_fp8_compute_binding.py | 46 +- .../att_resources/test_full_att_cpu_kernel.py | 126 ++++- .../att_resources/test_full_att_transfer.py | 61 ++- .../common/att_resources/test_full_state.py | 98 +++- .../test_full_storage_initialization.py | 24 +- .../att_resources/test_gqa_pd_kernel.py | 75 ++- .../common/att_resources/test_gqa_transfer.py | 7 +- .../common/att_resources/test_layer_groups.py | 51 ++- .../test_layer_state_selection.py | 17 +- .../att_resources/test_manager_helpers.py | 109 +++-- .../att_resources/test_model_calibration.py | 37 +- .../test_model_initialization.py | 14 +- .../att_resources/test_model_resources.py | 73 ++- .../test_model_storage_definition.py | 51 ++- .../att_resources/test_mtp_state_commit.py | 8 +- .../att_resources/test_network_transport.py | 52 ++- .../att_resources/test_no_attention_draft.py | 7 +- .../att_resources/test_page_pool_storage.py | 22 +- .../test_quantized_full_state.py | 122 ++++- .../test_quantized_paged_storage.py | 8 +- .../att_resources/test_role_transfers.py | 43 +- .../att_resources/test_runtime_transport.py | 161 ++++++- .../att_resources/test_state_creation.py | 84 +++- .../att_resources/test_state_identity.py | 40 +- .../att_resources/test_state_queries.py | 7 +- .../common/att_resources/test_tail_state.py | 55 ++- .../common/att_resources/test_tp_config.py | 58 ++- .../attention/test_fp8_logical_layer.py | 22 +- .../common/basemodel/test_overlap_utils.py | 7 +- unit_tests/common/test_req_manager_page.py | 29 +- unit_tests/models/glm5_next/test_cache.py | 82 +++- unit_tests/models/glm5_next/test_pd_cache.py | 158 +++++-- .../mode_backend/test_dp_big_pages.py | 45 +- .../mode_backend/test_nccl_control_channel.py | 7 +- .../test_paged_side_allocations.py | 41 +- .../model_infer/test_cpu_resource_layout.py | 93 +++- .../test_hybrid_cpu_state_cache.py | 60 ++- .../model_infer/test_multimodal_chunking.py | 23 +- 90 files changed, 2907 insertions(+), 750 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 01a5013a07..ad81d08514 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -31,6 +31,8 @@ class AttResourceManager: 新增 Python 文件、类和接口按 **公共函数在前、私有辅助函数在后** 排列。已有基础流程文件以保留功能、减少 review 差异为优先，保留原有函数相对顺序，不为满足分组规范整体重排。 +新增和改动较多的文件不使用函数声明中的独立 `*` 分隔符。函数参数不少于 4 个时（不计 `self`、`cls`），声明每行一个参数；调用处不少于 4 个实参时，使用 `参数名=值` 并逐行排列。外部 API 的位置参数约束及 `*args`、`**kwargs` 保持其原有语义；改动较少的存量文件不作整文件格式调整。 + - 模块中先定义公共函数，再定义以下划线开头的私有辅助函数（双下划线协议函数除外）。文件名以 `_` 开头不改变文件内函数的排列规则。脚本的 `if __name__ == "__main__"` 执行入口仍放在全部定义之后。 - 类的方法区依次排列：`__new__`、`__init__`、其他公共方法与属性、私有辅助方法。`__getstate__`、`__post_init__` 等 Python 双下划线协议方法属于公共协议，不作为私有辅助方法处理。 - `async def`、`classmethod`、`staticmethod`、`property`、抽象方法均按方法名分类。属性 getter/setter 等存在定义依赖的方法保持必要顺序，装饰器和关联注释随函数一起移动。 diff --git a/lightllm/common/att_resources/allocators/full_att_kv_allocator.py b/lightllm/common/att_resources/allocators/full_att_kv_allocator.py index 8bb62d5fa6..545ed600b6 100644 --- a/lightllm/common/att_resources/allocators/full_att_kv_allocator.py +++ b/lightllm/common/att_resources/allocators/full_att_kv_allocator.py @@ -19,15 +19,25 @@ class FullAttKVAllocator: - def __init__(self, size: int, page_size: int = 1, *, pin_memory: bool = True): + def __init__(self, size: int, page_size: int = 1, pin_memory: bool = True): assert size >= 0 and page_size > 0 and size % page_size == 0, "KV capacity must contain complete pages" assert size + page_size - 1 <= np.iinfo(np.int32).max, "KV token indices including HOLD must fit in int32" self.size = size self.page_size = page_size - self.mem_state = torch.arange(self.size, dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self.mem_state = torch.arange( + end=self.size, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) self.mark_start = 0 self.mark_end = self.size - self._mem_state_return = torch.empty(self.size * 3, dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self._mem_state_return = torch.empty( + size=(self.size * 3,), + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) self._return_start = 0 self.can_use_mem_size = self.size rank_in_node = get_current_rank_in_node() diff --git a/lightllm/common/att_resources/allocators/state_page_allocator.py b/lightllm/common/att_resources/allocators/state_page_allocator.py index cfc58ff9c2..24bc53bbec 100644 --- a/lightllm/common/att_resources/allocators/state_page_allocator.py +++ b/lightllm/common/att_resources/allocators/state_page_allocator.py @@ -7,7 +7,7 @@ class StatePageAllocator: - def __init__(self, size: int, *, reserved_slots: int = 0): + def __init__(self, size: int, reserved_slots: int = 0): self.size = size self.reserved_slots = reserved_slots if not 0 <= self.reserved_slots <= self.size: diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/att_resources/att_resource_manager.py index c72822694e..a27a940e11 100644 --- a/lightllm/common/att_resources/att_resource_manager.py +++ b/lightllm/common/att_resources/att_resource_manager.py @@ -84,7 +84,6 @@ def get_att_state( model_name: str, att_type: str, logical_layer_id: int, - *, state_class: type[FullAttentionState | BoundedAttentionState | AttentionTailState] | None = None, ) -> FullAttentionState | BoundedAttentionState | AttentionTailState: """按模型角色、算法/方案字符串和逻辑层选择唯一 State，缓存查询结果。 @@ -161,7 +160,6 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: def initialize_resources( self, args: StartArgs | EasyDict, - *, token_capacity: int | None, mem_fraction: float, device: str | torch.device | int = "cuda", @@ -212,8 +210,8 @@ def initialize_resources( # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num self.initialize_storage( - token_capacity, - args.page_size, + token_capacity=token_capacity, + page_size=args.page_size, big_page_tokens=big_page_tokens, small_page_capacity=args.linear_att_cache_size or 0, device=device, @@ -229,7 +227,6 @@ def initialize_storage( self, token_capacity: int, page_size: int, - *, big_page_tokens: int, small_page_capacity: int = 0, device: str | torch.device | int = "cuda", @@ -390,7 +387,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: # Bounded attention 状态页：供 radix/CPU cache 复用；未启用对应页池时操作报错。 # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 - def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """将 prefill 运行态保存到全部参与 State 的同类页池、同一组槽位。""" bounded_states = self.get_bounded_attention_states() assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" @@ -401,7 +398,7 @@ def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, for state in bounded_states: state.save_bounded_att_pages(req_indices, page_kind=page_kind, slot_ids=slot_ids) - def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """将调用方选定的页快照恢复到对应请求槽位的运行态。""" bounded_states = self.get_bounded_attention_states() assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" @@ -418,7 +415,6 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, - *, role: str, ) -> None: """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" @@ -485,7 +481,6 @@ def write_req_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layouts: Mapping[str, StateTransferLayout], cpu_page_readies: List[bool] | None = None, big_page_slot_ids: List[int] | None = None, @@ -516,8 +511,11 @@ def write_req_to_cpu( # 阶段一：统一写出本批 full-attention token 数据。 for state in full_states: state.write_to_cpu( - token_indices, cpu_page_buffers, cpu_page_indices, - layout=layouts[state.resource_id], cpu_page_readies=cpu_page_readies, + token_indices=token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + layout=layouts[state.resource_id], + cpu_page_readies=cpu_page_readies, ) # 阶段二：分别从大小页池取快照，写入对应 CPU 页面，不修改请求运行态。 @@ -525,13 +523,21 @@ def write_req_to_cpu( layout = layouts[state.resource_id] if big_page_slot_ids: state.write_page_state_to_cpu( - page_kind="big", slot_ids=big_page_slot_ids, cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=big_cpu_page_indices, layout=layout, cpu_page_readies=big_page_readies, + page_kind="big", + slot_ids=big_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=big_cpu_page_indices, + layout=layout, + cpu_page_readies=big_page_readies, ) if small_page_slot_ids: state.write_page_state_to_cpu( - page_kind="small", slot_ids=small_page_slot_ids, cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=small_cpu_page_indices, layout=layout, cpu_page_readies=small_page_readies, + page_kind="small", + slot_ids=small_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=small_cpu_page_indices, + layout=layout, + cpu_page_readies=small_page_readies, ) def read_req_from_cpu( @@ -539,7 +545,6 @@ def read_req_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layouts: Mapping[str, StateTransferLayout], big_page_slot_ids: List[int] | None = None, req_idx: int | None = None, @@ -567,12 +572,19 @@ def read_req_from_cpu( # 阶段一：恢复 full-attention token 数据；各 State 负责对应布局的搬运约束。 for state in full_states: - state.read_from_cpu(cpu_page_buffers, cpu_page_indices, token_indices, layout=layouts[state.resource_id]) + state.read_from_cpu( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + token_indices=token_indices, + layout=layouts[state.resource_id], + ) # 阶段二：完整页和末尾碎页统一读入大页池，碎页使用追加的 LOAD 槽位。 for state in bounded_states: state.read_page_state_from_cpu( - cpu_page_buffers, cpu_page_indices, big_page_slot_ids=load_big_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + big_page_slot_ids=load_big_page_slot_ids, layout=layouts[state.resource_id], ) @@ -688,7 +700,6 @@ def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTrans def write_to_pd_buffer( self, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, kind: Literal["kv", "att_state"], layouts: Mapping[str, StateTransferLayout], @@ -703,7 +714,10 @@ def write_to_pd_buffer( assert token_indices is not None, "KV PD pages require token indices" for state in self.get_full_attention_states(): state.write_to_pd_buffer( - token_indices, pd_page_buffers, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + token_indices=token_indices, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + layout=layouts[state.resource_id], ) case "att_state": @@ -716,13 +730,19 @@ def write_to_pd_buffer( # 阶段一：写出 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): state.write_to_pd_buffer( - req_idx, pd_page_buffers, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + req_idx=req_idx, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + layout=layouts[state.resource_id], ) # 阶段二：按前缀长度写出 tail 逻辑窗口，适配两端各自的 MTP ring。 for state in tail_states: state.write_to_pd_buffer( - req_idx, pd_page_buffers, pd_page_index=pd_page_index, state_token_len=state_token_len, + req_idx=req_idx, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + state_token_len=state_token_len, layout=layouts[state.resource_id], ) @@ -732,7 +752,6 @@ def write_to_pd_buffer( def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, kind: Literal["kv", "att_state"], layouts: Mapping[str, StateTransferLayout], @@ -747,7 +766,10 @@ def read_from_pd_buffer( assert token_indices is not None, "KV PD pages require token indices" for state in self.get_full_attention_states(): state.read_from_pd_buffer( - pd_page_buffers, token_indices, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + pd_page_buffers=pd_page_buffers, + token_indices=token_indices, + pd_page_index=pd_page_index, + layout=layouts[state.resource_id], ) case "att_state": @@ -760,13 +782,19 @@ def read_from_pd_buffer( # 阶段一：恢复 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): state.read_from_pd_buffer( - pd_page_buffers, req_idx, pd_page_index=pd_page_index, layout=layouts[state.resource_id] + pd_page_buffers=pd_page_buffers, + req_idx=req_idx, + pd_page_index=pd_page_index, + layout=layouts[state.resource_id], ) # 阶段二：按前缀长度恢复 tail 逻辑窗口，适配两端各自的 MTP ring。 for state in tail_states: state.read_from_pd_buffer( - pd_page_buffers, req_idx, pd_page_index=pd_page_index, state_token_len=state_token_len, + pd_page_buffers=pd_page_buffers, + req_idx=req_idx, + pd_page_index=pd_page_index, + state_token_len=state_token_len, layout=layouts[state.resource_id], ) diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/att_resources/att_tail/base_state.py index 2127c98c3f..abf90c9c42 100644 --- a/lightllm/common/att_resources/att_tail/base_state.py +++ b/lightllm/common/att_resources/att_tail/base_state.py @@ -69,16 +69,24 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: @abstractmethod def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: """按 state_token_len 对应的 prefill 终点导出逻辑窗口，不导出推测候选历史。""" ... @abstractmethod def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, - *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: """按 state_token_len 将逻辑窗口恢复到本地运行态，适配本地 MTP 容量。""" ... diff --git a/lightllm/common/att_resources/att_tail/create_kpool_state.py b/lightllm/common/att_resources/att_tail/create_kpool_state.py index 28e87d3235..b4158e3785 100644 --- a/lightllm/common/att_resources/att_tail/create_kpool_state.py +++ b/lightllm/common/att_resources/att_tail/create_kpool_state.py @@ -11,8 +11,12 @@ def create_kpool_state( config: KPoolStateConfig, request_capacity: int, mtp_step: int = 0, - *, device: str | torch.device = "cuda", ) -> KPoolAttentionTailState: """按池化配置和请求容量创建 K-pool 尾部运行态，包含 MTP 额外槽位。""" - return KPoolAttentionTailState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) + return KPoolAttentionTailState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py index 69c2fcd8f1..2c1f7560b6 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -53,7 +53,6 @@ def __init__( config: KPoolStateConfig, request_capacity: int, mtp_step: int = 0, - *, device: str | torch.device = "cuda", ) -> None: """为各层、各请求创建原始 key/gate ring，包含 MTP 所需额外历史。""" @@ -119,30 +118,62 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: axes=("layer", "window", "channel"), shape=(len(config.logical_layer_ids), window_size, channels), axis_slices=( - AxisSlice(axis="window", src_start=0, dst_start=0, length=window_size), - AxisSlice(axis="channel", src_start=0, dst_start=0, length=channels), + AxisSlice( + axis="window", + src_start=0, + dst_start=0, + length=window_size, + ), + AxisSlice( + axis="channel", + src_start=0, + dst_start=0, + length=channels, + ), ), write_enabled=config.tp_rank == 0, ) return (tail_region,) def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: """将 prefill 末尾的逻辑窗口写入 PD 页；state_token_len 为已执行的前缀长度。""" - copies = self._pd_buffer_copies(pd_page_buffers, pd_page_index, req_idx, state_token_len, layout, write=True) + copies = self._pd_buffer_copies( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + req_idx=req_idx, + state_token_len=state_token_len, + layout=layout, + write=True, + ) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, - *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: """按前缀长度将逻辑窗口放回本地 ring，清零其余候选历史，允许两端 MTP 步数不同。""" # TODO：PD 恢复可按目标 ring 槽位一次写入有效历史或零，融合清零和逐窗口复制； # 避免同一 kernel 内先清零再恢复造成竞争，同时清除窗口外的旧候选历史。 - copies = self._pd_buffer_copies(pd_page_buffers, pd_page_index, req_idx, state_token_len, layout, write=False) + copies = self._pd_buffer_copies( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + req_idx=req_idx, + state_token_len=state_token_len, + layout=layout, + write=False, + ) self.tail_buffer[:, req_idx].fill_(0) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) @@ -158,7 +189,6 @@ def _pd_buffer_copies( req_idx: int, state_token_len: int, layout: StateTransferLayout, - *, write: bool, ) -> list[tuple[torch.Tensor, torch.Tensor]]: """为 PD prefill 交接生成复制视图，允许 P/D 使用不同的 MTP 步数。 @@ -197,6 +227,10 @@ def _pd_buffer_copies( assert state_token_len >= 0, "PD tail requires a nonnegative prefix length" # helper 按 L 和本地 ring 实际容量逐位置取视图，不创建索引 Tensor 或读取 GPU 游标。 return att_tail_buffer_copies( - self.tail_buffer[:, req_idx], pd_page_buffers, layout, - page_index=pd_page_index, state_token_len=state_token_len, write=write, + tail=self.tail_buffer[:, req_idx], + page_buffers=pd_page_buffers, + layout=layout, + page_index=pd_page_index, + state_token_len=state_token_len, + write=write, ) diff --git a/lightllm/common/att_resources/att_tail/transfer_utils.py b/lightllm/common/att_resources/att_tail/transfer_utils.py index 85eb47fc61..c77174cb2e 100644 --- a/lightllm/common/att_resources/att_tail/transfer_utils.py +++ b/lightllm/common/att_resources/att_tail/transfer_utils.py @@ -9,7 +9,6 @@ def att_tail_buffer_copies( tail: torch.Tensor, page_buffers: torch.Tensor, layout: StateTransferLayout, - *, page_index: int, state_token_len: int, write: bool, diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/att_resources/bounded_att/base_state.py index 6b47f39e04..b529184bec 100644 --- a/lightllm/common/att_resources/bounded_att/base_state.py +++ b/lightllm/common/att_resources/bounded_att/base_state.py @@ -45,7 +45,6 @@ def initialize_page_pools( self, big_page_capacity: int, small_page_capacity: int = 0, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: @@ -72,12 +71,12 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... @abstractmethod - def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """将 prefill 完成后的基准运行态保存到指定页槽位，前缀关联由调用方维护。""" ... @abstractmethod - def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """从调用方选定的页槽位恢复运行态和控制信息，保留源快照数据。""" ... @@ -110,7 +109,6 @@ def write_page_state_to_cpu( slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -122,7 +120,6 @@ def read_page_state_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: @@ -167,8 +164,11 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: @abstractmethod def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将指定槽位的 prefill 基准运行态写入 PD 页。 @@ -178,8 +178,11 @@ def write_to_pd_buffer( @abstractmethod def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, - *, pd_page_index: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" ... diff --git a/lightllm/common/att_resources/bounded_att/create_gdn_state.py b/lightllm/common/att_resources/bounded_att/create_gdn_state.py index 83ed3db150..33375c7372 100644 --- a/lightllm/common/att_resources/bounded_att/create_gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/create_gdn_state.py @@ -12,8 +12,12 @@ def create_gdn_state( config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, - *, device: str | torch.device = "cuda", ) -> GDNState: """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" - return GDNState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) + return GDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/att_resources/bounded_att/create_kda_state.py b/lightllm/common/att_resources/bounded_att/create_kda_state.py index 989ecce54c..c1934c666b 100644 --- a/lightllm/common/att_resources/bounded_att/create_kda_state.py +++ b/lightllm/common/att_resources/bounded_att/create_kda_state.py @@ -12,8 +12,12 @@ def create_kda_state( config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, - *, device: str | torch.device = "cuda", ) -> KDAState: """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" - return KDAState(config=config, request_capacity=request_capacity, mtp_step=mtp_step, device=device) + return KDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py index 1b1c0745c8..12095f12fc 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py +++ b/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py @@ -14,7 +14,6 @@ def __init__( self, size: int, config: ConvSSMStateConfig, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py index 917b106583..1cedb23e59 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -32,7 +32,10 @@ class GDNState(BoundedAttentionState): # ======================================================================== def __init__( - self, config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, *, + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, device: str | torch.device = "cuda", ) -> None: """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" @@ -71,7 +74,6 @@ def initialize_page_pools( self, big_page_capacity: int, small_page_capacity: int = 0, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: @@ -79,12 +81,22 @@ def initialize_page_pools( assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" self.big_pages_pool = ( - GDNPagesPool(size=big_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + GDNPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) if big_page_capacity else None ) self.small_pages_pool = ( - GDNPagesPool(size=small_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + GDNPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) if small_page_capacity else None ) @@ -112,7 +124,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" @@ -125,7 +137,7 @@ def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) - def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" @@ -167,7 +179,6 @@ def write_page_state_to_cpu( slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -189,7 +200,13 @@ def write_page_state_to_cpu( if cpu_page_readies is not None and cpu_page_readies[index]: continue fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=True)) + copies.extend(bounded_att_buffer_copies( + fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_index=page, + write=True, + )) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) @@ -197,7 +214,6 @@ def read_page_state_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: @@ -214,7 +230,13 @@ def read_page_state_from_cpu( # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 for page, slot in zip(cpu_page_indices, big_page_slot_ids): fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=False)) + copies.extend(bounded_att_buffer_copies( + fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_index=page, + write=False, + )) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) @@ -281,28 +303,46 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: return self._get_transfer_regions() def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 - copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=True) + copies = bounded_att_buffer_copies( + fields=fields, + page_buffers=pd_page_buffers, + layout=layout, + page_index=pd_page_index, + write=True, + ) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, - *, pd_page_index: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 - copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=False) + copies = bounded_att_buffer_copies( + fields=fields, + page_buffers=pd_page_buffers, + layout=layout, + page_index=pd_page_index, + write=False, + ) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) if self.mtp_state_index is not None: @@ -323,7 +363,10 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.resource_id, config.tp_rank, config.tp_world_size, conv_shape, + self.resource_id, + config.tp_rank, + config.tp_world_size, + conv_shape, ) # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 @@ -354,15 +397,25 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dst_start=target_channel_start, length=local_channel_count, ), - AxisSlice(axis="window", src_start=0, dst_start=0, length=config.conv_width), + AxisSlice( + axis="window", + src_start=0, + dst_start=0, + length=config.conv_width, + ), ), write_enabled=write_enabled, ) regions.append(conv_region) logger.debug( "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", - self.resource_id, name, local_channel_start, target_channel_start, - local_channel_count, replica_count, write_enabled, + self.resource_id, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, ) local_channel_start += local_channel_count global_channel_start += global_head_count * head_dim @@ -382,16 +435,36 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, axis_slices=( - AxisSlice(axis="head", src_start=0, dst_start=head_start, length=config.num_v_heads), - AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=config.key_dim), - AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=config.value_dim), + AxisSlice( + axis="head", + src_start=0, + dst_start=head_start, + length=config.num_v_heads, + ), + AxisSlice( + axis="key_dim", + src_start=0, + dst_start=0, + length=config.key_dim, + ), + AxisSlice( + axis="value_dim", + src_start=0, + dst_start=0, + length=config.value_dim, + ), ), write_enabled=write_enabled, ) regions.append(ssm_region) logger.debug( "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.resource_id, head_start, config.num_v_heads, write_enabled, ssm_region.byte_start, ssm_region.byte_end, + self.resource_id, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, ) return tuple(regions) diff --git a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py index e70d52e376..9ac96b1661 100644 --- a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py +++ b/lightllm/common/att_resources/bounded_att/kda_pages_pool.py @@ -14,7 +14,6 @@ def __init__( self, size: int, config: ConvSSMStateConfig, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py index 42c85ba1b7..d92c8bf7d0 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -32,7 +32,10 @@ class KDAState(BoundedAttentionState): # ======================================================================== def __init__( - self, config: ConvSSMStateConfig, request_capacity: int, mtp_step: int = 0, *, + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, device: str | torch.device = "cuda", ) -> None: """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" @@ -71,7 +74,6 @@ def initialize_page_pools( self, big_page_capacity: int, small_page_capacity: int = 0, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> None: @@ -79,12 +81,22 @@ def initialize_page_pools( assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" self.big_pages_pool = ( - KDAPagesPool(size=big_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + KDAPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) if big_page_capacity else None ) self.small_pages_pool = ( - KDAPagesPool(size=small_page_capacity, config=self.config, device=device, pin_memory=pin_memory) + KDAPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) if small_page_capacity else None ) @@ -112,7 +124,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" @@ -125,7 +137,7 @@ def save_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) - def restore_bounded_att_pages(self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int]) -> None: + def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" @@ -167,7 +179,6 @@ def write_page_state_to_cpu( slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -189,7 +200,13 @@ def write_page_state_to_cpu( if cpu_page_readies is not None and cpu_page_readies[index]: continue fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=True)) + copies.extend(bounded_att_buffer_copies( + fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_index=page, + write=True, + )) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) @@ -197,7 +214,6 @@ def read_page_state_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: @@ -214,7 +230,13 @@ def read_page_state_from_cpu( # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 for page, slot in zip(cpu_page_indices, big_page_slot_ids): fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies(fields, cpu_page_buffers, layout, page_index=page, write=False)) + copies.extend(bounded_att_buffer_copies( + fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_index=page, + write=False, + )) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) @@ -281,28 +303,46 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: return self._get_transfer_regions() def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 - copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=True) + copies = bounded_att_buffer_copies( + fields=fields, + page_buffers=pd_page_buffers, + layout=layout, + page_index=pd_page_index, + write=True, + ) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, - *, pd_page_index: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + layout: StateTransferLayout, ) -> None: """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 - copies = bounded_att_buffer_copies(fields, pd_page_buffers, layout, page_index=pd_page_index, write=False) + copies = bounded_att_buffer_copies( + fields=fields, + page_buffers=pd_page_buffers, + layout=layout, + page_index=pd_page_index, + write=False, + ) for src_tensor, dst_tensor in copies: dst_tensor.copy_(src_tensor, non_blocking=True) if self.mtp_state_index is not None: @@ -323,7 +363,10 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.resource_id, config.tp_rank, config.tp_world_size, conv_shape, + self.resource_id, + config.tp_rank, + config.tp_world_size, + conv_shape, ) # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 @@ -354,15 +397,25 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dst_start=target_channel_start, length=local_channel_count, ), - AxisSlice(axis="window", src_start=0, dst_start=0, length=config.conv_width), + AxisSlice( + axis="window", + src_start=0, + dst_start=0, + length=config.conv_width, + ), ), write_enabled=write_enabled, ) regions.append(conv_region) logger.debug( "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", - self.resource_id, name, local_channel_start, target_channel_start, - local_channel_count, replica_count, write_enabled, + self.resource_id, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, ) local_channel_start += local_channel_count global_channel_start += global_head_count * head_dim @@ -382,16 +435,36 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, axis_slices=( - AxisSlice(axis="head", src_start=0, dst_start=head_start, length=config.num_v_heads), - AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=config.key_dim), - AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=config.value_dim), + AxisSlice( + axis="head", + src_start=0, + dst_start=head_start, + length=config.num_v_heads, + ), + AxisSlice( + axis="key_dim", + src_start=0, + dst_start=0, + length=config.key_dim, + ), + AxisSlice( + axis="value_dim", + src_start=0, + dst_start=0, + length=config.value_dim, + ), ), write_enabled=write_enabled, ) regions.append(ssm_region) logger.debug( "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.resource_id, head_start, config.num_v_heads, write_enabled, ssm_region.byte_start, ssm_region.byte_end, + self.resource_id, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, ) return tuple(regions) diff --git a/lightllm/common/att_resources/bounded_att/transfer_utils.py b/lightllm/common/att_resources/bounded_att/transfer_utils.py index 61f07a4b1a..0736008dcb 100644 --- a/lightllm/common/att_resources/bounded_att/transfer_utils.py +++ b/lightllm/common/att_resources/bounded_att/transfer_utils.py @@ -15,7 +15,6 @@ def bounded_att_buffer_copies( fields: Mapping[str, TensorField], page_buffers: torch.Tensor, layout: StateTransferLayout, - *, page_index: int, write: bool, ) -> list[tuple[torch.Tensor, torch.Tensor]]: diff --git a/lightllm/common/att_resources/bounded_att/utils.py b/lightllm/common/att_resources/bounded_att/utils.py index ca0fc3fb7b..d8e4c2fbb7 100644 --- a/lightllm/common/att_resources/bounded_att/utils.py +++ b/lightllm/common/att_resources/bounded_att/utils.py @@ -26,7 +26,6 @@ def create_page_pool_tensor( size: tuple[int, ...], dtype: torch.dtype, - *, device: str | torch.device | int = "cpu", pin_memory: bool = True, ) -> torch.Tensor: @@ -43,7 +42,12 @@ def create_page_pool_tensor( # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 if not shared: - tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=pin_memory) + tensor = torch.empty( + size=size, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" logger.info( f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " @@ -94,7 +98,12 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 - cleanup = finalize(storage, _unregister_page_pool_storage, ptr, device_id) + cleanup = finalize( + storage, + _unregister_page_pool_storage, + ptr, + device_id, + ) cleanup.atexit = False logger.info( f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index e15a4232c1..5e89139b44 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -109,7 +109,6 @@ def __init__( self, layout_config: FullAttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """只准备 State 配置，不分配 KV 或 HOLD 页；存储由后续统一容量初始化。""" @@ -176,7 +175,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -189,7 +187,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" @@ -226,7 +223,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -238,7 +234,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: diff --git a/lightllm/common/att_resources/full_att/create_gqa_state.py b/lightllm/common/att_resources/full_att/create_gqa_state.py index c176ae1daa..042e9042e2 100644 --- a/lightllm/common/att_resources/full_att/create_gqa_state.py +++ b/lightllm/common/att_resources/full_att/create_gqa_state.py @@ -14,7 +14,6 @@ def create_gqa_state( config: FullAttentionConfig, - *, kv_type: Literal["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"] = "None", page_size: int = 1, device: str | torch.device = "cuda", diff --git a/lightllm/common/att_resources/full_att/create_mla_state.py b/lightllm/common/att_resources/full_att/create_mla_state.py index e412b12e20..5ebd16c0a8 100644 --- a/lightllm/common/att_resources/full_att/create_mla_state.py +++ b/lightllm/common/att_resources/full_att/create_mla_state.py @@ -14,7 +14,6 @@ def create_mla_state( config: FullAttentionConfig, - *, kv_type: Literal["None", "fp8kv_dsa"] = "None", packed: bool = False, indexer_dim: int = 128, diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index 1253efc78d..ff7ba301d4 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -42,7 +42,6 @@ def __init__( self, layout_config: FullAttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """只记录布局配置，KV 与 HOLD 页由后续统一容量初始化分配。""" @@ -157,7 +156,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -179,7 +177,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" @@ -233,7 +230,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -253,7 +249,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -280,7 +275,10 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + self.resource_id, + config.tp_rank, + config.tp_world_size, + page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 field_specs = ( @@ -335,11 +333,20 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, field_name, head_start, local_head_count, replica_count, - write_enabled, field_region.byte_start, field_region.byte_end, + self.resource_id, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, ) logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + "%s transfer regions built: fields=%d, total_bytes=%d", + self.resource_id, + len(regions), + byte_offset, ) return tuple(regions) diff --git a/lightllm/common/att_resources/full_att/gqa_transfer.py b/lightllm/common/att_resources/full_att/gqa_transfer.py index 3f4537d151..920ef23715 100644 --- a/lightllm/common/att_resources/full_att/gqa_transfer.py +++ b/lightllm/common/att_resources/full_att/gqa_transfer.py @@ -18,7 +18,6 @@ def transfer_gqa_pages( token_indices: torch.Tensor, page_buffers: torch.Tensor, page_indices: list[int], - *, layout: StateTransferLayout, write: bool, grid_num: int = 36, diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index 3dcc961e67..d71c7e0c67 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -52,7 +52,6 @@ def __init__( self, layout_config: GroupQuantizedAttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """按分组量化配置准备 KV 与动态 scale 存储。""" @@ -206,7 +205,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -230,7 +228,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" @@ -298,7 +295,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -320,7 +316,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -349,7 +344,10 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + self.resource_id, + config.tp_rank, + config.tp_world_size, + page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 field_specs = ( @@ -406,11 +404,20 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, field_name, head_start, local_head_count, replica_count, - write_enabled, field_region.byte_start, field_region.byte_end, + self.resource_id, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, ) logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + "%s transfer regions built: fields=%d, total_bytes=%d", + self.resource_id, + len(regions), + byte_offset, ) return tuple(regions) diff --git a/lightllm/common/att_resources/full_att/group_quantized_transfer.py b/lightllm/common/att_resources/full_att/group_quantized_transfer.py index b4376c7235..cb4ec6b4f1 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_transfer.py +++ b/lightllm/common/att_resources/full_att/group_quantized_transfer.py @@ -20,7 +20,6 @@ def transfer_group_quantized_pages( token_indices: torch.Tensor, page_buffers: torch.Tensor, page_indices: list[int], - *, layout: StateTransferLayout, write: bool, grid_num: int = 36, diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index 658692656b..aa49881630 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -32,7 +32,6 @@ def __init__( self, layout_config: FullAttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """按 MLA 布局配置准备压缩 KV 存储。""" @@ -135,7 +134,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -156,7 +154,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" @@ -206,7 +203,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -225,7 +221,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -251,7 +246,10 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + self.resource_id, + config.tp_rank, + config.tp_world_size, + page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 field_specs = (("kv", self.storage_dtype, config.head_dim, "dim"),) @@ -303,11 +301,20 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, field_name, head_start, local_head_count, replica_count, - write_enabled, field_region.byte_start, field_region.byte_end, + self.resource_id, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, ) logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + "%s transfer regions built: fields=%d, total_bytes=%d", + self.resource_id, + len(regions), + byte_offset, ) return tuple(regions) diff --git a/lightllm/common/att_resources/full_att/mla_transfer.py b/lightllm/common/att_resources/full_att/mla_transfer.py index 42145c3825..45f68e56e7 100644 --- a/lightllm/common/att_resources/full_att/mla_transfer.py +++ b/lightllm/common/att_resources/full_att/mla_transfer.py @@ -17,7 +17,6 @@ def transfer_mla_pages( token_indices: torch.Tensor, page_buffers: torch.Tensor, page_indices: list[int], - *, layout: StateTransferLayout, write: bool, grid_num: int = 36, diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index f8c94527a9..0c86ff5784 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -73,7 +73,6 @@ def __init__( self, layout_config: PackedMLAAttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """按 packed MLA 配置准备 KV 与 indexer 存储。""" @@ -208,7 +207,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -232,7 +230,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" @@ -295,7 +292,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -317,7 +313,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -346,7 +341,10 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + self.resource_id, + config.tp_rank, + config.tp_world_size, + page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 if config.fp8_nope_dim: @@ -409,11 +407,20 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, field_name, head_start, local_head_count, replica_count, - write_enabled, field_region.byte_start, field_region.byte_end, + self.resource_id, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, ) logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + "%s transfer regions built: fields=%d, total_bytes=%d", + self.resource_id, + len(regions), + byte_offset, ) return tuple(regions) diff --git a/lightllm/common/att_resources/full_att/packed_mla_transfer.py b/lightllm/common/att_resources/full_att/packed_mla_transfer.py index b0310eaa37..08e757899b 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_transfer.py +++ b/lightllm/common/att_resources/full_att/packed_mla_transfer.py @@ -17,7 +17,6 @@ def transfer_packed_mla_pages( token_indices: torch.Tensor, page_buffers: torch.Tensor, page_indices: list[int], - *, layout: StateTransferLayout, kv_dim: int, indexer_offset: int, diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index ad2e73cdb9..0e61d53bae 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -45,7 +45,6 @@ def __init__( self, layout_config: StaticFP8AttentionConfig, page_size: int = 1, - *, device: str | torch.device = "cuda", ) -> None: """只准备静态 FP8 布局；标定参数由后续 bind_calibration 统一绑定。""" @@ -70,7 +69,6 @@ def __init__( def bind_calibration( self, calibration: Mapping[str, object] | None, - *, storage_layer_start: int, ) -> None: """在分配 KV 前绑定标定表中本 State 的物理层区间，并切出当前 TP 的 head。 @@ -243,7 +241,6 @@ def write_to_cpu( token_indices: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], - *, layout: StateTransferLayout, cpu_page_readies: list[bool] | None = None, ) -> None: @@ -265,7 +262,6 @@ def read_from_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], token_indices: torch.Tensor, - *, layout: StateTransferLayout, ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" @@ -319,7 +315,6 @@ def write_to_pd_buffer( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -339,7 +334,6 @@ def read_from_pd_buffer( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, - *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -366,7 +360,10 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, config.tp_rank, config.tp_world_size, page_tokens, + self.resource_id, + config.tp_rank, + config.tp_world_size, + page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 field_specs = ( @@ -421,11 +418,20 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, field_name, head_start, local_head_count, replica_count, - write_enabled, field_region.byte_start, field_region.byte_end, + self.resource_id, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, ) logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", self.resource_id, len(regions), byte_offset + "%s transfer regions built: fields=%d, total_bytes=%d", + self.resource_id, + len(regions), + byte_offset, ) return tuple(regions) diff --git a/lightllm/common/att_resources/full_att/static_fp8_transfer.py b/lightllm/common/att_resources/full_att/static_fp8_transfer.py index 1674f7d5e0..ef2db76f17 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_transfer.py +++ b/lightllm/common/att_resources/full_att/static_fp8_transfer.py @@ -18,7 +18,6 @@ def transfer_static_fp8_pages( token_indices: torch.Tensor, page_buffers: torch.Tensor, page_indices: list[int], - *, layout: StateTransferLayout, write: bool, grid_num: int = 36, diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index 4378bcce91..78f4e2887b 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -76,8 +76,12 @@ def prefill_att( ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, - *, logical_layer_id: int, + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, ) -> torch.Tensor: self.backend: Fp8Fa3AttBackend = self.backend # for typing @@ -85,8 +89,12 @@ def _fp8_prefill_att( q_head_dim = q.shape[2] k_head_num = k.shape[1] k_head_dim = k.shape[2] - cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) + 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) state = getattr( self.infer_state.att_resource_manager, f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", @@ -172,15 +180,18 @@ def _fp8_decode_att( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - alloc_func=torch.empty, - *, logical_layer_id: int, + alloc_func=torch.empty, ): k_head_num = k.shape[1] k_head_dim = k.shape[2] - cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) + 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) state = getattr( self.infer_state.att_resource_manager, diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 41e8e590fa..e16e87ab97 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -62,8 +62,12 @@ def prefill_att( ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty, - *, logical_layer_id: int, + self, + 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="cuda") k = k.unsqueeze(1).view(torch.float8_e4m3fn) @@ -76,8 +80,8 @@ def _fp8_prefill_att( k_descale = state.cpu_scales[storage_layer_index][0] v_descale = state.cpu_scales[storage_layer_index][1] self.prefill_wrapper.run( - q, - (k, v), + q=q, + paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor, @@ -119,9 +123,8 @@ def _fp8_decode_att( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - alloc_func=torch.empty, - *, logical_layer_id: int, + alloc_func=torch.empty, ): o_tensor = alloc_func(q.shape, q.dtype, device="cuda") @@ -136,8 +139,8 @@ def _fp8_decode_att( k_descale = state.cpu_scales[storage_layer_index][0] v_descale = state.cpu_scales[storage_layer_index][1] self.decode_wrapper.run( - q, - (k, v), + q=q, + paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor, diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index c29f95e64a..34aefaa408 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -10,7 +10,7 @@ def get_resource_config( - model_config: dict[str, object], *, model_class: type[TpPartBaseModel] | None = None, + model_config: dict[str, object], model_class: type[TpPartBaseModel] | None = None, ) -> dict[str, object]: """优先采用模型专有配置转换；默认展开 text_config 并补齐常见字段别名。""" build_config = getattr(model_class, "_resource_config", None) diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 056a63f3d2..3ab394548e 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -88,8 +88,12 @@ def _verify_params(self): @classmethod def _create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state @@ -126,9 +130,12 @@ def _create_full_att_states( global_head_num=per_token_k_width // head_dim, head_dim=head_dim, dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] @@ -204,7 +211,12 @@ def _init_to_get_rotary_gemma4(self): sliding_partial = sliding_params.get("partial_rotary_factor", 1.0) sliding_rot_dim = int(sliding_head_dim * sliding_partial) inv_freq_sliding = 1.0 / ( - sliding_theta ** (torch.arange(0, sliding_rot_dim, 2, dtype=torch.float32) / sliding_rot_dim) + sliding_theta ** (torch.arange( + start=0, + end=sliding_rot_dim, + step=2, + dtype=torch.float32, + ) / sliding_rot_dim) ) freqs_s = torch.outer(t, inv_freq_sliding) self._cos_cached_sliding = torch.cos(freqs_s).to(self.data_type).cuda() @@ -224,7 +236,12 @@ def _init_to_get_rotary_gemma4(self): if rope_type == "proportional": rope_angles = int(full_partial * full_head_dim // 2) inv_freq_rot = 1.0 / ( - full_theta ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float32) / full_head_dim) + full_theta ** (torch.arange( + start=0, + end=2 * rope_angles, + step=2, + dtype=torch.float32, + ) / full_head_dim) ) nope_angles = full_head_dim // 2 - rope_angles if nope_angles > 0: @@ -233,7 +250,12 @@ def _init_to_get_rotary_gemma4(self): inv_freq_full = inv_freq_rot else: full_rot_dim = int(full_head_dim * full_partial) - inv_freq_full = 1.0 / (full_theta ** (torch.arange(0, full_rot_dim, 2, dtype=torch.float32) / full_rot_dim)) + inv_freq_full = 1.0 / (full_theta ** (torch.arange( + start=0, + end=full_rot_dim, + step=2, + dtype=torch.float32, + ) / full_rot_dim)) freqs_f = torch.outer(t, inv_freq_full) self._cos_cached_full = torch.cos(freqs_f).to(self.data_type).cuda() diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 3a66cfc784..40d242f3f4 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -27,7 +27,7 @@ class Glm5NextResourceConfig(Qwen3NextResourceConfig): INDEX_PADDING_BYTES = 144 @classmethod - def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict, *, include_draft: bool = True): + def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict, include_draft: bool = True): config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index cbbcc22f6b..f8ffa16bcb 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -36,9 +36,14 @@ class Glm5NextTpPartModel(TpPartBaseModel): @classmethod def _create_bounded_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig @@ -63,7 +68,9 @@ def _create_bounded_att_states( tp_rank=tp_rank, tp_world_size=tp_world_size, ), - request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, ) return [state] @@ -101,8 +108,12 @@ def _init_some_value(self): @classmethod def _create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state @@ -121,17 +132,26 @@ def _create_full_att_states( global_head_num=1, head_dim=config["kv_lora_rank"], dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + packed=True, + page_size=args.page_size, + device="meta", ) return [state] @classmethod def _create_att_tail_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig @@ -147,10 +167,15 @@ def _create_att_tail_states( state = create_kpool_state( config=KPoolStateConfig( logical_layer_ids=logical_layer_ids, - pool_size=config["index_kpool"], head_dim=config["index_head_dim"], - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + pool_size=config["index_kpool"], + head_dim=config["index_head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, ) return [state] diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index fb2cdc2a11..eb5bd01fd1 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -44,7 +44,12 @@ def _init_weights(self, start_layer_index=None): self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ self.pre_post_weight.main_norm_weight_ = self.main_model.pre_post_weight.final_norm_weight_ self.trans_layers_weight = [ - self.transformer_weight_class(self.config["num_hidden_layers"], self.data_type, self.config, self.quant_cfg) + self.transformer_weight_class( + layer_num=self.config["num_hidden_layers"], + data_type=self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) ] def _init_infer_layer(self, start_layer_index=None): @@ -69,8 +74,12 @@ def _init_some_value(self): @classmethod def _create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state @@ -87,17 +96,26 @@ def _create_full_att_states( global_head_num=1, head_dim=config["kv_lora_rank"], dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + packed=True, + page_size=args.page_size, + device="meta", ) return [state] @classmethod def _create_att_tail_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig @@ -112,10 +130,15 @@ def _create_att_tail_states( state = create_kpool_state( config=KPoolStateConfig( logical_layer_ids=logical_layer_ids, - pool_size=config["index_kpool"], head_dim=config["index_head_dim"], - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + pool_size=config["index_kpool"], + head_dim=config["index_head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, ) return [state] diff --git a/lightllm/models/qwen3next/config.py b/lightllm/models/qwen3next/config.py index 5c78a7cc7c..c5961e44e1 100644 --- a/lightllm/models/qwen3next/config.py +++ b/lightllm/models/qwen3next/config.py @@ -125,7 +125,6 @@ def from_model_config( cls, model_cfg: dict[str, object], args: StartArgs | EasyDict, - *, include_draft: bool = True, ): llm_config = model_cfg.get("text_config", model_cfg) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 18340e6239..1de13e2d4d 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -41,9 +41,14 @@ def __init__(self, kvargs) -> None: @classmethod def _create_bounded_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig @@ -69,7 +74,9 @@ def _create_bounded_att_states( tp_rank=tp_rank, tp_world_size=tp_world_size, ), - request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, ) return [state] @@ -92,8 +99,12 @@ def _init_config(self): @classmethod def _create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state @@ -112,9 +123,12 @@ def _create_full_att_states( global_head_num=cfg.full_att_all_num_kv_heads, head_dim=cfg.full_att_head_dim, dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 00e0c6b00c..0979697e4a 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -81,7 +81,10 @@ def build_shared_kv_trans_tasks( rank_max_radix_cache_lens = np.max(available_lens, axis=1) # 为本 DP 负责的请求选择可用前缀最长的来源。 for req_index, req, max_req_radix_cache_len, req_dp_rank in zip( - list(range(len(reqs))), reqs, rank_max_radix_cache_lens, req_dp_ranks + list(range(len(reqs))), + reqs, + rank_max_radix_cache_lens, + req_dp_ranks, ): # 当前请求是本 dp_rank 负责的 is_current_dp_handle = req_dp_rank == self.dp_rank_in_node diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index dbd48b37f6..728c002dd0 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -138,7 +138,6 @@ def read_req_from_cpu( req: InferReq, cpu_page_indices: List[int], page_lengths: List[int], - *, start_token_len: int, ) -> None: """恢复已选定的请求前缀页面，分配资源并登记完整大页引用。 @@ -179,9 +178,9 @@ def read_req_from_cpu( # 阶段三：统一恢复各类 State；异常直接向上传播，不在此处回收或继续执行。 resources.read_req_from_cpu( - self.cpu_cache_client.cpu_kv_cache_tensor, - cpu_page_indices, - token_indices, + cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, + cpu_page_indices=cpu_page_indices, + token_indices=token_indices, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=big_page_slot_ids, req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index 6a47d6b98f..f805f61f9f 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -11,7 +11,7 @@ from typing import Literal -def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, *, device: str | torch.device | int): +def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, device: str | torch.device | int): if page_count <= 0: raise ValueError("PD buffer must contain pages") return torch.empty((page_count, resources.get_pd_page_bytes()), dtype=torch.uint8, device=device) @@ -20,15 +20,14 @@ def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, *, def copy_page( node_att_managers: Sequence[AttResourceManager], pd_page_buffers: torch.Tensor, - *, pd_page_index: int, dp_rank_in_node: int, tp_world_size: int, token_indices: torch.Tensor, + write_to_pd_buffer: bool, task_kind: Literal["kv", "att_state"] = "kv", req_idx: int | None = None, state_token_len: int | None = None, - write_to_pd_buffer: bool, ) -> None: """选择节点内指定 DP 的 TP 组，检查整组布局后提交 PD 页面读写。 diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py index e031d2e850..88db8df336 100644 --- a/lightllm/server/router/model_infer/shared_att_resources_utils.py +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -20,7 +20,6 @@ def publish_att_resources( resources: AttResourceManager, request_mapping: torch.Tensor, - *, kind: Literal["pd", "dp"], reader_count: int, ) -> None: @@ -56,7 +55,7 @@ def publish_att_resources( def load_att_resources( - rank_in_node: int, *, kind: Literal["pd", "dp"], + rank_in_node: int, kind: Literal["pd", "dp"], ) -> tuple[AttResourceManager, torch.Tensor]: """读取指定用途的资源；DP 接收方随后通过 manager 统一注册共享页池。""" assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index c6813eebc3..5e2b996441 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -41,8 +41,11 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, ) bounded_states = model_class._create_bounded_att_states( - model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, - request_capacity=1, device="meta", + model_config=config, + args=args, + model_logical_layer_ids=logical_layer_ids, + request_capacity=1, + device="meta", ) # tail 不占 CPU 页；无需创建 tail，也不需要实际模型或临时 CUDA 运行态。 for state in (*full_states, *bounded_states): diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 4ca6c08ad9..ab62f5c4a7 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -52,7 +52,12 @@ @torch.inference_mode() def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpPartBaseModel] = ()): - executor = StaticBenchmarkExecutor(args, model, draft_models, TokenSource(args)) + executor = StaticBenchmarkExecutor( + args=args, + model=model, + draft_models=draft_models, + token_source=TokenSource(args), + ) resources = model.att_resource_manager (full_state,) = resources.get_full_attention_states("main") for draft in draft_models: @@ -74,12 +79,22 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP source_req = executor._alloc_req_indexes(1) resources.reset_bounded_att_state(source_req.tolist()) resources.reset_attention_tail_state(source_req.tolist()) - tokens = np.random.default_rng(42).integers(100, 10000, (1, prefix_len), dtype=np.int64) + tokens = np.random.default_rng(42).integers( + low=100, + high=10000, + size=(1, prefix_len), + dtype=np.int64, + ) for page, length in zip(pages, lengths): batch = executor._make_prefill_input(tokens[:, page * page_tokens : length], source_req, page * page_tokens) output = model.forward(batch) if draft_models: - executor._fill_mtp_prefill_kv(None, batch, output, executor._argmax_ids(output.logits)) + executor._fill_mtp_prefill_kv( + case=None, + main_prefill_input=batch, + main_output=output, + first_next_ids=executor._argmax_ids(output.logits), + ) if hybrid: resources.save_bounded_att_pages(source_req.tolist(), page_kind="big", slot_ids=[source_slots[page]]) torch.cuda.synchronize() @@ -88,9 +103,9 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP cpu = torch.zeros((len(pages), resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) resources.write_req_to_cpu( - source_indices, - cpu, - pages, + token_indices=source_indices, + cpu_page_buffers=cpu, + cpu_page_indices=pages, layouts=resources.get_cpu_transfer_layouts(), page_kind="big" if hybrid else None, slot_ids=source_slots, @@ -103,7 +118,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP if runtime_layouts: pd["att_state"] = torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") resources.write_to_pd_buffer( - pd["kv"], + pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", layouts=full_att_layouts, @@ -111,7 +126,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP ) if "att_state" in pd: resources.write_to_pd_buffer( - pd["att_state"], + pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", layouts=runtime_layouts, @@ -137,9 +152,9 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP if hybrid: allocated_slots.extend(slots) resources.read_req_from_cpu( - cpu, - pages, - indices, + cpu_page_buffers=cpu, + cpu_page_indices=pages, + token_indices=indices, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=slots, req_idx=int(request[0]) if getattr(resources, "cpu_tail_restore", False) else None, @@ -157,7 +172,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP ) else: resources.read_from_pd_buffer( - pd["kv"], + pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", layouts=full_att_layouts, @@ -165,7 +180,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP ) if "att_state" in pd: resources.read_from_pd_buffer( - pd["att_state"], + pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", layouts=runtime_layouts, @@ -189,12 +204,12 @@ def decode(request: torch.Tensor): request_rows, mtp_index = executor._build_mtp_decode_index_tensors(request, step_width) lengths = executor._build_mtp_seq_len(torch.tensor([sequence_length], dtype=torch.int32), step_width) batch = executor._make_decode_input( - step_width, - request_rows, - mtp_index, - lengths, - next_token.repeat(step_width), - int(lengths[-1]), + batch_size=step_width, + req_idx=request_rows, + mtp_index=mtp_index, + seq_len=lengths, + input_ids=next_token.repeat(step_width), + max_kv_seq_len=int(lengths[-1]), ) output = model.forward(batch) logits = [output.logits.clone()] @@ -214,7 +229,12 @@ def decode(request: torch.Tensor): for kind, request in targets.items(): actual = decode(request) for restored_logits, reference_logits in zip(actual, reference): - torch.testing.assert_close(restored_logits, reference_logits, rtol=1e-4, atol=1e-4) + torch.testing.assert_close( + actual=restored_logits, + expected=reference_logits, + rtol=1e-4, + atol=1e-4, + ) assert torch.equal(restored_logits.argmax(dim=-1), reference_logits.argmax(dim=-1)) errors[kind].append(float((actual[0] - reference[0]).abs().max())) if draft_models: diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 4245367ead..758b2dab9e 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -55,19 +55,33 @@ def main(): ) tokens = args.tokens_per_page * args.pages full = GQAFullAttentionState( - FullAttentionConfig(tuple((layer,) for layer in tuple(range(12))), 1, 256, torch.bfloat16, 0, 8), + FullAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in tuple(range(12))), + global_head_num=1, + head_dim=256, + dtype=torch.bfloat16, + tp_rank=0, + tp_world_size=8, + ), ) full.initialize_storage(tokens) resources = AttResourceManager(request_capacity=1) bounded_state = create_gdn_state( config=ConvSSMStateConfig( logical_layer_ids=tuple(layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval), - global_k_heads=cfg.global_linear_k_heads, global_v_heads=cfg.global_linear_v_heads, - key_dim=cfg.head_linear_k_dim, value_dim=cfg.head_linear_v_dim, - conv_width=cfg.conv_kernel_size - 1, conv_dtype=cfg.conv_state_dtype, ssm_dtype=cfg.ssm_state_dtype, - tp_rank=0, tp_world_size=cfg.tp_world_size, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=0, + tp_world_size=cfg.tp_world_size, ), - request_capacity=1, mtp_step=0, device="cuda", + request_capacity=1, + mtp_step=0, + device="cuda", ) bounded_state.initialize_page_pools(big_page_capacity=args.pages + 2, small_page_capacity=0) resources.set_att_state("main", full) @@ -87,9 +101,9 @@ def main(): torch.cuda.synchronize() token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") resources.write_req_to_cpu( - token_indices, - cpu_page_buffers, - pages, + token_indices=token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=pages, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=slots, ) @@ -97,9 +111,9 @@ def main(): def restore(): resources.read_req_from_cpu( - cpu_page_buffers, - pages, - token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=pages, + token_indices=token_indices, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=slots, req_idx=0, diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index ed78f7fedd..b82a563c17 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -19,12 +19,14 @@ class MainDraftTestResourceManager(AttResourceManager): - def __init__(self, device: str | torch.device | int = "cpu", *, with_tail: bool = True, tail_pool_size: int = 4): + def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): super().__init__(request_capacity=3) self.set_att_state("main", make_state(device=device)) draft_state = GDNState( - replace(self.main_kda_conv_fp32_ssm_fp32_state.config, global_v_heads=2), - request_capacity=3, mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, device=device, + config=replace(self.main_kda_conv_fp32_ssm_fp32_state.config, global_v_heads=2), + request_capacity=3, + mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, + device=device, ) draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") self.set_att_state("draft0", draft_state) @@ -39,8 +41,10 @@ def __init__(self, device: str | torch.device | int = "cpu", *, with_tail: bool if with_tail: tail = make_tail(device=device) state = KPoolAttentionTailState( - replace(tail.config, pool_size=tail_pool_size), - request_capacity=tail.request_capacity, mtp_step=tail.mtp_step, device=tail.tail_buffer.device, + config=replace(tail.config, pool_size=tail_pool_size), + request_capacity=tail.request_capacity, + mtp_step=tail.mtp_step, + device=tail.tail_buffer.device, ) self.set_att_state("main", state) self.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 463da5008c..9f4625d462 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -52,7 +52,6 @@ def bind_test_fp8_calibration( def rebuild_full_state( template: FullAttentionState, - *, config: FullAttentionConfig | None = None, token_capacity: int | None = None, page_size: int | None = None, @@ -98,7 +97,7 @@ def main_linear_state(resources: AttResourceManager): return state -def prepare_model_states(model: TpPartBaseModel, *, logical_layer_ids: Sequence[int] | None = None): +def prepare_model_states(model: TpPartBaseModel, logical_layer_ids: Sequence[int] | None = None): config = model.config.get("text_config", model.config) if logical_layer_ids is None: logical_layer_ids = range(config.get("num_hidden_layers", config.get("n_layer"))) @@ -126,7 +125,9 @@ def create_full_attention_roles( config = replace(template.config, logical_layer_groups=layers[start:end]) if isinstance(template, StaticFP8FullAttentionState): state = rebuild_full_state( - template, config=config, scales=template.scales[start:end].clone(), + template=template, + config=config, + scales=template.scales[start:end].clone(), q_scales=None if template.q_scales is None else template.q_scales[start:end].clone(), ) else: @@ -191,7 +192,6 @@ def build_hybrid_resources( mtp_step: int, full_state: FullAttentionState, full_allocator: FullAttKVAllocator, - *, big_page_capacity: int, small_page_capacity: int, draft_full_state: FullAttentionState | None = None, @@ -222,7 +222,9 @@ def build_hybrid_resources( tp_rank=tp_rank, tp_world_size=config.tp_world_size, ), - request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + request_capacity=request_capacity, + mtp_step=mtp_step, + device="cuda", ) bounded_state.initialize_page_pools( big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=True, @@ -236,10 +238,15 @@ def build_hybrid_resources( tail = create_kpool_state( config=KPoolStateConfig( logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), - pool_size=config.index_kpool, head_dim=config.index_head_dim, dtype=config.full_att_dtype, - tp_rank=tp_rank, tp_world_size=config.tp_world_size, + pool_size=config.index_kpool, + head_dim=config.index_head_dim, + dtype=config.full_att_dtype, + tp_rank=tp_rank, + tp_world_size=config.tp_world_size, ), - request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + request_capacity=request_capacity, + mtp_step=mtp_step, + device="cuda", ) manager.set_att_state(role, tail) manager.full_att_allocator = full_allocator diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index 68433786d0..ea01afd357 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -246,7 +246,13 @@ def start_runtime(): if via_backend: with pytest.raises(RuntimeReady): backend.init_model( - dict(args=args, world_size=1, load_way="HF", weight_dir=str(tmp_path), max_total_token_num=None) + dict( + args=args, + world_size=1, + load_way="HF", + weight_dir=str(tmp_path), + max_total_token_num=None, + ) ) else: assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 72e6316ffc..3b43e74763 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -24,16 +24,26 @@ def make_state( cls: type[KDAState] | type[GDNState] = KDAState, - *, device: str | torch.device | int = "cpu", heads: int = 1, mtp_step: int = 2, ): config = ConvSSMStateConfig( - logical_layer_ids=(1, 3), global_k_heads=heads, global_v_heads=heads, key_dim=2, value_dim=2, conv_width=2, - conv_dtype=torch.float32, ssm_dtype=torch.float32, + logical_layer_ids=(1, 3), + global_k_heads=heads, + global_v_heads=heads, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + state = cls( + config=config, + request_capacity=3, + mtp_step=mtp_step, + device=device, ) - state = cls(config, request_capacity=3, mtp_step=mtp_step, device=device) manager = AttResourceManager(request_capacity=3) manager.set_att_state("main", state) state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") @@ -61,8 +71,18 @@ def page_layout(state: KDAState | GDNState): axes=("channel", "layer", "window"), shape=conv_shape, axis_slices=( - AxisSlice(axis="channel", src_start=0, dst_start=0, length=cfg.conv_channels), - AxisSlice(axis="window", src_start=0, dst_start=0, length=cfg.conv_width), + AxisSlice( + axis="channel", + src_start=0, + dst_start=0, + length=cfg.conv_channels, + ), + AxisSlice( + axis="window", + src_start=0, + dst_start=0, + length=cfg.conv_width, + ), ), ), TransferRegion( @@ -73,9 +93,24 @@ def page_layout(state: KDAState | GDNState): axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, axis_slices=( - AxisSlice(axis="head", src_start=0, dst_start=0, length=cfg.num_v_heads), - AxisSlice(axis="key_dim", src_start=0, dst_start=0, length=cfg.key_dim), - AxisSlice(axis="value_dim", src_start=0, dst_start=0, length=cfg.value_dim), + AxisSlice( + axis="head", + src_start=0, + dst_start=0, + length=cfg.num_v_heads, + ), + AxisSlice( + axis="key_dim", + src_start=0, + dst_start=0, + length=cfg.key_dim, + ), + AxisSlice( + axis="value_dim", + src_start=0, + dst_start=0, + length=cfg.value_dim, + ), ), ), ), @@ -94,10 +129,21 @@ def test_prefill_state_is_saved_and_restored( if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") config = ConvSSMStateConfig( - logical_layer_ids=(1, 3), global_k_heads=1, global_v_heads=1, key_dim=2, value_dim=2, conv_width=2, - conv_dtype=torch.float32, ssm_dtype=torch.float32, + logical_layer_ids=(1, 3), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + state = cls( + config=config, + request_capacity=3, + mtp_step=2, + device=device, ) - state = cls(config, request_capacity=3, mtp_step=2, device=device) assert state.resource_id is None and not hasattr(config, "resource_id") assert state.big_pages_pool is None and state.small_pages_pool is None conv, ssm = state.get_layer_view(3) @@ -132,12 +178,23 @@ def test_prefill_state_is_saved_and_restored( def test_gdn_dimensions_and_page_restore(): config = replace( - make_state().config, global_k_heads=1, global_v_heads=2, key_dim=2, value_dim=3 + make_state().config, + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + ) + state = GDNState( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", ) - state = GDNState(config, request_capacity=3, mtp_step=2, device="cpu") state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2, pin_memory=False) conv = torch.arange(40, dtype=config.conv_dtype).reshape(2, 10, 2) - ssm = torch.arange(24, dtype=config.ssm_dtype).reshape(2, 2, 2, 3) + ssm = torch.arange(24, dtype=config.ssm_dtype).reshape( + shape=(2, 2, 2, 3), + ) state.conv_state[:, 0, :, :2].copy_(conv) state.ssm_state[:, 0].copy_(ssm) for page_kind in ("big", "small"): @@ -178,16 +235,38 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): source.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) layout = page_layout(source) buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8) - source.write_page_state_to_cpu("big", [0], buffer, [2], layout=layout) + source.write_page_state_to_cpu( + page_kind="big", + slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[2], + layout=layout, + ) assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() target = make_state(cls) - target.read_page_state_from_cpu(buffer, [2], big_page_slot_ids=[1], layout=layout) + target.read_page_state_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + big_page_slot_ids=[1], + layout=layout, + ) target.restore_bounded_att_pages([2], page_kind="big", slot_ids=[1]) assert (target.conv_state[:, 2, :, :2] == 13).all() assert (target.ssm_state[:, 6] == 29).all() # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 - target.read_page_state_from_cpu(buffer, [2], big_page_slot_ids=[0], layout=layout) - target.write_page_state_to_cpu("big", [0], buffer, [0], layout=layout) + target.read_page_state_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + big_page_slot_ids=[0], + layout=layout, + ) + target.write_page_state_to_cpu( + page_kind="big", + slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[0], + layout=layout, + ) assert torch.equal(buffer[0], buffer[2]) @@ -206,11 +285,26 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( source.ssm_state[:, 0].fill_(47) layout = page_layout(source) # PD 只传 Conv/SSM，仍保留非零区域起点以覆盖组合页偏移。 - allocation = torch.full((3, layout.byte_end), 254, dtype=torch.uint8, device=device) + allocation = torch.full( + size=(3, layout.byte_end), + fill_value=254, + dtype=torch.uint8, + device=device, + ) page = allocation[1] - source.write_to_pd_buffer(0, allocation, pd_page_index=1, layout=layout) + source.write_to_pd_buffer( + req_idx=0, + pd_page_buffers=allocation, + pd_page_index=1, + layout=layout, + ) target = make_state(cls, device=device, mtp_step=1) - target.read_from_pd_buffer(allocation, 2, pd_page_index=1, layout=layout) + target.read_from_pd_buffer( + pd_page_buffers=allocation, + req_idx=2, + pd_page_index=1, + layout=layout, + ) assert (target.conv_state[:, 2, :, :2] == 23).all() assert (target.ssm_state[:, 4] == 47).all() assert (allocation[0] == 254).all() and (allocation[2] == 254).all() diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index d032d2fabd..4aea99aa7b 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -41,7 +41,12 @@ def definition(kind: str): - config = FullAttentionConfig(((0,), (2,)), 2, 16, torch.bfloat16) + config = FullAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + ) if kind in ("int4", "int8"): return GroupQuantizedFullAttentionState( GroupQuantizedAttentionConfig(**vars(config), bits=4 if kind == "int4" else 8), @@ -64,7 +69,10 @@ def definition(kind: str): fp8 = kind == "packed_fp8" kv_dim = 512 if kind == "packed_glm" else 576 config = replace( - config, global_head_num=1, head_dim=800 if fp8 else kv_dim + 72, dtype=torch.uint8 if fp8 else config.dtype + config, + global_head_num=1, + head_dim=800 if fp8 else kv_dim + 72, + dtype=torch.uint8 if fp8 else config.dtype, ) return PackedMLAFullAttentionState( PackedMLAAttentionConfig(**vars(config), kv_dim=kv_dim, fp8_nope_dim=512 if fp8 else 0), page_size=4, device="meta" @@ -74,7 +82,7 @@ def definition(kind: str): return GQAFullAttentionState(config, page_size=4, device="meta") -def allocated_bytes(state: FullAttentionState, *, token_only: bool = False): +def allocated_bytes(state: FullAttentionState, token_only: bool = False): buffers = list(full_state_test_buffers(state).values()) if not token_only and isinstance(state, StaticFP8FullAttentionState): buffers.append(state.scales) @@ -96,7 +104,12 @@ def test_budget_matches_token_stores_including_hold_and_padding( pd_args: EasyDict, ): template = definition(kind) - state = rebuild_full_state(template, config=template.config, token_capacity=20, device="cpu") + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=20, + device="cpu", + ) actual = allocated_bytes(state, token_only=True) planning = AttResourceManager(request_capacity=1) planning.set_att_state("main", rebuild_full_state(template, config=template.config)) @@ -118,7 +131,12 @@ def test_budget_matches_token_stores_including_hold_and_padding( manager.finish_state_creation() (state,) = manager.get_full_attention_states("main") assert manager.get_pd_page_bytes() == pd_page_bytes - args = StartArgs(page_size=4, run_mode="prefill", pd_kv_page_size=8, pd_kv_page_num=3) + args = StartArgs( + page_size=4, + run_mode="prefill", + pd_kv_page_size=8, + pd_kv_page_num=3, + ) reserve = manager.get_pd_page_bytes() * args.pd_kv_page_num monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) @@ -133,12 +151,27 @@ def test_budget_matches_token_stores_including_hold_and_padding( resources.finish_state_creation() if expected is None: with pytest.raises(MemoryError, match="HOLD"): - resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + resources.initialize_resources( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) elif expected == 0: with pytest.raises(AssertionError, match="at least one complete page"): - resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + resources.initialize_resources( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) else: - resources.initialize_resources(args, token_capacity=None, mem_fraction=1.0, device="cpu") + resources.initialize_resources( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) assert resources.full_att_allocator.size == expected assert resources.get_pd_page_bytes() == manager.get_pd_page_bytes() @@ -151,7 +184,13 @@ def test_shared_allocator_accounts_for_distinct_state_stores(): manager.finish_state_creation() expected = sum( allocated_bytes( - rebuild_full_state(state, config=state.config, token_capacity=12, page_size=4, device="cpu"), + rebuild_full_state( + template=state, + config=state.config, + token_capacity=12, + page_size=4, + device="cpu", + ), token_only=True, ) for state in (main, draft) @@ -166,15 +205,32 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( pd_args: EasyDict, ): config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=2, global_v_heads=2 if state_cls is KDAState else 4, key_dim=3, - value_dim=3 if state_cls is KDAState else 5, conv_width=1, conv_dtype=torch.bfloat16, + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=2 if state_cls is KDAState else 4, + key_dim=3, + value_dim=3 if state_cls is KDAState else 5, + conv_width=1, + conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) tail = KPoolAttentionTailState( - KPoolStateConfig(logical_layer_ids=(1,), pool_size=4, head_dim=3, dtype=torch.bfloat16), - request_capacity=3, mtp_step=3, device="cpu", + config=KPoolStateConfig( + logical_layer_ids=(1,), + pool_size=4, + head_dim=3, + dtype=torch.bfloat16, + ), + request_capacity=3, + mtp_step=3, + device="cpu", + ) + reference = state_cls( + config=config, + request_capacity=3, + mtp_step=3, + device="cpu", ) - reference = state_cls(config, request_capacity=3, mtp_step=3, device="cpu") conv = reference.conv_state[:, 0, :, : config.conv_width] ssm = reference.ssm_state[:, 0] bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() @@ -191,11 +247,25 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8) layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "tail"} - manager.write_to_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + manager.write_to_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) reference.conv_state.zero_() reference.ssm_state.zero_() tail.tail_buffer.zero_() - manager.read_from_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + manager.read_from_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) assert (reference.conv_state[:, 0, :, : config.conv_width] == 2).all() assert (reference.ssm_state[:, 0] == 3).all() expected_tail = torch.zeros_like(tail.tail_buffer[:, 0]) @@ -222,7 +292,12 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( (GQAFullAttentionState, "main", 3, torch.bfloat16), (GQAFullAttentionState, "main", 1, torch.float32), ): - config = FullAttentionConfig(((0,),), 1, dim, dtype) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=dim, + dtype=dtype, + ) manager.set_att_state(role, state_cls(config, page_size=1, device="meta")) manager.finish_state_creation() # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 @@ -230,7 +305,12 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage(4, 1, big_page_tokens=4, device="cuda") + manager.initialize_storage( + token_capacity=4, + page_size=1, + big_page_tokens=4, + device="cuda", + ) assert manager.get_pd_page_bytes() == 72 states = manager.get_full_attention_states() for i, state in enumerate(states, 1): @@ -240,12 +320,24 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( page = pages[1] layouts = manager.get_pd_full_att_transfer_layouts() assert [min(region.byte_start for region in layouts[state.resource_id].regions) for state in states] == [0, 8, 48] - manager.write_to_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + manager.write_to_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) assert (page[:6].view(torch.bfloat16) == 1).all() for state in states: for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() - manager.read_from_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) + manager.read_from_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) for i, state in enumerate(states, 1): assert (read_kv_storage(state)[:, :3] == i).all() assert not pages[0].any() @@ -254,7 +346,10 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( @pytest.mark.parametrize("layout_first", [False, True]) @pytest.mark.parametrize("transfer", ["cpu", "pd"]) def test_capacity_uses_quantized_field_padding_and_preserves_page_boundaries( - monkeypatch: pytest.MonkeyPatch, pd_args: EasyDict, layout_first: bool, transfer: str + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + layout_first: bool, + transfer: str, ): if not torch.cuda.is_available(): pytest.skip("CUDA required") @@ -269,8 +364,12 @@ def _get_transfer_regions(self, page_tokens: int): pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 config = GroupQuantizedAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=128, - dtype=torch.float32, bits=8, group_size=128, + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + dtype=torch.float32, + bits=8, + group_size=128, ) manager = AttResourceManager(request_capacity=1) prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") @@ -295,7 +394,11 @@ def _get_transfer_regions(self, page_tokens: int): layouts = get_layouts() assert get_layouts() is layouts pages = torch.full( - (3, 1600), 91, dtype=torch.uint8, device="cpu" if transfer == "cpu" else "cuda", pin_memory=transfer == "cpu" + size=(3, 1600), + fill_value=91, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", ) for index, state in enumerate(manager.get_full_attention_states()): layout = layouts[state.resource_id] @@ -306,11 +409,33 @@ def _get_transfer_regions(self, page_tokens: int): field.tensor.zero_() field.tensor[:, :3].fill_(index * 4 + field_index + 1) if transfer == "cpu": - manager.write_req_to_cpu(torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), pages, [1], layouts=layouts) - manager.read_req_from_cpu(pages, [1], torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), layouts=layouts) + manager.write_req_to_cpu( + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + ) + manager.read_req_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + layouts=layouts, + ) else: - manager.write_to_pd_buffer(pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda")) - manager.read_from_pd_buffer(pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda")) + manager.write_to_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) + manager.read_from_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + ) for state in manager.get_full_attention_states(): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): torch.testing.assert_close(field.tensor[:, 4:7], field.tensor[:, :3]) @@ -336,12 +461,23 @@ def get_pd_transfer_regions(self): manager = AttResourceManager(request_capacity=1) if full_page_tokens is not None: pd_args.pd_kv_page_size = full_page_tokens - config = FullAttentionConfig(((0,),), 1, 1, torch.float32) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) state = PaddedGDNState( ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, - conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, ), request_capacity=1, device="cpu", ) @@ -365,10 +501,24 @@ def get_pd_transfer_regions(self): pages = torch.full((3, expected), 91, dtype=torch.uint8) state.conv_state.fill_(3) state.ssm_state.fill_(5) - manager.write_to_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + manager.write_to_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) state.conv_state.zero_() state.ssm_state.zero_() - manager.read_from_pd_buffer(pages, pd_page_index=1, kind="att_state", layouts=layouts, req_idx=0, state_token_len=8) + manager.read_from_pd_buffer( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) assert (state.conv_state == 3).all() and (state.ssm_state == 5).all() assert (pages[[0, 2]] == 91).all() and (pages[1, 80:] == 91).all() assert len(warnings) == int(full_page_tokens == 8) @@ -398,7 +548,12 @@ def _get_transfer_regions(self, page_tokens: int): pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig(((0,),), 1, 3, torch.bfloat16) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + dtype=torch.bfloat16, + ) manager.set_att_state("main", PaddedState(config, page_size=1, device="meta")) manager.finish_state_creation() get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes @@ -406,7 +561,12 @@ def _get_transfer_regions(self, page_tokens: int): monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage(8, 4, big_page_tokens=8, device="cpu") + manager.initialize_storage( + token_capacity=8, + page_size=4, + big_page_tokens=8, + device="cpu", + ) assert get_page_bytes() == (80 if transfer == "cpu" else 72) layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() assert layouts["main_gqa_bf16_state"].byte_end == 72 @@ -420,20 +580,38 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( pd_args: EasyDict, ): manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig(((0,),), 1, 1, torch.float32) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) full_state = GQAFullAttentionState(config, page_size=1, device="meta") manager.set_att_state("main", full_state) for cls, role in ((KDAState, "main"), (GDNState, "draft0")): config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, - conv_dtype=torch.float32, ssm_dtype=torch.float32, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, ) state = cls(config, request_capacity=1, device="meta") manager.set_att_state(role, state) if with_tail: state = KPoolAttentionTailState( - KPoolStateConfig(logical_layer_ids=(0,), pool_size=4, head_dim=1, dtype=torch.float32), - request_capacity=1, mtp_step=0, device="meta", + config=KPoolStateConfig( + logical_layer_ids=(0,), + pool_size=4, + head_dim=1, + dtype=torch.float32, + ), + request_capacity=1, + mtp_step=0, + device="meta", ) manager.set_att_state("main", state) # 两份运行态各占 16 字节数据，tail 另占 32 字节，无页头及长度、版本字段。 @@ -453,7 +631,11 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("page_tokens", [1, 3]) def test_cpu_page_bytes_derive_from_cached_transfer_layouts( - monkeypatch: pytest.MonkeyPatch, cpu_args: EasyDict, kind: str, tp_world_size: int, page_tokens: int + monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + kind: str, + tp_world_size: int, + page_tokens: int, ): cpu_args.cpu_cache_token_page_size = page_tokens template = definition(kind) @@ -462,15 +644,32 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None manager = AttResourceManager(request_capacity=3) for role in ("main", "draft0"): - manager.set_att_state(role, rebuild_full_state(template, config=config, scales=scales, q_scales=q_scales)) + manager.set_att_state(role, rebuild_full_state( + template=template, + config=config, + scales=scales, + q_scales=q_scales, + )) for state_cls, v_heads, value_dim in ((KDAState, 1, 1), (GDNState, 2, 3)): # 奇数长度的 BF16 Conv 与 FP32 SSM 之间需要补齐，MTP 候选区不应计入页大小。 bounded_config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=v_heads, key_dim=1, value_dim=value_dim, conv_width=1, - conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, tp_rank=tp_world_size - 1, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=v_heads, + key_dim=1, + value_dim=value_dim, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + tp_rank=tp_world_size - 1, tp_world_size=tp_world_size, ) - bounded = state_cls(bounded_config, request_capacity=3, mtp_step=2, device="meta") + bounded = state_cls( + config=bounded_config, + request_capacity=3, + mtp_step=2, + device="meta", + ) manager.set_att_state("main", bounded) manager.finish_state_creation() # CPU 容量和布局都只依赖 State 的 CPU 字段描述，不读取 GPU 存储开销或 PD 布局。 @@ -503,8 +702,14 @@ def get_cpu_transfer_regions(self): manager = AttResourceManager(request_capacity=1) state = PaddedGDNState( ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, - conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, ), request_capacity=1, device="meta", ) @@ -525,7 +730,12 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea states = [] for encoding in ("gqa", kind): template = definition(encoding) - state = rebuild_full_state(template, config=template.config, token_capacity=8, device="cuda") + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device="cuda", + ) for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.view(torch.uint8).fill_(211) states.append(state) @@ -546,12 +756,27 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): field.tensor[:, :3].fill_(value) field.tensor[:, 4:7].zero_() - pages = torch.full((3, manager.get_cpu_page_bytes()), 165, dtype=torch.uint8, pin_memory=True) - manager.write_req_to_cpu(torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), pages, [1], layouts=layouts) + pages = torch.full( + size=(3, manager.get_cpu_page_bytes()), + fill_value=165, + dtype=torch.uint8, + pin_memory=True, + ) + manager.write_req_to_cpu( + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + ) torch.cuda.synchronize() assert (pages[1, data_end:gqa_start] == 165).all() assert (pages[[0, 2]] == 165).all() - manager.read_req_from_cpu(pages, [1], torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), layouts=layouts) + manager.read_req_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + layouts=layouts, + ) for value, state in enumerate((packed, gqa), 3): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): assert (field.tensor[:, 4:7] == value).all() @@ -565,10 +790,21 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas manager.bounded_att_small_page_allocator = StatePageAllocator(1) for state_cls, role, value_dim in ((KDAState, "draft0", 1), (GDNState, "main", 3)): config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=value_dim, conv_width=1, - conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=value_dim, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = state_cls( + config=config, + request_capacity=1, + mtp_step=2, + device="cpu", ) - state = state_cls(config, request_capacity=1, mtp_step=2, device="cpu") state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=False) manager.set_att_state(role, state) manager.finish_state_creation() @@ -585,14 +821,23 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas manager.save_bounded_att_pages([0], page_kind="big", slot_ids=source_slots) pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8) manager.write_req_to_cpu( - torch.tensor([], dtype=torch.long), pages, [1], layouts=layouts, big_page_slot_ids=source_slots + token_indices=torch.tensor([], dtype=torch.long), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + big_page_slot_ids=source_slots, ) assert (pages[0] == 87).all() assert (pages[1, 12:16] == 87).all() manager.reset_bounded_att_state([0]) target_slots = manager.bounded_att_big_page_allocator.alloc(1) manager.read_req_from_cpu( - pages, [1], torch.tensor([], dtype=torch.long), layouts=layouts, big_page_slot_ids=target_slots, req_idx=0 + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([], dtype=torch.long), + layouts=layouts, + big_page_slot_ids=target_slots, + req_idx=0, ) for i, state in enumerate(states): assert (state.conv_state[:, 0, :, :1] == i + 3).all() @@ -608,12 +853,23 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation pd_args.pd_kv_page_size = 3 pd_args.cpu_cache_token_page_size = 3 manager = AttResourceManager(request_capacity=1) - config = FullAttentionConfig(((0,),), 1, 3, torch.bfloat16) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + dtype=torch.bfloat16, + ) manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) manager.set_att_state("main", GDNState( ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=1, global_v_heads=1, key_dim=1, value_dim=1, conv_width=1, - conv_dtype=torch.bfloat16, ssm_dtype=torch.bfloat16, + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.bfloat16, ), request_capacity=1, device="cpu", )) @@ -626,7 +882,12 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage(8, 4, big_page_tokens=8, device="cpu") + manager.initialize_storage( + token_capacity=8, + page_size=4, + big_page_tokens=8, + device="cpu", + ) after = getter() after_layout = after[state.resource_id] assert after is not before and getter() is after @@ -660,7 +921,10 @@ def get_cpu_transfer_regions(self): pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 config = FullAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=2, dtype=torch.float32 + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=2, + dtype=torch.float32, ) manager = AttResourceManager(request_capacity=1) prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") @@ -687,14 +951,49 @@ def get_cpu_transfer_regions(self): page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() pages = torch.zeros( - (1, page_bytes), dtype=torch.uint8, device="cpu" if transfer == "cpu" else "cuda", pin_memory=transfer == "cpu" + size=(1, page_bytes), + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", ) if transfer == "cpu": - manager.write_req_to_cpu(torch.arange(4, dtype=torch.long, device="cuda"), pages, [0], layouts=layouts) - manager.read_req_from_cpu(pages, [0], torch.arange(4, 8, dtype=torch.long, device="cuda"), layouts=layouts) + manager.write_req_to_cpu( + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[0], + layouts=layouts, + ) + manager.read_req_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=torch.arange( + start=4, + end=8, + dtype=torch.long, + device="cuda", + ), + layouts=layouts, + ) else: - manager.write_to_pd_buffer(pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, dtype=torch.long, device="cuda")) - manager.read_from_pd_buffer(pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, 8, dtype=torch.long, device="cuda")) + manager.write_to_pd_buffer( + pd_page_buffers=pages, + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + ) + manager.read_from_pd_buffer( + pd_page_buffers=pages, + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.arange( + start=4, + end=8, + dtype=torch.long, + device="cuda", + ), + ) for state in manager.get_full_attention_states(): torch.testing.assert_close(read_kv_storage(state)[:, 4:8], read_kv_storage(state)[:, :4]) assert not read_kv_storage(state)[:, 8:].any() diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index f280477a92..f4fb55fed7 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -30,7 +30,12 @@ @pytest.mark.parametrize("capture", [False, True]) def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, capture: bool): template = definition(kind) - state = rebuild_full_state(template, config=replace(template.config, logical_layer_groups=((3,), (8,))), token_capacity=8, device="cuda") + state = rebuild_full_state( + template=template, + config=replace(template.config, logical_layer_groups=((3,), (8,))), + token_capacity=8, + device="cuda", + ) for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() manager = AttResourceManager(request_capacity=2) @@ -45,13 +50,28 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, infer.mem_index = torch.tensor([1, 5], dtype=torch.int32, device="cuda") heads = 1 if kind == "packed_glm" else 2 * state.config.head_num dim = state.config.kv_dim if kind == "packed_glm" else state.config.head_dim - kv = torch.full((2, heads, dim), 0.5, dtype=torch.bfloat16, device="cuda") + kv = torch.full( + size=(2, heads, dim), + fill_value=0.5, + dtype=torch.bfloat16, + device="cuda", + ) layer = SimpleNamespace(layer_num_=8) - TransformerLayerInferTpl._post_cache_kv(layer, kv, infer, None) + TransformerLayerInferTpl._post_cache_kv( + self=layer, + cache_kv=kv, + infer_state=infer, + layer_weight=None, + ) if capture: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - TransformerLayerInferTpl._post_cache_kv(layer, kv, infer, None) + TransformerLayerInferTpl._post_cache_kv( + self=layer, + cache_kv=kv, + infer_state=infer, + layer_weight=None, + ) for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() graph.replay() @@ -76,13 +96,27 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, def test_packed_fp8_prefill_combines_history_and_fresh_kv_with_sparse_remapping(): template = definition("packed_fp8") - state = rebuild_full_state(template, config=template.config, token_capacity=8, device="cuda") + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device="cuda", + ) for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() history_ids = torch.tensor([3, 7], dtype=torch.int32, device="cuda") - history = torch.randn(2, 1, 576, dtype=torch.bfloat16, device="cuda") + history = torch.randn( + size=(2, 1, 576), + dtype=torch.bfloat16, + device="cuda", + ) state.write_full_att_kv(0, history_ids, history) - fresh = torch.full((1, 1, 576), 123, dtype=torch.bfloat16, device="cuda") + fresh = torch.full( + size=(1, 1, 576), + fill_value=123, + dtype=torch.bfloat16, + device="cuda", + ) topk = torch.tensor([[7, 5, 3, 7, -1], [5, -1, -1, -1, -1]], dtype=torch.int32, device="cuda") compact, mapped = get_prefill_kv_cache_and_remap_indices_triton( packed_kv=state.get_layer_view(0).kv, @@ -96,7 +130,17 @@ def test_packed_fp8_prefill_combines_history_and_fresh_kv_with_sparse_remapping( assert mapped[0, 0] == mapped[0, 3] assert torch.equal(compact[mapped[0, 1]], fresh[0]) assert torch.equal(compact[mapped[1, 0]], fresh[0]) - torch.testing.assert_close(compact[mapped[0, 0], :, :512], history[1, :, :512], atol=0.15, rtol=0.08) - torch.testing.assert_close(compact[mapped[0, 2], :, :512], history[0, :, :512], atol=0.15, rtol=0.08) + torch.testing.assert_close( + actual=compact[mapped[0, 0], :, :512], + expected=history[1, :, :512], + atol=0.15, + rtol=0.08, + ) + torch.testing.assert_close( + actual=compact[mapped[0, 2], :, :512], + expected=history[0, :, :512], + atol=0.15, + rtol=0.08, + ) assert torch.equal(compact[mapped[0, 0], :, 512:], history[1, :, 512:]) assert torch.equal(compact[mapped[0, 2], :, 512:], history[0, :, 512:]) diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 176bb7592e..9f943a64df 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -45,8 +45,12 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( scales[1] += 10 q_scales = None if kind == "static_tensor" else template.q_scales + value states.append(rebuild_full_state( - template, config=replace(template.config, logical_layer_groups=((3,), (8,))), token_capacity=8, device="cuda", - scales=scales, q_scales=q_scales, + template=template, + config=replace(template.config, logical_layer_groups=((3,), (8,))), + token_capacity=8, + device="cuda", + scales=scales, + q_scales=q_scales, )) manager = AttResourceManager(request_capacity=2) manager.set_att_state("main", states[0]) @@ -54,7 +58,12 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) # Another encoding in the same role must not supply the FP8 calibration. plain = definition("gqa") - plain = rebuild_full_state(plain, config=plain.config, token_capacity=8, device="cuda") + plain = rebuild_full_state( + template=plain, + config=plain.config, + token_capacity=8, + device="cuda", + ) manager.set_att_state(role, plain) manager.finish_state_creation() selected = getattr(manager, f"{role}_{template.get_state_type()}_state") @@ -98,9 +107,19 @@ def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.devic return torch.empty(shape, dtype=dtype, device=device) if prefill: - att.prefill_att(q, k, v, alloc_func=allocate) + att.prefill_att( + q=q, + k=k, + v=v, + alloc_func=allocate, + ) else: - att.decode_att(q, k, v, alloc_func=allocate) + att.decode_att( + q=q, + k=k, + v=v, + alloc_func=allocate, + ) if kind == "static_head": heads = selected.config.head_num torch.testing.assert_close(observed["q_scales"], selected.q_scales[1]) @@ -119,10 +138,18 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 scale_count = 2 if scale_mode == "per_tensor" else 2 * head_num - scales = torch.arange(1, scale_count + 1, dtype=torch.float32, device="cuda").unsqueeze(0) / 4 + scales = torch.arange( + start=1, + end=scale_count + 1, + dtype=torch.float32, + device="cuda", + ).unsqueeze(0) / 4 state = StaticFP8FullAttentionState( layout_config=StaticFP8AttentionConfig( - logical_layer_groups=((5,),), global_head_num=head_num, head_dim=128, dtype=torch.bfloat16, + logical_layer_groups=((5,),), + global_head_num=head_num, + head_dim=128, + dtype=torch.bfloat16, scale_mode=scale_mode, ), page_size=4, device="cuda", @@ -149,7 +176,10 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu # 算子的原合并输出入口仍应保持相同的量化结果。 combined = torch.zeros((12, 2 * head_num, 128), dtype=torch.uint8, device="cuda") destindex_copy_kv_fp8( - KV=kv, DestLoc=indices, scales=scales[0], Out=combined.view(torch.float8_e4m3fn), + KV=kv, + DestLoc=indices, + scales=scales[0], + Out=combined.view(torch.float8_e4m3fn), is_per_tensor_quant=scale_mode == "per_tensor", ) assert torch.equal(combined[indices], actual) diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index a991630b1a..e28b973a78 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -33,7 +33,12 @@ def test_direct_cpu_pages_preserve_bytes_without_gather(kind: str, token_major: regions.append(replace(region, axes=axes, shape=tuple(region.shape[region.axes.index(a)] for a in axes))) layout = replace(layout, regions=tuple(regions)) # 页池本身是更大联合页的视图，检验真实 page stride，而不是按字段长度推算下一页。 - pages = torch.full((4, layout.byte_end + 64), 91, dtype=torch.uint8, pin_memory=True) + pages = torch.full( + size=(4, layout.byte_end + 64), + fill_value=91, + dtype=torch.uint8, + pin_memory=True, + ) page_view = pages[:, 32:-32] src_ids = torch.tensor([6, 6, 1, 1, 7, 7, 0, 0, 3, 3], dtype=torch.int32, device="cuda")[::2] dst_ids = torch.tensor([2, 2, 6, 6, 4, 4, 1, 1, 7, 7], dtype=torch.int64, device="cuda")[::2] @@ -67,8 +72,18 @@ def reject_index_contiguous(tensor: torch.Tensor, *args: object, **kwargs: objec patch.setattr(torch.Tensor, "index_select", reject_gather) patch.setattr(torch.Tensor, "index_copy_", reject_gather) with torch.cuda.stream(stream): - source.write_to_cpu(src_ids, page_view, [2, 0], layout=layout) - target.read_from_cpu(page_view, [2, 0], dst_ids, layout=layout) + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=[2, 0], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=[2, 0], + token_indices=dst_ids, + layout=layout, + ) stream.synchronize() for name, tensor in full_state_test_buffers(target).items(): assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) @@ -92,21 +107,41 @@ def test_direct_cpu_pages_only_update_selected_head_slice_and_skip_replica_write regions = [] for region in layout.regions: slices = tuple( - AxisSlice(axis="head", src_start=1, dst_start=0, length=1) if item.axis == "head" else item + AxisSlice( + axis="head", + src_start=1, + dst_start=0, + length=1, + ) if item.axis == "head" else item for item in region.axis_slices ) regions.append(replace(region, axis_slices=slices)) layout = replace(layout, regions=tuple(regions)) pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) ids = torch.tensor([1, 5], dtype=torch.int64, device="cuda") - state.write_to_cpu(ids, pages, [0], layout=layout) + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) torch.cuda.synchronize() snapshot = pages.clone() state.k_storage.fill_(19) state.v_storage.fill_(19) read_only = replace(layout, regions=tuple(replace(region, write_enabled=False) for region in layout.regions)) - state.write_to_cpu(ids, pages, [0], layout=read_only) - state.read_from_cpu(pages, [0], ids, layout=read_only) + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=read_only, + ) + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=read_only, + ) torch.cuda.synchronize() assert torch.equal(pages, snapshot) assert (state.k_storage[:, ids, 1] == 3).all() and (state.v_storage[:, ids, 1] == 7).all() @@ -133,10 +168,20 @@ def test_direct_cpu_pages_accept_registered_external_memory(): try: assert pages.is_pinned() ids = torch.tensor([1, 5], dtype=torch.int32, device="cuda") - state.write_to_cpu(ids, pages, [0], layout=layout) + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) state.k_storage.zero_() state.v_storage.zero_() - state.read_from_cpu(pages, [0], ids, layout=layout) + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=layout, + ) torch.cuda.synchronize() assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() finally: @@ -155,9 +200,19 @@ def test_cpu_cache_transfer_rejects_cpu_state(kind: str, write: bool): ids = torch.tensor([0], dtype=torch.int64) with pytest.raises(AssertionError, match="CUDA State storage"): if write: - state.write_to_cpu(ids, pages, [0], layout=layout) + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) else: - state.read_from_cpu(pages, [0], ids, layout=layout) + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=layout, + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -170,13 +225,28 @@ def test_gqa_single_cpu_page_and_empty_transfer(token_count: int) -> None: target.k_storage.fill_(19) target.v_storage.fill_(19) layout = transfer_layout(source) - pages = torch.full(size=(3, layout.byte_end), fill_value=91, dtype=torch.uint8, pin_memory=True) + pages = torch.full( + size=(3, layout.byte_end), + fill_value=91, + dtype=torch.uint8, + pin_memory=True, + ) src_ids = torch.tensor(data=[5], dtype=torch.int32, device="cuda")[:token_count] dst_ids = torch.tensor(data=[2], dtype=torch.int64, device="cuda")[:token_count] page_indices = [1] if token_count else [] - source.write_to_cpu(src_ids, pages, page_indices, layout=layout) - target.read_from_cpu(pages, page_indices, dst_ids, layout=layout) + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + token_indices=dst_ids, + layout=layout, + ) torch.cuda.synchronize() if token_count: assert (target.k_storage[:, 2] == 3).all() and (target.v_storage[:, 2] == 7).all() @@ -207,10 +277,30 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_tokens_from_other_dp(src, dst, state) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.write_to_cpu(cpu_ids, cpu_pages, [0], layout=layout) + state.write_to_cpu( + token_indices=cpu_ids, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layout=layout, + ) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.read_from_cpu(cpu_pages, [0], cpu_ids, layout=layout) + state.read_from_cpu( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + token_indices=cpu_ids, + layout=layout, + ) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.write_to_pd_buffer(cpu_ids, pd_pages, pd_page_index=0, layout=layout) + state.write_to_pd_buffer( + token_indices=cpu_ids, + pd_page_buffers=pd_pages, + pd_page_index=0, + layout=layout, + ) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.read_from_pd_buffer(pd_pages, cpu_ids, pd_page_index=0, layout=layout) + state.read_from_pd_buffer( + pd_page_buffers=pd_pages, + token_indices=cpu_ids, + pd_page_index=0, + layout=layout, + ) diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py index 68610c1825..7cfeee833c 100644 --- a/unit_tests/common/att_resources/test_full_att_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -114,11 +114,31 @@ def reject_gather(*args: object, **kwargs: object) -> None: patch.setattr(torch.Tensor, "tolist", reject_gather) with torch.cuda.stream(stream): if transfer == "cpu": - source.write_to_cpu(token_indices=src_ids, cpu_page_buffers=page_view, cpu_page_indices=page_indices, layout=layout) - target.read_from_cpu(cpu_page_buffers=page_view, cpu_page_indices=page_indices, token_indices=dst_ids, layout=layout) + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=page_indices, + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=page_indices, + token_indices=dst_ids, + layout=layout, + ) else: - source.write_to_pd_buffer(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout) - target.read_from_pd_buffer(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, layout=layout) + source.write_to_pd_buffer( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=3, + layout=layout, + ) + target.read_from_pd_buffer( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=3, + layout=layout, + ) stream.synchronize() # 比较整行原始字节：packed padding 沿用目标填充值，字段数据来自源存储。 @@ -239,7 +259,11 @@ def _peer_transfer_worker( @pytest.mark.parametrize("readies", [[False, True, False], [True, False, True], [True, True, True], [False, False, False]]) @pytest.mark.parametrize("grid_num", [1, 36]) def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( - kind: str, readies: list[bool], grid_num: int, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + kind: str, + readies: list[bool], + grid_num: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ) -> None: pd_args.cpu_cache_token_page_size = 5 state = rebuild_full_state(definition(kind), token_capacity=32, device="cuda") @@ -252,15 +276,28 @@ def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( buffer.view(torch.uint8).random_(0, 256) regions = state.get_cpu_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, tokens_per_page=5, byte_start=0, - byte_end=max(region.byte_end for region in regions), regions=regions, + resource_id=state.resource_id, + tokens_per_page=5, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, ) # 页编号不连续，token 索引带 stride；最后一页只写三个 token。 page_indices = [3, 0, 4] token_indices = torch.arange(26, dtype=torch.int32, device="cuda")[::2] - baseline = torch.full((6, layout.byte_end + 32), 173, dtype=torch.uint8, pin_memory=True) + baseline = torch.full( + size=(6, layout.byte_end + 32), + fill_value=173, + dtype=torch.uint8, + pin_memory=True, + ) pages = torch.full_like(baseline, 173, pin_memory=True) - state.write_to_cpu(token_indices, baseline[:, 16:-16], page_indices, layout=layout) + state.write_to_cpu( + token_indices=token_indices, + cpu_page_buffers=baseline[:, 16:-16], + cpu_page_indices=page_indices, + layout=layout, + ) torch.cuda.synchronize() expected = pages.clone() for page, ready in zip(page_indices, readies): @@ -274,7 +311,11 @@ def reject_allocation(*args: object, **kwargs: object) -> None: if all(readies): patch.setattr(torch, "tensor", reject_allocation) state.write_to_cpu( - token_indices, pages[:, 16:-16], page_indices, layout=layout, cpu_page_readies=readies, + token_indices=token_indices, + cpu_page_buffers=pages[:, 16:-16], + cpu_page_indices=page_indices, + layout=layout, + cpu_page_readies=readies, ) torch.cuda.synchronize() # 整页比较同时覆盖 KV、量化 scale、packed indexer 字段及前后 padding。 diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index f5a03d59ef..8b6e1ed742 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -24,7 +24,12 @@ def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): - prepared_full_state = cls(FullAttentionConfig(((2,), (5,)), 1, 4, torch.float32), page_size=4, device=device) + prepared_full_state = cls(FullAttentionConfig( + logical_layer_groups=((2,), (5,)), + global_head_num=1, + head_dim=4, + dtype=torch.float32, + ), page_size=4, device=device) prepared_full_state.initialize_storage(12) return prepared_full_state @@ -42,8 +47,18 @@ def layout_for(state: FullAttentionState): axes=("token", "layer", "head", "dim"), shape=(4, 2, 1, 4), axis_slices=( - AxisSlice(axis="head", src_start=0, dst_start=0, length=1), - AxisSlice(axis="dim", src_start=0, dst_start=0, length=4), + AxisSlice( + axis="head", + src_start=0, + dst_start=0, + length=1, + ), + AxisSlice( + axis="dim", + src_start=0, + dst_start=0, + length=4, + ), ), ) ) @@ -96,18 +111,43 @@ def reject_host_list(tensor: torch.Tensor): with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) layout = layout_for(state) - pages = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, device=device) + pages = torch.full( + size=(3, layout.byte_end), + fill_value=255, + dtype=torch.uint8, + device=device, + ) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - state.write_to_pd_buffer(ids, pages, pd_page_index=1, layout=layout) + state.write_to_pd_buffer( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - state.read_from_pd_buffer(pages, ids, pd_page_index=1, layout=layout) + state.read_from_pd_buffer( + pd_page_buffers=pages, + token_indices=ids, + pd_page_index=1, + layout=layout, + ) return - state.write_to_pd_buffer(ids, pages, pd_page_index=1, layout=layout) + state.write_to_pd_buffer( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) target = full_state(cls, device) for storage_buffer in full_state_test_buffers(target).values(): storage_buffer.zero_() - target.read_from_pd_buffer(pages, torch.tensor([5, 6], dtype=torch.long, device=device), pd_page_index=1, layout=layout) + target.read_from_pd_buffer( + pd_page_buffers=pages, + token_indices=torch.tensor([5, 6], dtype=torch.long, device=device), + pd_page_index=1, + layout=layout, + ) assert torch.equal(read_kv_storage(target)[:, [5, 6]], read_kv_storage(state)[:, [1, 3]]) assert (pages[0] == 255).all() and (pages[2] == 255).all() assert not read_kv_storage(target)[:, 0].any() @@ -130,12 +170,27 @@ def test_cpu_partial_last_page(state_offset: int, extra_page_bytes: int): for r in layout.regions ), ) - buffer = torch.full((4, layout.byte_end + extra_page_bytes), 252, dtype=torch.uint8, pin_memory=True) - state.write_to_cpu(torch.tensor([0, 1, 2, 3, 7], dtype=torch.long, device="cuda"), buffer, [2, 0], layout=layout) + buffer = torch.full( + size=(4, layout.byte_end + extra_page_bytes), + fill_value=252, + dtype=torch.uint8, + pin_memory=True, + ) + state.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 7], dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[2, 0], + layout=layout, + ) target = full_state(device="cuda") for storage_buffer in full_state_test_buffers(target).values(): storage_buffer.zero_() - target.read_from_cpu(buffer, [2, 0], torch.tensor([6, 7, 8, 9, 10], dtype=torch.long, device="cuda"), layout=layout) + target.read_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2, 0], + token_indices=torch.tensor([6, 7, 8, 9, 10], dtype=torch.long, device="cuda"), + layout=layout, + ) assert torch.equal(read_kv_storage(target)[:, [6, 7, 8, 9, 10]], read_kv_storage(state)[:, [0, 1, 2, 3, 7]]) assert (buffer[[1, 3]] == 252).all() assert (buffer[:, :layout.byte_start] == 252).all() and (buffer[:, layout.byte_end:] == 252).all() @@ -148,9 +203,19 @@ def test_invalid_late_region_shape_cannot_partially_write_pd_page(): region = replace(region, byte_end=region.byte_end - region.dtype.itemsize) bad = replace(layout, regions=(*layout.regions[:-1], region)) # 直接构造 Tensor 视图时，字节数与 shape 不匹配仍会在提交复制前失败。 - buffer = torch.full((layout.byte_end + 256,), 53, dtype=torch.uint8, device="cuda") + buffer = torch.full( + size=(layout.byte_end + 256,), + fill_value=53, + dtype=torch.uint8, + device="cuda", + ) with pytest.raises(RuntimeError, match="shape"): - state.write_to_pd_buffer(torch.tensor([0], dtype=torch.long, device="cuda"), buffer.unsqueeze(0), pd_page_index=0, layout=bad) + state.write_to_pd_buffer( + token_indices=torch.tensor([0], dtype=torch.long, device="cuda"), + pd_page_buffers=buffer.unsqueeze(0), + pd_page_index=0, + layout=bad, + ) assert (buffer == 53).all() @@ -158,7 +223,12 @@ def test_invalid_late_region_shape_cannot_partially_write_pd_page(): def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") - config = FullAttentionConfig(((1,), (3,), (5,), (7,), (9,), (11,)), 1, 4, torch.float32) + config = FullAttentionConfig( + logical_layer_groups=((1,), (3,), (5,), (7,), (9,), (11,)), + global_head_num=1, + head_dim=4, + dtype=torch.float32, + ) source = GQAFullAttentionState(config, page_size=4, device=device) source.initialize_storage(12) target = GQAFullAttentionState(config, page_size=4, device=device) diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index b8a985d661..d9f4bdfa94 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -38,14 +38,25 @@ def test_model_storage_write_and_copy_preserve_all_fields( encoding: str, model_cls: type[TpPartBaseModel], ): - _, (resources,) = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, 1) + _, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=1, + ) (state,) = resources.get_full_attention_states("main") config = state.config for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") dim = config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else config.head_dim - kv = torch.randn(2, (1 if "kv" in full_state_test_buffers(state) else 2) * config.head_num, dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn( + size=(2, (1 if "kv" in full_state_test_buffers(state) else 2) * config.head_num, dim), + device="cuda", + dtype=torch.bfloat16, + ) state.write_full_att_kv(config.logical_layer_groups[-1][0], ids, kv) if len(config.logical_layer_groups) > 1: assert not read_kv_storage(state)[0].any() @@ -71,7 +82,14 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( encoding: str, model_cls: type[TpPartBaseModel], ): - _, (reference,) = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, 1) + _, (reference,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=1, + ) (state,) = reference.get_full_attention_states("main") args = get_env_start_args() args.run_mode, args.pd_kv_page_size, args.pd_kv_page_num = mode, 8, 3 diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py index 934803990f..762208c425 100644 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -30,7 +30,12 @@ def test_gqa_pd_direct_copy_between_tp_layouts( pd_args.pd_kv_page_size = 4 # P 保存完整 heads，D 只恢复 rank 1 的一半 heads；两侧外部页格式相同。 source = GQAFullAttentionState( - layout_config=FullAttentionConfig(logical_layer_groups=((2,), (5,)), global_head_num=4, head_dim=128, dtype=dtype), + layout_config=FullAttentionConfig( + logical_layer_groups=((2,), (5,)), + global_head_num=4, + head_dim=128, + dtype=dtype, + ), page_size=4, device="cuda", ) @@ -65,7 +70,12 @@ def test_gqa_pd_direct_copy_between_tp_layouts( regions=tuple(regions), )) layout, target_layout = layouts - pages = torch.full(size=(3, layout.byte_end + 64), fill_value=91, dtype=torch.uint8, device="cuda") + pages = torch.full( + size=(3, layout.byte_end + 64), + fill_value=91, + dtype=torch.uint8, + device="cuda", + ) page_view = pages[:, 32:-32] # 页池视图的真实跨度大于可见宽度。 src_ids = torch.tensor(data=[6, 9, 1, 9, 7, 9], dtype=torch.int32, device="cuda")[::2][:token_count] dst_ids = torch.tensor(data=[2, 6, 4], dtype=torch.int64, device="cuda")[:token_count] @@ -81,8 +91,18 @@ def reject_gather(*args: object, **kwargs: object) -> None: patch.setattr(torch.Tensor, "index_copy_", reject_gather) patch.setattr(torch.Tensor, "tolist", reject_gather) with torch.cuda.stream(stream): - source.write_to_pd_buffer(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=1, layout=layout) - target.read_from_pd_buffer(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=1, layout=target_layout) + source.write_to_pd_buffer( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=1, + layout=layout, + ) + target.read_from_pd_buffer( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=1, + layout=target_layout, + ) stream.synchronize() for tensor, expected_bytes in zip((target.k_storage, target.v_storage), expected): assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected_bytes) @@ -98,7 +118,10 @@ def reject_gather(*args: object, **kwargs: object) -> None: # 反向验证分片写出：rank 1 只修改全局 heads [2, 4)，rank 0 区域保持原始字节。 pages.fill_(91) target.write_to_pd_buffer( - token_indices=dst_ids, pd_page_buffers=page_view, pd_page_index=2, layout=target_layout, + token_indices=dst_ids, + pd_page_buffers=page_view, + pd_page_index=2, + layout=target_layout, ) for region, expected_bytes in zip(target_layout.regions, expected): external = page_view[2, region.byte_start:region.byte_end].view(dtype).view(region.shape) @@ -115,7 +138,12 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: # MQA 的 KV head 在 TP ranks 间复制，非首 rank 禁止写出但仍能读回。 state = GQAFullAttentionState( layout_config=FullAttentionConfig( - logical_layer_groups=((0,), (2,)), global_head_num=1, head_dim=128, dtype=torch.bfloat16, tp_world_size=2, tp_rank=1, + logical_layer_groups=((0,), (2,)), + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, + tp_world_size=2, + tp_rank=1, ), page_size=4, device="cuda", @@ -125,16 +153,29 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: state.v_storage.fill_(19) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, tokens_per_page=4, byte_start=0, - byte_end=max(region.byte_end for region in regions), regions=regions, + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, ) pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda") for region, value in zip(regions, (3, 7)): pages[1, region.byte_start:region.byte_end].view(region.dtype).fill_(value) snapshot = pages.clone() ids = torch.tensor(data=[5, 1], dtype=torch.int32, device="cuda") - state.write_to_pd_buffer(token_indices=ids, pd_page_buffers=pages, pd_page_index=1, layout=layout) - state.read_from_pd_buffer(pd_page_buffers=pages, token_indices=ids, pd_page_index=1, layout=layout) + state.write_to_pd_buffer( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_buffer( + pd_page_buffers=pages, + token_indices=ids, + pd_page_index=1, + layout=layout, + ) assert torch.equal(pages, snapshot) assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() assert (state.k_storage[:, [0, 2, 3, 4, 6, 7]] == 19).all() @@ -153,7 +194,12 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt pytest.skip("CUDA peer access required") pd_args.pd_kv_page_size = 4 state = GQAFullAttentionState( - layout_config=FullAttentionConfig(logical_layer_groups=((0,), (2,)), global_head_num=2, head_dim=128, dtype=torch.bfloat16), + layout_config=FullAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=128, + dtype=torch.bfloat16, + ), page_size=4, device="cuda:1", ) @@ -165,8 +211,11 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt torch.cuda.synchronize(device=1) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, tokens_per_page=4, byte_start=0, - byte_end=max(region.byte_end for region in regions), regions=regions, + resource_id=state.resource_id, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, ) monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) payload = bytes(ForkingPickler.dumps(state)) diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index 6f767e34f0..d86a671d50 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -18,7 +18,12 @@ @pytest.mark.parametrize("page_device", ["cpu", "cuda"]) def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_args: "EasyDict") -> None: pd_args.pd_kv_page_size = 32 - config = FullAttentionConfig(logical_layer_groups=((0,), (2,), (4,)), global_head_num=3, head_dim=128, dtype=torch.bfloat16) + config = FullAttentionConfig( + logical_layer_groups=((0,), (2,), (4,)), + global_head_num=3, + head_dim=128, + dtype=torch.bfloat16, + ) source = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") target = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") source.initialize_storage(token_capacity=128) diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index bffbe941d6..0570effb52 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -28,7 +28,12 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") -> None: template = definition(kind) config = replace(template.config, logical_layer_groups=((0, 4), (8, 11))) - state = rebuild_full_state(template, config=config, token_capacity=8, device="cpu") + state = rebuild_full_state( + template=template, + config=config, + token_capacity=8, + device="cpu", + ) assert all(state.has_logical_layer(logical_layer_id) for logical_layer_id in (0, 4, 8, 11)) assert not state.has_logical_layer(1) assert all(buffer.shape[0] == 2 for buffer in full_state_test_buffers(state).values()) @@ -55,14 +60,24 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - @pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) def test_invalid_layer_groups_are_rejected(groups: tuple) -> None: with pytest.raises(AssertionError): - FullAttentionConfig(logical_layer_groups=groups, global_head_num=1, head_dim=128, dtype=torch.bfloat16) + FullAttentionConfig( + logical_layer_groups=groups, + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, + ) def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "EasyDict") -> None: pd_args.data_type = "bfloat16" config = dict( - num_hidden_layers=6, num_attention_heads=4, num_key_value_heads=2, hidden_size=512, - head_dim=128, global_head_dim=128, num_kv_shared_layers=2, + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=512, + head_dim=128, + global_head_dim=128, + num_kv_shared_layers=2, layer_types=["sliding_attention", "full_attention"] * 3, ) (state,) = Gemma4TpPartModel._create_full_att_states( @@ -83,10 +98,16 @@ def test_shared_layers_transfer_each_physical_layer_once( pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 template = definition(kind) source = rebuild_full_state( - template, config=replace(template.config, logical_layer_groups=((0, 4), (8,))), token_capacity=8, device="cuda", + template=template, + config=replace(template.config, logical_layer_groups=((0, 4), (8,))), + token_capacity=8, + device="cuda", ) target = rebuild_full_state( - template, config=replace(template.config, logical_layer_groups=((20, 24), (28,))), token_capacity=8, device="cuda", + template=template, + config=replace(template.config, logical_layer_groups=((20, 24), (28,))), + token_capacity=8, + device="cuda", ) source_buffers = full_state_test_buffers(source) target_buffers = full_state_test_buffers(target) @@ -105,21 +126,31 @@ def test_shared_layers_transfer_each_physical_layer_once( if transfer == "cpu": pages = torch.zeros((2, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) managers[0].write_req_to_cpu( - token_indices=src_ids, cpu_page_buffers=pages, cpu_page_indices=[1], + token_indices=src_ids, + cpu_page_buffers=pages, + cpu_page_indices=[1], layouts=managers[0].get_cpu_transfer_layouts(), ) managers[1].read_req_from_cpu( - cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=dst_ids, + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=dst_ids, layouts=managers[1].get_cpu_transfer_layouts(), ) else: pages = torch.zeros((2, managers[0].get_pd_page_bytes()), dtype=torch.uint8, device="cuda") managers[0].write_to_pd_buffer( - pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=src_ids, + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + token_indices=src_ids, layouts=managers[0].get_pd_full_att_transfer_layouts(), ) managers[1].read_from_pd_buffer( - pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=dst_ids, + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + token_indices=dst_ids, layouts=managers[1].get_pd_full_att_transfer_layouts(), ) torch.cuda.synchronize() diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 87b6dd68c3..5f5522e328 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -26,7 +26,10 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) state = GQAFullAttentionState( layout_config=FullAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=128, dtype=torch.bfloat16, + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, ), device="meta", ) monkeypatch.setattr( @@ -47,7 +50,10 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): manager.set_att_state(model_role, GQAFullAttentionState( layout_config=FullAttentionConfig( - logical_layer_groups=groups, global_head_num=1, head_dim=128, dtype=dtype, + logical_layer_groups=groups, + global_head_num=1, + head_dim=128, + dtype=dtype, ), device="meta", )) model = TpPartBaseModel.__new__(TpPartBaseModel) @@ -154,7 +160,12 @@ def test_mla_selection_excludes_same_layer_tail() -> None: assert not tail.has_logical_layer(0) assert manager.get_att_state("main", "mla", 7) is state assert manager.get_att_state("main", "kpool", 7) is tail - assert manager.get_att_state("main", "", 7, state_class=FullAttentionState) is state + assert manager.get_att_state( + model_name="main", + att_type="", + logical_layer_id=7, + state_class=FullAttentionState, + ) is state with pytest.raises(AssertionError, match="expected exactly one"): manager.get_att_state("main", "", 7) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index e43f32b5e1..0c8d653627 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -22,11 +22,11 @@ from easydict import EasyDict -def composition(device: str | torch.device | int = "cpu", *, with_tail: bool = True, tail_pool_size: int = 4): +def composition(device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): return MainDraftTestResourceManager(device, with_tail=with_tail, tail_pool_size=tail_pool_size) -def layouts_for(manager: AttResourceManager, *, kind: str = "cpu"): +def layouts_for(manager: AttResourceManager, kind: str = "cpu"): layouts = [] if kind != "att_state": layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] @@ -109,8 +109,10 @@ def test_disabled_domains(): def test_reset_tail_clears_all_tail_states_and_preserves_linear(): first = make_tail() second = KPoolAttentionTailState( - replace(first.config, pool_size=3), - request_capacity=first.request_capacity, mtp_step=first.mtp_step, device=first.tail_buffer.device, + config=replace(first.config, pool_size=3), + request_capacity=first.request_capacity, + mtp_step=first.mtp_step, + device=first.tail_buffer.device, ) first.tail_buffer.fill_(5) second.tail_buffer.fill_(7) @@ -155,9 +157,9 @@ def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | tor (3, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=device == "cuda" ) source.write_req_to_cpu( - torch.arange(8, dtype=torch.long, device=device), - pages, - [2, 0], + token_indices=torch.arange(8, dtype=torch.long, device=device), + cpu_page_buffers=pages, + cpu_page_indices=[2, 0], layouts=layouts, big_page_slot_ids=slots, ) @@ -168,9 +170,14 @@ def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | tor assert tail.resource_id not in layouts target_slots = target.bounded_att_big_page_allocator.alloc(2) target.read_req_from_cpu( - pages, - [2, 0], - torch.arange(4, 12, dtype=torch.long, device=device), + cpu_page_buffers=pages, + cpu_page_indices=[2, 0], + token_indices=torch.arange( + start=4, + end=12, + dtype=torch.long, + device=device, + ), layouts=layouts, big_page_slot_ids=target_slots, req_idx=2, @@ -208,8 +215,12 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i source.save_bounded_att_pages([page_count - 1], page_kind="small", slot_ids=small_slots) torch.cuda.synchronize() source.write_req_to_cpu( - torch.arange(token_count, device="cuda"), pages, page_ids, - layouts=layouts, big_page_slot_ids=source_slots[:token_count // 4], small_page_slot_ids=small_slots, + token_indices=torch.arange(token_count, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=layouts, + big_page_slot_ids=source_slots[:token_count // 4], + small_page_slot_ids=small_slots, ) torch.cuda.synchronize() @@ -221,8 +232,12 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i original_slots = slots.copy() free_count = allocator.free_count target.read_req_from_cpu( - pages, page_ids, torch.arange(4, 4 + token_count, device="cuda"), - layouts=layouts, big_page_slot_ids=slots, req_idx=2, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + token_indices=torch.arange(4, 4 + token_count, device="cuda"), + layouts=layouts, + big_page_slot_ids=slots, + req_idx=2, ) torch.cuda.synchronize() assert slots == original_slots @@ -252,9 +267,9 @@ def test_cpu_restore_resets_selected_tail_after_all_pages(tail_pool_size: int): layouts = layouts_for(source) pages = torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) source.write_req_to_cpu( - torch.arange(8, dtype=torch.long, device="cuda"), - pages, - [0, 1], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[0, 1], layouts=layouts, big_page_slot_ids=slots, ) @@ -264,9 +279,9 @@ def test_cpu_restore_resets_selected_tail_after_all_pages(tail_pool_size: int): tail.tail_buffer.fill_(37) # 缓存流程保证最终前缀 L=8 是可恢复边界；中间页 L=4 不单独重置 tail。 target.read_req_from_cpu( - pages, - [0, 1], - torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[0, 1], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), layouts=layouts, big_page_slot_ids=target_slots, req_idx=0, @@ -286,16 +301,30 @@ def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.d page = torch.empty(max(layout.byte_end for layout in layouts.values()), dtype=torch.uint8, device=device) src_args = {"token_indices": torch.tensor([0, 2], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 1, "state_token_len": 8} dst_args = {"token_indices": torch.tensor([5, 7], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 2, "state_token_len": 8} - source.write_to_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=layouts, **src_args) - target.read_from_pd_buffer(page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=layouts, **dst_args) + source.write_to_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=layouts, + **src_args, + ) + target.read_from_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=layouts, + **dst_args, + ) dp_tail = dp_target.get_attention_tail_states()[0] dp_tail.tail_buffer.fill_(19) dp_target.copy_from_other_dp( remote_att_manager=source, src_token_indexes=torch.tensor([0, 2], dtype=torch.int32, device=device), dst_token_indexes=torch.tensor([5, 7], dtype=torch.long, device=device), - src_big_page_ids=[], dst_big_page_ids=[], - src_req_idx=1, dst_req_idx=2, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=1, + dst_req_idx=2, ) for a, b in zip( (target.main_gqa_fp32_state, target.draft0_mla_fp32_state), @@ -327,8 +356,10 @@ def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): remote_att_manager=source, src_token_indexes=torch.empty(0, dtype=torch.long), dst_token_indexes=torch.empty(0, dtype=torch.long), - src_big_page_ids=[], dst_big_page_ids=[], - src_req_idx=1, dst_req_idx=2, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=1, + dst_req_idx=2, ) assert all(not state.tail_buffer[:, 2].any() for state in tails) assert all((state.tail_buffer[:, :2] == 19).all() for state in tails) @@ -365,7 +396,8 @@ def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: boo dst_token_indexes=torch.tensor([5, 7] if copy_tokens else [], dtype=torch.long, device=device), src_big_page_ids=[2, 0] if copy_pages else [], dst_big_page_ids=[0, 1] if copy_pages else [], - src_req_idx=1, dst_req_idx=2, + src_req_idx=1, + dst_req_idx=2, ) for state in target.get_full_attention_states(): @@ -418,10 +450,20 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo layouts = layouts_for(source) page_ids = [2, 0, 3] token_ids = torch.arange(10, device="cuda") - baseline = torch.full((4, max(layout.byte_end for layout in layouts.values())), 139, dtype=torch.uint8, pin_memory=True) + baseline = torch.full( + size=(4, max(layout.byte_end for layout in layouts.values())), + fill_value=139, + dtype=torch.uint8, + pin_memory=True, + ) pages = torch.full_like(baseline, 139, pin_memory=True) source.write_req_to_cpu( - token_ids, baseline, page_ids, layouts=layouts, big_page_slot_ids=big_slots, small_page_slot_ids=small_slots, + token_indices=token_ids, + cpu_page_buffers=baseline, + cpu_page_indices=page_ids, + layouts=layouts, + big_page_slot_ids=big_slots, + small_page_slot_ids=small_slots, ) torch.cuda.synchronize() expected = pages.clone() @@ -432,8 +474,13 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo masked_big_slots = [-999 if readies[i] else slot for i, slot in enumerate(big_slots)] masked_small_slots = [-999] if readies[2] else small_slots source.write_req_to_cpu( - token_ids, pages, page_ids, layouts=layouts, cpu_page_readies=readies, - big_page_slot_ids=masked_big_slots, small_page_slot_ids=masked_small_slots, + token_indices=token_ids, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=layouts, + cpu_page_readies=readies, + big_page_slot_ids=masked_big_slots, + small_page_slot_ids=masked_small_slots, ) torch.cuda.synchronize() assert torch.equal(pages, expected) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index 1472454222..c91fe76fd0 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -103,9 +103,11 @@ def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encod states = [] if not (model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att")): states = model._create_full_att_states( - model_config=model.config, args=model.args, + model_config=model.config, + args=model.args, 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_, + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, ) for state in states: assert state.resource_id is None @@ -139,7 +141,11 @@ def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( q_scales = scales[:, :2] + 20 calibration_path = tmp_path / "calibration.json" calibration_path.write_text(json.dumps(dict( - num_layers=4, num_head=2, quant_type="per_head", scales_shape=[4, 4], scales=scales.tolist(), + num_layers=4, + num_head=2, + quant_type="per_head", + scales_shape=[4, 4], + scales=scales.tolist(), q_calibration=dict(num_head=2, scales_shape=[4, 2], scales=q_scales.tolist()), ))) args = StartArgs(kv_quant_calibration_config_path=str(calibration_path)) @@ -163,8 +169,12 @@ def count_open(file: str, *args: object, **kwargs: object): ): state = create_gqa_state( config=FullAttentionConfig( - logical_layer_groups=groups, global_head_num=4, head_dim=16, dtype=torch.bfloat16, - tp_rank=tp_rank, tp_world_size=tp_world_size, + logical_layer_groups=groups, + global_head_num=4, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), kv_type=encoding, device="meta", ) manager.set_att_state(role, state) @@ -188,11 +198,19 @@ def count_open(file: str, *args: object, **kwargs: object): @pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) def test_cpu_layout_does_not_require_calibration_file( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, cpu_args: EasyDict, encoding: str, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + encoding: str, ): from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts - config = dict(model_type="llama", num_attention_heads=2, hidden_size=32, num_hidden_layers=2) + config = dict( + model_type="llama", + num_attention_heads=2, + hidden_size=32, + num_hidden_layers=2, + ) (tmp_path / "config.json").write_text(json.dumps(config)) cpu_args.model_dir = str(tmp_path) cpu_args.llm_kv_type = encoding @@ -212,7 +230,10 @@ def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): state = create_gqa_state( config=FullAttentionConfig( - logical_layer_groups=((0, 1), (2,)), global_head_num=2, head_dim=16, dtype=torch.bfloat16, + logical_layer_groups=((0, 1), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, ), kv_type=encoding, device="cpu", ) manager = AttResourceManager(request_capacity=1) diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 17783eb6b6..8296c60bf9 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -38,7 +38,12 @@ def clear_mtp_capacity_cache(): @pytest.mark.parametrize("glm", [False, True]) @pytest.mark.parametrize("mtp", [0, 2]) @pytest.mark.parametrize("mode", ["normal", "prefill"]) -def test_real_model_resource_hooks_and_shared_draft(monkeypatch: pytest.MonkeyPatch, glm: bool, mtp: int, mode: str): +def test_real_model_resource_hooks_and_shared_draft( + monkeypatch: pytest.MonkeyPatch, + glm: bool, + mtp: int, + mode: str, +): args = StartArgs( run_mode=mode, data_type="bfloat16", @@ -117,7 +122,12 @@ def check_runtime_before_profile(device: str | torch.device | int): draft.args, draft.data_type = args, torch.bfloat16 draft.tp_rank_, draft.tp_world_size_ = 0, 1 draft.max_req_num, draft.max_total_token_num = model.max_req_num, None - draft.config = dict(model.config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft.config = dict( + model.config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + ) draft.max_seq_length, draft.batch_max_tokens = model.max_seq_length, model.batch_max_tokens prepare_model_states(draft, logical_layer_ids=(4,)) assert draft.att_resource_manager is model.att_resource_manager diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index af505e8d34..57a7127ed7 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -32,7 +32,13 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int = 0): +def make_resources( + glm: bool, + tp: int, + rank: int, + mtp: int, + draft_layers: int = 0, +): gk, gv = (4, 4) if glm else (2, 4) local_k, local_v = max(1, gk // tp), max(1, gv // tp) heads = 1 if glm else max(1, 2 // tp) @@ -60,7 +66,13 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int if glm: full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), 1, 584, torch.bfloat16, rank, tp, kv_dim=512 + logical_layer_groups=tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), + global_head_num=1, + head_dim=584, + dtype=torch.bfloat16, + tp_rank=rank, + tp_world_size=tp, + kv_dim=512, ), page_size=4, ) @@ -68,14 +80,22 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int else: full = GQAFullAttentionState( FullAttentionConfig( - tuple((layer,) for layer in (1, 3, *range(4, 4 + draft_layers))), 2, 16, torch.bfloat16, rank, tp + logical_layer_groups=tuple((layer,) for layer in (1, 3, *range(4, 4 + draft_layers))), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=rank, + tp_world_size=tp, ), page_size=4, ) full.initialize_storage(16) main, draft = create_full_attention_roles(full, 4) resources = build_hybrid_resources( - config=config, request_capacity=3, mtp_step=mtp, tp_rank=rank, + config=config, + request_capacity=3, + mtp_step=mtp, + tp_rank=rank, full_state=main, full_allocator=FullAttKVAllocator(16, 4), draft_full_state=draft, @@ -88,9 +108,24 @@ def make_resources(glm: bool, tp: int, rank: int, mtp: int, *, draft_layers: int @pytest.mark.parametrize("glm", [False, True]) @pytest.mark.parametrize("tp", [1, 2, 4]) @pytest.mark.parametrize("mtp", [0, 2]) -def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: int, mtp: int, cpu_args: EasyDict): - sources = [make_resources(glm, tp, rank, mtp) for rank in range(tp)] - targets = [make_resources(glm, tp, rank, mtp) for rank in range(tp)] +def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( + glm: bool, + tp: int, + mtp: int, + cpu_args: EasyDict, +): + sources = [make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) for rank in range(tp)] + targets = [make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) for rank in range(tp)] expected = [] for rank, (resources, full) in enumerate(sources): cfg = main_linear_state(resources).config @@ -117,9 +152,9 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: i buffer = torch.zeros((2, size), dtype=torch.uint8, pin_memory=True) for resources, full in sources: resources.write_req_to_cpu( - torch.arange(8, dtype=torch.long, device="cuda"), - buffer, - [1], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[1], layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=[0], ) @@ -129,9 +164,14 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail(glm: bool, tp: i resources.main_kpool_c4_bf16_state.tail_buffer.fill_(13) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] resources.read_req_from_cpu( - buffer, - [1], - torch.arange(8, 16, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[1], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=[slot], req_idx=1, @@ -155,6 +195,11 @@ def test_cpu_layout_query_does_not_require_cpu_cache_enabled(monkeypatch: pytest args = StartArgs(enable_cpu_cache=False, cpu_cache_token_page_size=8) # 未启用缓存时也可查询全局页配置。 monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) get_env_start_args.cache_clear() - resources, _ = make_resources(True, 1, 0, 0) + resources, _ = make_resources( + glm=True, + tp=1, + rank=0, + mtp=0, + ) assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_transfer_layouts().values()) get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index aad1f5d245..241b6c3291 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -106,7 +106,13 @@ def test_model_and_cpu_service_use_identical_storage_geometry( local_heads: int, ): # Exercise raw HF aliases as well as a wrapped multimodal config. - config = dict(model_type="llama", n_head=8, n_embd=128, n_layer=2, **overrides) + config = dict( + model_type="llama", + n_head=8, + n_embd=128, + n_layer=2, + **overrides, + ) wrapped = {"model_type": "gemma4", "text_config": config} if model_cls is Gemma4TpPartModel else config (tmp_path / "config.json").write_text(json.dumps(wrapped)) monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: model_cls) @@ -151,7 +157,12 @@ def test_gpu_profile_counts_token_fields_hold_and_transfer_reserve( reference.set_att_state("main", state) reference.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) reference.finish_state_creation() - args = StartArgs(page_size=4, run_mode="prefill", pd_kv_page_size=8, pd_kv_page_num=3) + args = StartArgs( + page_size=4, + run_mode="prefill", + pd_kv_page_size=8, + pd_kv_page_num=3, + ) reserve = reference.get_pd_page_bytes() * args.pd_kv_page_num required = state.get_token_bytes() * (32 + args.page_size) # One byte below the full 32-token budget must drop a complete page. @@ -170,11 +181,21 @@ def memory_info(device: str | torch.device | int): monkeypatch.setattr( "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - template = rebuild_full_state(state, config=state.config, token_capacity=0, device="meta") + template = rebuild_full_state( + template=state, + config=state.config, + token_capacity=0, + device="meta", + ) resources = AttResourceManager(request_capacity=1) resources.set_att_state("main", rebuild_full_state(template, config=template.config)) resources.finish_state_creation() - resources.initialize_resources(args, token_capacity=None, mem_fraction=0.75, device="cpu") + resources.initialize_resources( + args=args, + token_capacity=None, + mem_fraction=0.75, + device="cpu", + ) assert resources.full_att_allocator.size == 28 assert resources.get_pd_page_bytes() == reference.get_pd_page_bytes() assert observed_devices == [torch.device("cuda", 3)] @@ -189,12 +210,19 @@ def test_cpu_layout_uses_all_model_factory_states( class MixedStorageModel(LlamaTpPartModel): @classmethod def _create_full_att_states( - cls, model_config: dict[str, object], args: StartArgs | EasyDict, *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list[FullAttentionState]: (template,) = super()._create_full_att_states( - model_config=model_config, args=args, model_logical_layer_ids=model_logical_layer_ids, - tp_rank=tp_rank, tp_world_size=tp_world_size, + model_config=model_config, + args=args, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ) return [ GQAFullAttentionState( @@ -209,7 +237,12 @@ def _create_full_att_states( ), ] - config = dict(model_type="llama", num_attention_heads=8, hidden_size=128, num_hidden_layers=3) + config = dict( + model_type="llama", + num_attention_heads=8, + hidden_size=128, + num_hidden_layers=3, + ) (tmp_path / "config.json").write_text(json.dumps(config)) monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: MixedStorageModel) cpu_layouts = build_cpu_cache_layouts() diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index dce30194c1..edc514cb6d 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -96,10 +96,10 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i tokens.extend(99 if offset == accepted - 1 else 10 + offset for offset in range(width)) expected_mask.extend(int(offset < accepted) for offset in range(width)) accept_lengths, accepted_index = verify_mtp_tokens( - backend, - torch.tensor(tokens, device="cuda"), - torch.tensor(reqs, dtype=torch.int32, device="cuda"), - torch.tensor(starts, dtype=torch.int32, device="cuda"), + backend=backend, + next_token_ids=torch.tensor(tokens, device="cuda"), + b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), + b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), ) assert accept_lengths.tolist() == [3, 1, 2] diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 3c2ad55397..f450ed9ce9 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -157,14 +157,24 @@ def _worker(connection: Connection, backend: str, sender: bool): torch.set_num_threads(2) set_env_start_args(StartArgs(run_mode="prefill" if sender else "decode", pd_kv_page_size=32)) get_env_start_args.cache_clear() - resources, full = make_resources(True, 1, 0, 2) + resources, full = make_resources( + glm=True, + tp=1, + rank=0, + mtp=2, + ) buffer = allocate_transfer_buffer(resources, 3, device=f"cuda:{device}") buffer.fill_(93 if sender else 11) copy_stream = torch.cuda.Stream() if backend == "nccl": from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter - transporter = NcclKVTransporter(os.getpid(), device, buffer, host_ip="127.0.0.1") + transporter = NcclKVTransporter( + node_id=os.getpid(), + tp_idx=device, + kv_move_buffer=buffer, + host_ip="127.0.0.1", + ) else: from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter @@ -173,10 +183,10 @@ def _worker(connection: Connection, backend: str, sender: bool): ( "metadata", PDAgentMetadata( - transporter.agent_name, - transporter.agent_metadata, - transporter.num_pages, - transporter.local_page_mem_desc, + agent_name=transporter.agent_name, + agent_metadata=transporter.agent_metadata, + num_pages=transporter.num_pages, + page_reg_desc=transporter.local_page_mem_desc, ), ) ) @@ -192,7 +202,13 @@ def _worker(connection: Connection, backend: str, sender: bool): (state.conv_state, state.ssm_state, resources.main_kpool_c4_bf16_state.tail_buffer) ): tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) - expected = _snapshot(resources, full, [3, 0, 7, 1, 9], 1, 0) + expected = _snapshot( + resources=resources, + full=full, + indices=[3, 0, 7, 1, 9], + req_idx=1, + cursor=0, + ) for kind in ("kv", "att_state"): task = _task(peer, kind) with torch.cuda.stream(copy_stream): @@ -211,7 +227,12 @@ def _worker(connection: Connection, backend: str, sender: bool): # The production sender publishes a page only after its copy event completes. copy_stream.synchronize() transporter.send_write_request_task_to_decode_node(task) - ready = _notification(transporter, "ready", pending, task.get_key()) + ready = _notification( + transporter=transporter, + stage="ready", + pending=pending, + key=task.get_key(), + ) assert ready.get_key() == task.get_key() task.dst_page_index = ready.dst_page_index task.xfer_handle = transporter.write_blocks_paged(task) @@ -233,7 +254,12 @@ def _worker(connection: Connection, backend: str, sender: bool): assert request.page_kind == kind request.dst_page_index = 2 if kind == "kv" else 1 transporter.send_write_ready_task_to_prefill_node(request) - done = _notification(transporter, "done", pending, request.get_key()) + done = _notification( + transporter=transporter, + stage="done", + pending=pending, + key=request.get_key(), + ) assert done.get_key() == request.get_key() # Match production: unpack on a separate stream after the done notification. with torch.cuda.stream(copy_stream): @@ -250,7 +276,13 @@ def _worker(connection: Connection, backend: str, sender: bool): write_to_pd_buffer=False, ) copy_stream.synchronize() - expected = _snapshot(resources, full, [8, 12, 10, 15, 11], 2, 0) + expected = _snapshot( + resources=resources, + full=full, + indices=[8, 12, 10, 15, 11], + req_idx=2, + cursor=0, + ) assert main_linear_state(resources).mtp_state_index[2] == 0 assert not main_linear_state(resources).conv_state[:, :2].any() assert not main_linear_state(resources).ssm_state[:, :6].any() diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index e47073939c..0887dcaeff 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -24,7 +24,12 @@ def draft_model(mode: str): draft.args = StartArgs(mtp_mode=mode, data_type="bfloat16") draft.max_req_num = 2 draft.layers_infer = [SimpleNamespace(layer_num_=2)] - draft.config = dict(num_attention_heads=2, num_key_value_heads=1, hidden_size=8, num_hidden_layers=1) + draft.config = dict( + num_attention_heads=2, + num_key_value_heads=1, + hidden_size=8, + num_hidden_layers=1, + ) draft.tp_rank_, draft.tp_world_size_ = 0, 1 main = LlamaTpPartModel.__new__(LlamaTpPartModel) main.args, main.config = draft.args, dict(draft.config, num_hidden_layers=2) diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index f813bb5802..17c3b44141 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -28,9 +28,14 @@ def test_page_pool_allocation_strategy( pytest.skip("CUDA required for pinned memory") cpu_args.enable_dp_prompt_cache_fetch = dp_fetch config = ConvSSMStateConfig( - logical_layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, - key_dim=4, value_dim=4, conv_width=3, - conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + logical_layer_ids=(0, 2), + global_k_heads=1, + global_v_heads=1, + key_dim=4, + value_dim=4, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, ) pool = pool_cls(size=2, config=config, pin_memory=pin_memory) assert pool.conv_state.shape == (2, 2, 12, 3) @@ -99,9 +104,14 @@ def test_pinned_page_pool_survives_queue_on_another_gpu( cpu_args.enable_dp_prompt_cache_fetch = True torch.cuda.set_device(0) config = ConvSSMStateConfig( - logical_layer_ids=(0, 2), global_k_heads=1, global_v_heads=1, - key_dim=4, value_dim=4, conv_width=3, - conv_dtype=torch.float32, ssm_dtype=torch.float32, + logical_layer_ids=(0, 2), + global_k_heads=1, + global_v_heads=1, + key_dim=4, + value_dim=4, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, ) pool = pool_cls(size=2, config=config) tensors = (pool.conv_state, pool.ssm_state) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 5be91d3172..2560d0eefa 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -37,16 +37,31 @@ def make_quant_state(kind: str, device: str | torch.device | int): if kind in ("int4", "int8"): - config = GroupQuantizedAttentionConfig(((1,), (3,)), 2, 16, torch.float16, bits=int(kind[-1])) + config = GroupQuantizedAttentionConfig( + logical_layer_groups=((1,), (3,)), + global_head_num=2, + head_dim=16, + dtype=torch.float16, + bits=int(kind[-1]), + ) prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) prepared_full_state.initialize_storage(8) return prepared_full_state if kind.startswith("fp8_"): per_head = kind == "fp8_head" config = StaticFP8AttentionConfig( - ((1,), (3,)), 2, 16, torch.uint8, scale_mode="per_head" if per_head else "per_tensor" + logical_layer_groups=((1,), (3,)), + global_head_num=2, + head_dim=16, + dtype=torch.uint8, + scale_mode="per_head" if per_head else "per_tensor", ) - scales = torch.arange(1, 9 if per_head else 5, dtype=torch.float32, device=device).reshape(2, -1) / 16 + scales = torch.arange( + start=1, + end=9 if per_head else 5, + dtype=torch.float32, + device=device, + ).reshape(2, -1) / 16 prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device) bind_test_fp8_calibration( state=prepared_full_state, scales=scales, @@ -56,7 +71,11 @@ def make_quant_state(kind: str, device: str | torch.device | int): return prepared_full_state fp8 = kind == "packed_fp8" config = PackedMLAAttentionConfig( - ((1,), (3,)), 1, 800 if fp8 else 648, torch.uint8 if fp8 else torch.bfloat16, fp8_nope_dim=512 if fp8 else 0 + logical_layer_groups=((1,), (3,)), + global_head_num=1, + head_dim=800 if fp8 else 648, + dtype=torch.uint8 if fp8 else torch.bfloat16, + fp8_nope_dim=512 if fp8 else 0, ) prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) prepared_full_state.initialize_storage(8) @@ -78,7 +97,12 @@ def transfer_layout(state: FullAttentionState): axes=field.axes, shape=shape, axis_slices=tuple( - AxisSlice(axis=axis, src_start=0, dst_start=0, length=dim) + AxisSlice( + axis=axis, + src_start=0, + dst_start=0, + length=dim, + ) for axis, dim in zip(field.axes, shape) if axis not in ("layer", "token") ), ) @@ -121,28 +145,78 @@ def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( field.tensor[:, 0].fill_(i + 11) if device == "cuda": source.copy_full_att_tokens(torch.tensor([0, 2], dtype=torch.long, device=device), torch.tensor([4, 5], dtype=torch.long, device=device)) - assert_token_fields_equal(source, [0, 2], source, [4, 5]) + assert_token_fields_equal( + source=source, + source_ids=[0, 2], + target=source, + target_ids=[4, 5], + ) layout = transfer_layout(source) cpu = torch.zeros((2, layout.byte_end), dtype=torch.uint8, pin_memory=device == "cuda") if device == "cpu": with pytest.raises(AssertionError, match="CUDA State storage"): - source.write_to_cpu(torch.tensor([0, 1, 2, 3, 4], device=device), cpu, [1, 0], layout=layout) + source.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 4], device=device), + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + layout=layout, + ) else: - source.write_to_cpu(torch.tensor([0, 1, 2, 3, 4], dtype=torch.long, device=device), cpu, [1, 0], layout=layout) + source.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long, device=device), + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + layout=layout, + ) if device == "cuda": torch.cuda.synchronize() - target.read_from_cpu(cpu, [1, 0], torch.tensor([1, 2, 3, 4, 5], dtype=torch.long, device=device), layout=layout) - assert_token_fields_equal(source, [0, 1, 2, 3, 4], target, [1, 2, 3, 4, 5]) + target.read_from_cpu( + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + token_indices=torch.tensor([1, 2, 3, 4, 5], dtype=torch.long, device=device), + layout=layout, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 1, 2, 3, 4], + target=target, + target_ids=[1, 2, 3, 4, 5], + ) page = torch.empty(layout.byte_end, dtype=torch.uint8, device=device) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - source.write_to_pd_buffer(torch.tensor([0, 3], dtype=torch.long, device=device), page.unsqueeze(0), pd_page_index=0, layout=layout) + source.write_to_pd_buffer( + token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + layout=layout, + ) with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - target.read_from_pd_buffer(page.unsqueeze(0), torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, layout=layout) + target.read_from_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + pd_page_index=0, + layout=layout, + ) else: - source.write_to_pd_buffer(torch.tensor([0, 3], dtype=torch.long, device=device), page.unsqueeze(0), pd_page_index=0, layout=layout) - target.read_from_pd_buffer(page.unsqueeze(0), torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, layout=layout) - assert_token_fields_equal(source, [0, 3], target, [6, 7]) + source.write_to_pd_buffer( + token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + layout=layout, + ) + target.read_from_pd_buffer( + pd_page_buffers=page.unsqueeze(0), + token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + pd_page_index=0, + layout=layout, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 3], + target=target, + target_ids=[6, 7], + ) src_indices = torch.tensor([0, 1], dtype=torch.int32, device=device) dst_indices = torch.tensor([0, 2], dtype=torch.int64, device=device) @@ -170,7 +244,12 @@ def reject_index_conversion(tensor: torch.Tensor, *args, **kwargs): dst_token_indices=dst_indices, remote_dp_full_att_state=source, ) - assert_token_fields_equal(source, [0, 1], target, [0, 2]) + assert_token_fields_equal( + source=source, + source_ids=[0, 1], + target=target, + target_ids=[0, 2], + ) @pytest.mark.parametrize("kind", KINDS) @@ -188,7 +267,9 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): ids = torch.tensor([1, 6], dtype=torch.int32, device="cuda") state.write_full_att_kv(1, ids, kv) if isinstance(state, GroupQuantizedFullAttentionState): - quant = read_kv_storage(state)[0, ids].reshape(2, heads, -1, state.config.group_size * state.config.bits // 8) + quant = read_kv_storage(state)[0, ids].reshape( + shape=(2, heads, -1, state.config.group_size * state.config.bits // 8), + ) if state.config.bits == 4: raw = quant.to(torch.uint8) quant = torch.stack(((raw & 15).to(torch.int16) - 7, (raw >> 4).to(torch.int16) - 7), -1).flatten(-2) @@ -199,7 +280,12 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): if state.config.scale_mode == "per_tensor": scales = scales.repeat_interleave(state.config.head_num) dequant = read_kv_storage(state)[0, ids].view(torch.float8_e4m3fn).float() * scales[None, :, None] - torch.testing.assert_close(dequant, kv.float(), rtol=0.13, atol=0.03) + torch.testing.assert_close( + actual=dequant, + expected=kv.float(), + rtol=0.13, + atol=0.03, + ) else: view = state.get_layer_view(1) assert (view.indexer_k == 29).all() diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py index a7ce8744a7..e88cb3131a 100644 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -12,7 +12,13 @@ def test_quantized_kv_buffers_include_the_full_hold_page(): allocations = [] for bits in (8, 4): state = GroupQuantizedFullAttentionState( - GroupQuantizedAttentionConfig(((0,), (1,), (2,)), 2, 16, torch.float16, bits=bits), + GroupQuantizedAttentionConfig( + logical_layer_groups=((0,), (1,), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.float16, + bits=bits, + ), page_size=4, device="cpu", ) state.initialize_storage(32) diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 23378f26e5..8fcb74d304 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -21,8 +21,20 @@ def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( glm: bool, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, ): pd_args.pd_kv_page_size = 128 - source, _ = make_resources(glm, 1, 0, 2, draft_layers=2) - target, _ = make_resources(glm, 1, 0, 2, draft_layers=2) + source, _ = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + draft_layers=2, + ) + target, _ = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + draft_layers=2, + ) for index, role in enumerate(("main", "draft0")): state = getattr(source, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): @@ -39,14 +51,14 @@ def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( ("att_state", "get_pd_bounded_att_or_tail_transfer_layouts"), ): source.write_to_pd_buffer( - page.unsqueeze(0), + pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=getattr(source, getter)(), **({"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 0, "state_token_len": 8}), ) target.read_from_pd_buffer( - page.unsqueeze(0), + pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=getattr(target, getter)(), @@ -64,8 +76,10 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: remote_att_manager=source, src_token_indexes=src_indices, dst_token_indexes=dst_indices, - src_big_page_ids=[], dst_big_page_ids=[], - src_req_idx=0, dst_req_idx=2, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=2, ) for index, role in enumerate(("main", "draft0")): state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") @@ -85,18 +99,23 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: torch.cuda.synchronize() cpu = torch.empty((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) source.write_req_to_cpu( - torch.arange(8, dtype=torch.long, device="cuda"), - cpu, - [0], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=cpu, + cpu_page_indices=[0], layouts=source.get_cpu_transfer_layouts(), big_page_slot_ids=[slot], ) torch.cuda.synchronize() dst_slot = target.bounded_att_big_page_allocator.alloc(1)[0] target.read_req_from_cpu( - cpu, - [0], - torch.arange(8, 16, dtype=torch.long, device="cuda"), + cpu_page_buffers=cpu, + cpu_page_indices=[0], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), layouts=target.get_cpu_transfer_layouts(), big_page_slot_ids=[dst_slot], req_idx=1, diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index a7c6af7a42..60f6c0bdbc 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -31,7 +31,10 @@ @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_cpu_pd_and_dp_model_layouts_cover_quantized_fields( - kind: str, device: str, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + kind: str, + device: str, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") @@ -60,37 +63,99 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: page = allocate_transfer_buffer(managers[0], 1, device=device) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - copy_page(node_att_managers=[managers[0]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, write_to_pd_buffer=True) + copy_page( + node_att_managers=[managers[0]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=source_indices, + write_to_pd_buffer=True, + ) else: - copy_page(node_att_managers=[managers[0]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, write_to_pd_buffer=True) - copy_page(node_att_managers=[managers[1]], pd_page_buffers=page, pd_page_index=0, dp_rank_in_node=0, tp_world_size=1, token_indices=pd_indices, write_to_pd_buffer=False) - assert_token_fields_equal(src, [1, 3], dst, [4, 6]) + copy_page( + node_att_managers=[managers[0]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=source_indices, + write_to_pd_buffer=True, + ) + copy_page( + node_att_managers=[managers[1]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=pd_indices, + write_to_pd_buffer=False, + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[4, 6], + ) cpu_pages = torch.empty((1, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") if device == "cpu": with pytest.raises(AssertionError, match="CUDA State storage"): - managers[0].write_req_to_cpu(source_indices, cpu_pages, [0], layouts=managers[0].get_cpu_transfer_layouts()) + managers[0].write_req_to_cpu( + token_indices=source_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layouts=managers[0].get_cpu_transfer_layouts(), + ) else: - managers[0].write_req_to_cpu(source_indices, cpu_pages, [0], layouts=managers[0].get_cpu_transfer_layouts()) + managers[0].write_req_to_cpu( + token_indices=source_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layouts=managers[0].get_cpu_transfer_layouts(), + ) if device == "cuda": torch.cuda.synchronize() - managers[1].read_req_from_cpu(cpu_pages, [0], cpu_indices, layouts=managers[1].get_cpu_transfer_layouts()) - assert_token_fields_equal(src, [1, 3], dst, [5, 7]) + managers[1].read_req_from_cpu( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + token_indices=cpu_indices, + layouts=managers[1].get_cpu_transfer_layouts(), + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[5, 7], + ) if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): managers[1].copy_from_other_dp( - remote_att_manager=managers[0], src_token_indexes=source_indices, dst_token_indexes=dp_indices, - src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, dst_req_idx=0, + remote_att_manager=managers[0], + src_token_indexes=source_indices, + dst_token_indexes=dp_indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=0, ) return managers[1].copy_from_other_dp( remote_att_manager=managers[0], src_token_indexes=source_indices, dst_token_indexes=dp_indices, - src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, dst_req_idx=0, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=0, + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[0, 2], ) - assert_token_fields_equal(src, [1, 3], dst, [0, 2]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -108,9 +173,19 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans as module with torch.cuda.device(source_device): - source, source_full = make_resources(glm, 1, 0, 2) + source, source_full = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + ) with torch.cuda.device(0): - target, target_full = make_resources(glm, 1, 0, 2) + target, target_full = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + ) for i, field in enumerate(full_state_test_fields(source_full, full_state_test_buffers(source_full)).values()): if "token" in field.axes: field.tensor.fill_(i + 5) @@ -132,12 +207,20 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( source_pool = main_linear_state(source).big_pages_pool source_pool.conv_state[src_big_page_ids[0]].fill_(51) source_pool.ssm_state[src_big_page_ids[0]].fill_(53) - req = SimpleNamespace(req_idx=2, cur_kv_len=4, hold_kv_len=8, hybrid_len_to_big_page_id={}) + req = SimpleNamespace( + req_idx=2, + cur_kv_len=4, + hold_kv_len=8, + hybrid_len_to_big_page_id={}, + ) task = module.TransTask( - dst_req=req, src_dp_rank=1, src_req_idx=1, + dst_req=req, + src_dp_rank=1, + src_req_idx=1, src_token_indexes=torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=f"cuda:{source_device}"), dst_token_indexes=torch.tensor([8, 9, 10, 11]), - src_big_page_ids=src_big_page_ids, dst_big_page_ids=dst_big_page_ids, + src_big_page_ids=src_big_page_ids, + dst_big_page_ids=dst_big_page_ids, ) syncs = [] @@ -160,7 +243,12 @@ def barrier(group: str): worker.big_page_tokens = 4 worker.kv_trans([task]) assert syncs == ["node"] and req.cur_kv_len == 8 - assert_token_fields_equal(source_full, [0, 1, 2, 3], target_full, [8, 9, 10, 11]) + assert_token_fields_equal( + source=source_full, + source_ids=[0, 1, 2, 3], + target=target_full, + target_ids=[8, 9, 10, 11], + ) assert (main_linear_state(target).conv_state[:, 2, ..., :2] == 13).all() assert (main_linear_state(target).ssm_state[:, 6] == 17).all() assert main_linear_state(target).mtp_state_index[2] == 0 @@ -192,8 +280,11 @@ def barrier(group: str): (True, ("dp", "pd")), ]) def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( - monkeypatch: pytest.MonkeyPatch, p2p: bool, share_pools: bool, - ipc_kinds: tuple[str, ...], cpu_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, + p2p: bool, + share_pools: bool, + ipc_kinds: tuple[str, ...], + cpu_args: EasyDict, ): cpu_args.enable_dp_prompt_cache_fetch = share_pools from lightllm.utils.envs_utils import get_unique_server_name @@ -208,7 +299,12 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", mp.reductions.reduce_tensor.__code__) source_device = int(p2p) with torch.cuda.device(source_device): - resources, full = make_resources(True, 1, 0, 2) + resources, full = make_resources( + glm=True, + tp=1, + rank=0, + mtp=2, + ) for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.fill_(9) for storage_buffer in full_state_test_buffers(full).values(): @@ -226,7 +322,12 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( mapping[1, 0] = 7 torch.cuda.synchronize(source_device) for kind in ipc_kinds: - shared_att_resources_utils.publish_att_resources(resources, mapping, kind=kind, reader_count=1) + shared_att_resources_utils.publish_att_resources( + resources=resources, + request_mapping=mapping, + kind=kind, + reader_count=1, + ) ctx = mp.get_context("spawn") parent, child = ctx.Pipe() process = ctx.Process(target=_ipc_worker, args=(child, ipc_kinds)) @@ -263,8 +364,18 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( @pytest.mark.parametrize("source_tp,target_tp", [(1, 4), (4, 1), (2, 8), (8, 2)]) def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, target_tp: int, pd_args: EasyDict): pd_args.pd_kv_page_size = 128 - source = [make_resources(False, source_tp, rank, 2) for rank in range(source_tp)] - target = [make_resources(False, target_tp, rank, 2) for rank in range(target_tp)] + source = [make_resources( + glm=False, + tp=source_tp, + rank=rank, + mtp=2, + ) for rank in range(source_tp)] + target = [make_resources( + glm=False, + tp=target_tp, + rank=rank, + mtp=2, + ) for rank in range(target_tp)] for rank, (manager, full) in enumerate(source): for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): for local_head in range(full.config.head_num): diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index 78a8333562..d069551fc9 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -31,15 +31,27 @@ ], ) def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( - kv_type: str, expected_type: str, storage_dtype: torch.dtype, width: int, + kv_type: str, + expected_type: str, + storage_dtype: torch.dtype, + width: int, ): config = FullAttentionConfig( - logical_layer_groups=((3,), (7,)), global_head_num=8, head_dim=128, dtype=torch.bfloat16, tp_rank=2, tp_world_size=4, + logical_layer_groups=((3,), (7,)), + global_head_num=8, + head_dim=128, + dtype=torch.bfloat16, + tp_rank=2, + tp_world_size=4, ) scales = torch.full((2, 4 if kv_type == "fp8kv_sph" else 2), 0.5, dtype=torch.float32) q_scales = torch.ones((2, 2), dtype=torch.float32) if kv_type == "fp8kv_sph" else None state = create_gqa_state( - config=config, kv_type=kv_type, page_size=4, device="meta", group_size=16, + config=config, + kv_type=kv_type, + page_size=4, + device="meta", + group_size=16, ) assert state.get_state_type() == expected_type assert state.resource_id is None @@ -74,12 +86,27 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( ], ) def test_mla_creation_computes_packed_width_from_logical_dimensions( - packed: bool, kv_type: str, kv_dim: int, row_width: int, expected_type: str, + packed: bool, + kv_type: str, + kv_dim: int, + row_width: int, + expected_type: str, ): config = FullAttentionConfig( - logical_layer_groups=((2,), (6,)), global_head_num=1, head_dim=kv_dim, dtype=torch.bfloat16, tp_rank=1, tp_world_size=4, + logical_layer_groups=((2,), (6,)), + global_head_num=1, + head_dim=kv_dim, + dtype=torch.bfloat16, + tp_rank=1, + tp_world_size=4, + ) + state = create_mla_state( + config=config, + packed=packed, + kv_type=kv_type, + page_size=4, + device="meta", ) - state = create_mla_state(config=config, packed=packed, kv_type=kv_type, page_size=4, device="meta") assert state.get_state_type() == expected_type assert state.config.replicated and not config.replicated assert state.config.head_dim == row_width and config.head_dim == kv_dim @@ -95,10 +122,21 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( @pytest.mark.parametrize("create_state", [create_kda_state, create_gdn_state]) def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Callable[..., BoundedAttentionState]): config = ConvSSMStateConfig( - logical_layer_ids=(0, 2), global_k_heads=2, global_v_heads=2, - key_dim=8, value_dim=8, conv_width=3, conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, + logical_layer_ids=(0, 2), + global_k_heads=2, + global_v_heads=2, + key_dim=8, + value_dim=8, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = create_state( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", ) - state = create_state(config=config, request_capacity=3, mtp_step=2, device="cpu") assert state.config is config and state.resource_id is None assert state.conv_state.shape == (2, 3, 48, 5) assert state.ssm_state.shape == (2, 9, 2, 8, 8) @@ -107,19 +145,39 @@ def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Calla def test_kpool_creation_preserves_pool_and_mtp_dimensions(): - config = KPoolStateConfig(logical_layer_ids=(3, 7), pool_size=4, head_dim=128, dtype=torch.bfloat16) - state = create_kpool_state(config=config, request_capacity=3, mtp_step=2, device="cpu") + config = KPoolStateConfig( + logical_layer_ids=(3, 7), + pool_size=4, + head_dim=128, + dtype=torch.bfloat16, + ) + state = create_kpool_state( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", + ) assert state.config is config and state.resource_id is None assert state.tail_buffer.shape == (2, 3, 6, 256) assert state.get_state_type() == "kpool_c4_bf16" def test_create_rejects_encoding_from_another_logical_family(): - config = FullAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=576, dtype=torch.bfloat16) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=576, + dtype=torch.bfloat16, + ) with pytest.raises(AssertionError, match="unsupported GQA encoding"): create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") with pytest.raises(AssertionError, match="unsupported packed MLA encoding"): - create_mla_state(config=config, packed=True, kv_type="int4kv", device="meta") + create_mla_state( + config=config, + packed=True, + kv_type="int4kv", + device="meta", + ) with pytest.raises(AssertionError, match="unsupported MLA encoding"): create_mla_state(config=config, kv_type="fp8kv_dsa", device="meta") diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 5fe91caad1..e84d4e89d6 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -144,9 +144,14 @@ def test_set_att_state_rejects_invalid_draft_numbers(role: str): manager.set_att_state(role, definition("gqa")) -def allocated_state(kind: str, *, device: str = "cpu"): +def allocated_state(kind: str, device: str = "cpu"): template = definition(kind) - return rebuild_full_state(template, config=template.config, token_capacity=8, device=device) + return rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device=device, + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -184,8 +189,18 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp target.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) target.finish_state_creation() page = torch.zeros((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - source.write_req_to_cpu(torch.tensor([0, 1], dtype=torch.long, device="cuda"), page, [0], layouts=source.get_cpu_transfer_layouts()) - target.read_req_from_cpu(page, [0], torch.tensor([6, 7], dtype=torch.long, device="cuda"), layouts=source.get_cpu_transfer_layouts()) + source.write_req_to_cpu( + token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), + cpu_page_buffers=page, + cpu_page_indices=[0], + layouts=source.get_cpu_transfer_layouts(), + ) + target.read_req_from_cpu( + cpu_page_buffers=page, + cpu_page_indices=[0], + token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), + layouts=source.get_cpu_transfer_layouts(), + ) for src, dst in ((main, target_main), (draft, target_draft)): fields = {name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items()} for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): @@ -242,13 +257,19 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c layouts = manager.get_cpu_transfer_layouts() page = torch.zeros((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) manager.write_req_to_cpu( - torch.tensor([0, 1], dtype=torch.long, device="cuda"), page, [0], layouts=layouts + token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), + cpu_page_buffers=page, + cpu_page_indices=[0], + layouts=layouts, ) for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): field.tensor[:, 6:8].zero_() manager.read_req_from_cpu( - page, [0], torch.tensor([6, 7], dtype=torch.long, device="cuda"), layouts=layouts + cpu_page_buffers=page, + cpu_page_indices=[0], + token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), + layouts=layouts, ) for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): @@ -262,7 +283,12 @@ def test_token_copy_captures_large_index_sets_and_preserves_complete_buffers(kin state = rebuild_full_state(template, token_capacity=512, device="cuda") # 覆盖 int32/int64 和非连续索引；128 个元素足以触发旧 isin 实现的同步路径。 source = torch.arange(256, dtype=torch.int32, device="cuda")[::2] - target = torch.arange(256, 512, dtype=torch.int64, device="cuda")[::2] + target = torch.arange( + start=256, + end=512, + dtype=torch.int64, + device="cuda", + )[::2] buffers = tuple(full_state_test_buffers(state).values()) for buffer in buffers: buffer.view(torch.uint8).random_(0, 256) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 041c2c2411..166d93a64d 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -53,7 +53,12 @@ def check_created_before_allocation(manager: AttResourceManager, *args: object, "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) manager.finish_state_creation() - manager.initialize_resources(args, token_capacity=8, mem_fraction=1.0, device="cpu") + manager.initialize_resources( + args=args, + token_capacity=8, + mem_fraction=1.0, + device="cpu", + ) assert manager.get_full_attention_states() is cached[0] for role, meta in zip(("main", "draft0"), definitions): (state,) = manager.get_full_attention_states(role) diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py index 9840d4622e..490bd91aaf 100644 --- a/unit_tests/common/att_resources/test_tail_state.py +++ b/unit_tests/common/att_resources/test_tail_state.py @@ -9,10 +9,17 @@ ) -def make_tail(*, device: str | torch.device | int = "cpu", mtp_step: int = 2): +def make_tail( device: str | torch.device | int = "cpu", mtp_step: int = 2): return KPoolAttentionTailState( - KPoolStateConfig(logical_layer_ids=(3, 7), pool_size=4, head_dim=2, dtype=torch.float32), - request_capacity=3, mtp_step=mtp_step, device=device + config=KPoolStateConfig( + logical_layer_ids=(3, 7), + pool_size=4, + head_dim=2, + dtype=torch.float32, + ), + request_capacity=3, + mtp_step=mtp_step, + device=device, ) @@ -33,8 +40,18 @@ def tail_layout(state: KPoolAttentionTailState): axes=("layer", "window", "channel"), shape=shape, axis_slices=( - AxisSlice(axis="window", src_start=0, dst_start=0, length=shape[1]), - AxisSlice(axis="channel", src_start=0, dst_start=0, length=4), + AxisSlice( + axis="window", + src_start=0, + dst_start=0, + length=shape[1], + ), + AxisSlice( + axis="channel", + src_start=0, + dst_start=0, + length=4, + ), ), ), ), @@ -53,7 +70,10 @@ def test_tail_reset_clears_selected_requests_only(): @pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) @pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) def test_pd_logical_window_restores_different_mtp_rings( - device: str, source_step: int, target_step: int, length: int, + device: str, + source_step: int, + target_step: int, + length: int, ): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") @@ -65,9 +85,26 @@ def test_pd_logical_window_restores_different_mtp_rings( layout = tail_layout(source) assert layout == tail_layout(target) assert source.get_pd_transfer_regions() == target.get_pd_transfer_regions() - pages = torch.full((3, layout.byte_end + 16), 253, dtype=torch.uint8, device=device) - source.write_to_pd_buffer(1, pages, pd_page_index=1, state_token_len=length, layout=layout) - target.read_from_pd_buffer(pages, 2, pd_page_index=1, state_token_len=length, layout=layout) + pages = torch.full( + size=(3, layout.byte_end + 16), + fill_value=253, + dtype=torch.uint8, + device=device, + ) + source.write_to_pd_buffer( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) + target.read_from_pd_buffer( + pd_page_buffers=pages, + req_idx=2, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) expected = torch.zeros_like(target.tail_buffer[:, 2]) for position in range(max(0, length - 4), length): expected[:, position % (4 + target_step)] = source.tail_buffer[:, 1, position % (4 + source_step)] diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index ffa41946c8..33e7cb8eec 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -57,7 +57,13 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( if not torch.cuda.is_available(): pytest.skip("CUDA required") cpu_args.cpu_cache_token_page_size = 4 - config = FullAttentionConfig(((0,),), heads, 4, torch.float32, replicated=replicated) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=heads, + head_dim=4, + dtype=torch.float32, + replicated=replicated, + ) reference = state_cls(config, page_size=4, device="cuda") reference.initialize_storage(4) for storage_buffer in full_state_test_buffers(reference).values(): @@ -76,14 +82,24 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( assert layout.byte_end == reference_layout.byte_end for name, field in full_state_test_fields(source, full_state_test_buffers(source)).items(): field.tensor.copy_(fields[name][:, :, start:end]) - source.write_to_cpu(torch.arange(4, dtype=torch.long, device="cuda"), buffer, [0], layout=layout) + source.write_to_cpu( + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[0], + layout=layout, + ) for rank, (start, end) in enumerate(target_ranges): target = state_cls(replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda") target.initialize_storage(4) layout = cpu_layout(target) assert layout.byte_end == reference_layout.byte_end - target.read_from_cpu(buffer, [0], torch.arange(4, dtype=torch.long, device="cuda"), layout=layout) + target.read_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[0], + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + layout=layout, + ) for name, field in full_state_test_fields(target, full_state_test_buffers(target)).items(): torch.testing.assert_close(field.tensor[:, :4], fields[name][:, :4, start:end]) @@ -91,8 +107,16 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( def test_bounded_config_keeps_distinct_global_heads_when_local_heads_are_replicated(cpu_args: EasyDict): cpu_args.cpu_cache_token_page_size = 4 config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=3, conv_width=2, - conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_rank=7, tp_world_size=8, + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=4, + key_dim=2, + value_dim=3, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_rank=7, + tp_world_size=8, ) source = GDNState(config, request_capacity=1, device="meta") assert config.num_k_heads == config.num_v_heads == 1 @@ -120,18 +144,34 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( pd_args.pd_kv_page_size = page_tokens if bounded: config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=2, global_v_heads=4, key_dim=2, value_dim=2, conv_width=2, - conv_dtype=torch.float32, ssm_dtype=torch.float32, tp_world_size=8, + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=4, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=8, ) source = GDNState(config, request_capacity=1, device="meta") other_config = replace(config, logical_layer_ids=(1,), global_k_heads=4) else: - config = FullAttentionConfig(((0,),), 2, 4, torch.float32, tp_world_size=8) + config = FullAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=2, + head_dim=4, + dtype=torch.float32, + tp_world_size=8, + ) source = GQAFullAttentionState(config, page_size=4, device="meta") other_config = replace(config, logical_layer_groups=((1,),), global_head_num=4) if bounded: other = type(source)( - other_config, request_capacity=source.request_capacity, mtp_step=source.mtp_step, device="meta" + config=other_config, + request_capacity=source.request_capacity, + mtp_step=source.mtp_step, + device="meta", ) else: other = rebuild_full_state(source, config=other_config) diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index a87b88e503..381e18e202 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -30,8 +30,10 @@ def test_fp8_backend_selects_calibration_from_logical_layer( ) -> None: template = definition("static_head" if backend_name == "fa3" else "static_tensor") state = rebuild_full_state( - template, config=replace(template.config, logical_layer_groups=((3, 7), (8, 11))), - token_capacity=8, device="cpu", + template=template, + config=replace(template.config, logical_layer_groups=((3, 7), (8, 11))), + token_capacity=8, + device="cpu", ) state.scales[0].fill_(2) state.scales[1].fill_(5) @@ -81,8 +83,13 @@ def attend(**kwargs: object) -> torch.Tensor: attention.decode_max_q_seq_len = 1 attention.causal = True else: - def run(q: torch.Tensor, kv: tuple[torch.Tensor, torch.Tensor], *, - k_scale: float, v_scale: float, out: torch.Tensor) -> None: + def run( + q: torch.Tensor, + paged_kv_cache: tuple[torch.Tensor, torch.Tensor], + k_scale: float, + v_scale: float, + out: torch.Tensor, + ) -> None: assert k_scale == v_scale == expected_scale calls.append(operation) out.copy_(q) @@ -97,7 +104,12 @@ def run(q: torch.Tensor, kv: tuple[torch.Tensor, torch.Tensor], *, layer.alloc_tensor = lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu") q = torch.ones((1, state.config.head_num, state.config.head_dim)) if operation == "prefill": - output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) + output = layer._context_attention_kernel( + q=q, + kv=None, + infer_state=infer, + layer_weight=None, + ) else: output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) assert torch.equal(output, q) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index ac55020ee6..ef7c18a061 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -68,7 +68,12 @@ def record_to_cuda(self): def fake_gather(**kwargs): events.append("gather") - return torch.arange(40, 46, dtype=torch.int64, device="cuda") + return torch.arange( + start=40, + end=46, + dtype=torch.int64, + device="cuda", + ) monkeypatch.setattr(basemodel, "gather_token", fake_gather) diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index a834b8efff..3ae663ccae 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -150,7 +150,12 @@ def test_need_token_num_distinguishes_compute_and_page_allocation(): ], ) def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( - monkeypatch, hold_kv_len, can_alloc_token_num, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause + monkeypatch, + hold_kv_len, + can_alloc_token_num, + batch_max_tokens, + expected_alloc_sizes, + second_req_wait_pause, ): context, backend = _make_context(monkeypatch) backend.args.enable_cpu_cache = False @@ -466,7 +471,12 @@ def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch req_to_token_indexs=torch.arange(8, dtype=torch.int32)[None, :], ) infer_context.get_can_alloc_token_num = lambda: 4 - req = _make_paused_req(req_idx=0, target_kv_len=8, page_size=4, chunk_kv_len=4) + req = _make_paused_req( + req_idx=0, + target_kv_len=8, + page_size=4, + chunk_kv_len=4, + ) monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) monkeypatch.setattr(infer_batch, "custom_cat", lambda tensors: torch.cat(tensors)) @@ -482,7 +492,12 @@ def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): infer_context.args = SimpleNamespace(page_size=4) infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) infer_context.get_can_alloc_token_num = lambda: 12 - req = _make_paused_req(req_idx=0, target_kv_len=9, page_size=4, match_kv_len=8) + req = _make_paused_req( + req_idx=0, + target_kv_len=9, + page_size=4, + match_kv_len=8, + ) monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) infer_context.recover_paused_reqs([req], is_master_in_dp=True) @@ -517,7 +532,13 @@ def _make_req(req_idx): return req -def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv_len=None): +def _make_paused_req( + req_idx, + target_kv_len, + page_size, + match_kv_len=0, + chunk_kv_len=None, +): if chunk_kv_len is None: chunk_kv_len = target_kv_len req = SimpleNamespace( diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 2088c76ad8..421f552a9d 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -47,7 +47,12 @@ def test_nope_cache_config_uses_native_mla_width(): @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): +def test_hybrid_checkpoint_restore_and_packed_kv_copy( + monkeypatch, + small_page, + tp_world_size, + mtp_step, +): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -85,10 +90,10 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t ) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), - 1, - 584, - torch.bfloat16, + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + global_head_num=1, + head_dim=584, + dtype=torch.bfloat16, tp_world_size=tp_world_size, kv_dim=512, ), @@ -100,7 +105,9 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t main_full, draft_full = create_full_attention_roles(full, config.all_layer_num) resources = build_hybrid_resources( - config=config, request_capacity=4, mtp_step=mtp_step, + config=config, + request_capacity=4, + mtp_step=mtp_step, full_state=main_full, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft_full, @@ -117,7 +124,11 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t index_bytes = full.get_layer_view(3).indexer_k index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() - new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") + new_kv = torch.randn( + size=(2, 1, 512), + dtype=torch.bfloat16, + device="cuda", + ) destinations = torch.tensor([5, 9], dtype=torch.int32, device="cuda") full.write_full_att_kv(3, destinations, new_kv) assert torch.equal(att_kv[destinations], new_kv) @@ -171,27 +182,68 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t def test_gated_norm_gate_type_and_strided_gate(gate_type): from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward - x = torch.randn(12, 128, device="cuda", dtype=torch.bfloat16) - gate = torch.randn(3, 8, 128, device="cuda", dtype=torch.bfloat16)[:, :4] + x = torch.randn( + size=(12, 128), + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn( + size=(3, 8, 128), + device="cuda", + dtype=torch.bfloat16, + )[:, :4] weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) - actual = gated_rmsnorm_forward(x, weight, None, 1e-5, gate, gate_type=gate_type) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-5, + z=gate, + gate_type=gate_type, + ) expected = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5) * weight.float() z = gate.reshape(12, 128).float() expected *= z.sigmoid() if gate_type == "sigmoid" else torch.nn.functional.silu(z) - torch.testing.assert_close(actual, expected.bfloat16(), atol=0.008, rtol=0.008) + torch.testing.assert_close( + actual=actual, + expected=expected.bfloat16(), + atol=0.008, + rtol=0.008, + ) @pytest.mark.parametrize("add_one", [False, True]) def test_clamped_swiglu_preserves_gpt_oss_default(add_one): from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd - x = torch.linspace(-25, 25, 2048, device="cuda", dtype=torch.bfloat16).view(4, 512) - out = torch.empty(4, 256, device="cuda", dtype=torch.bfloat16) + x = torch.linspace( + start=-25, + end=25, + steps=2048, + device="cuda", + dtype=torch.bfloat16, + ).view(4, 512) + out = torch.empty( + size=(4, 256), + device="cuda", + dtype=torch.bfloat16, + ) kwargs = {} if add_one else {"clamp_up_add_one": False} alpha, limit = (1.702, 7.0) if add_one else (1.0, 10.0) - silu_and_mul_fwd(x, out, limit=limit, alpha=alpha, **kwargs) + silu_and_mul_fwd( + input=x, + output=out, + limit=limit, + alpha=alpha, + **kwargs, + ) gate, up = x.float().chunk(2, -1) gate = gate.clamp(max=limit) gate = (gate * torch.sigmoid(alpha * gate)).bfloat16().float() expected = gate * (up.clamp(-limit, limit) + int(add_one)) - torch.testing.assert_close(out, expected.bfloat16(), atol=0.008, rtol=0.008) + torch.testing.assert_close( + actual=out, + expected=expected.bfloat16(), + atol=0.008, + rtol=0.008, + ) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index cd5f7360a3..28fd6541a8 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -73,10 +73,10 @@ def create(tp, target_mtp_step=None): monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), - 1, - 584, - torch.bfloat16, + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + global_head_num=1, + head_dim=584, + dtype=torch.bfloat16, tp_rank=rank, tp_world_size=tp, kv_dim=512, @@ -86,7 +86,10 @@ def create(tp, target_mtp_step=None): full.initialize_storage(32) main, draft = create_full_attention_roles(full, config.all_layer_num) resources = build_hybrid_resources( - config=config, request_capacity=4, mtp_step=local_mtp_step, tp_rank=rank, + config=config, + request_capacity=4, + mtp_step=local_mtp_step, + tp_rank=rank, full_state=main, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft, @@ -111,7 +114,13 @@ def create(tp, target_mtp_step=None): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) @pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) -def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder, decode_mtp_step): +def test_pd_kv_and_runtime_state_roundtrip( + make_mems, + prefill_tp, + decode_tp, + remainder, + decode_mtp_step, +): source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) mtp_size = get_env_start_args().mtp_step + 1 dst_mtp_size = decode_mtp_step + 1 @@ -120,12 +129,30 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") - conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") - ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") + packed = torch.randint( + low=0, + high=256, + size=(layers, length, 1, 1168), + dtype=torch.uint8, + device="cuda", + ) + conv = torch.randn( + size=(6, 3, 8, 128, 3), + dtype=torch.bfloat16, + device="cuda", + ) + ssm = torch.randn( + size=(6, 8, 128, 128), + dtype=torch.float32, + device="cuda", + ) # Non-live slots deliberately contain stale data. The sequence length, not # the bytes in these slots, controls subsequent K-pool completion. - tail = torch.randn(layers, 3 + mtp_size, 256, dtype=torch.bfloat16, device="cuda") + tail = torch.randn( + size=(layers, 3 + mtp_size, 256), + dtype=torch.bfloat16, + device="cuda", + ) for rank, mem in enumerate(source): heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) offset = 0 @@ -227,7 +254,11 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem # Continue through the next pool boundary on P and on the restored D state. # Different physical KV/request slots must yield identical pooled FP8 bytes. count = 4 - remainder - raw = torch.randn(count, 256, dtype=torch.bfloat16, device="cuda") + raw = torch.randn( + size=(count, 256), + dtype=torch.bfloat16, + device="cuda", + ) ape = torch.randn(4, 128, device="cuda") tensor = lambda values: torch.tensor(values, dtype=torch.int32, device="cuda") outputs = [] @@ -271,48 +302,115 @@ def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, from lightllm.common.att_resources.types import StateTransferLayout torch.manual_seed(53) - config = KPoolStateConfig(logical_layer_ids=(0,), pool_size=4, head_dim=128, dtype=torch.bfloat16) - source = create_kpool_state(config=config, request_capacity=2, mtp_step=source_step, device="cuda") - target = create_kpool_state(config=config, request_capacity=2, mtp_step=target_step, device="cuda") - raw = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(64, 132, device="cuda", dtype=torch.uint8) + config = KPoolStateConfig( + logical_layer_ids=(0,), + pool_size=4, + head_dim=128, + dtype=torch.bfloat16, + ) + source = create_kpool_state( + config=config, + request_capacity=2, + mtp_step=source_step, + device="cuda", + ) + target = create_kpool_state( + config=config, + request_capacity=2, + mtp_step=target_step, + device="cuda", + ) + raw = torch.randn( + size=(64, 256), + device="cuda", + dtype=torch.bfloat16, + ) + ape = torch.randn( + size=(4, 128), + device="cuda", + dtype=torch.bfloat16, + ) + packed = torch.zeros( + size=(64, 132), + device="cuda", + dtype=torch.uint8, + ) ragged = torch.arange(64, device="cuda", dtype=torch.int32) def prefill(values, destination, tail): count = len(values) compress_pools( - raw=values, tail=tail, packed_buffer=destination, ape=ape, - lengths=torch.arange(1, count + 1, device="cuda", dtype=torch.int32), - starts=torch.zeros(count, device="cuda", dtype=torch.int32), ragged=ragged, + raw=values, + tail=tail, + packed_buffer=destination, + ape=ape, + lengths=torch.arange( + start=1, + end=count + 1, + device="cuda", + dtype=torch.int32, + ), + starts=torch.zeros(count, device="cuda", dtype=torch.int32), + ragged=ragged, req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), cu_q_lens=torch.tensor([0, count], device="cuda", dtype=torch.int32), - seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), max_q_len=count, + seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), + max_q_len=count, ) prefill(raw[:prefix], packed, source.get_layer_view(0)) regions = source.get_pd_transfer_regions() assert regions == target.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=source.resource_id, tokens_per_page=None, - byte_start=0, byte_end=regions[-1].byte_end, regions=regions, + resource_id=source.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, ) pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - source.write_to_pd_buffer(0, pages, pd_page_index=0, state_token_len=prefix, layout=layout) + source.write_to_pd_buffer( + req_idx=0, + pd_page_buffers=pages, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) target.tail_buffer.fill_(17) - target.read_from_pd_buffer(pages, 0, pd_page_index=0, state_token_len=prefix, layout=layout) + target.read_from_pd_buffer( + pd_page_buffers=pages, + req_idx=0, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) assert (target.tail_buffer[:, 1] == 17).all() width = target_step + 1 # 第一次验证后回访已接受位置；新增候选历史必须由 D 的本地 ring 正确保留。 for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: - raw[start:start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) - lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) + raw[start:start + count] = torch.randn( + size=(count, 256), + device="cuda", + dtype=torch.bfloat16, + ) + lengths = torch.arange( + start=start + 1, + end=start + count + 1, + device="cuda", + dtype=torch.int32, + ) compress_pools( - raw=raw[start:start + count], tail=target.get_layer_view(0), packed_buffer=packed, ape=ape, - lengths=lengths, starts=torch.zeros(count, device="cuda", dtype=torch.int32), ragged=ragged, + raw=raw[start:start + count], + tail=target.get_layer_view(0), + packed_buffer=packed, + ape=ape, + lengths=lengths, + starts=torch.zeros(count, device="cuda", dtype=torch.int32), + ragged=ragged, req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), cu_q_lens=torch.arange(count + 1, device="cuda", dtype=torch.int32), - seq_lens=lengths, max_q_len=1, + seq_lens=lengths, + max_q_len=1, mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py index 9464cc18a0..393d589e97 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -49,10 +49,17 @@ def test_fill_orders_radix_and_owned_pages_and_clears_reused_rows(monkeypatch: p monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) req = SimpleNamespace( - cur_kv_len=24, req_idx=7, shared_kv_node=node, + cur_kv_len=24, + req_idx=7, + shared_kv_node=node, hybrid_len_to_big_page_id={24: 3, 16: 12, 32: 88}, ) - second = SimpleNamespace(cur_kv_len=8, req_idx=5, shared_kv_node=None, hybrid_len_to_big_page_id={8: 6}) + second = SimpleNamespace( + cur_kv_len=8, + req_idx=5, + shared_kv_node=None, + hybrid_len_to_big_page_id={8: 6}, + ) module.fill_reqs_info([req, second]) assert module.shared_req_infos.arr[:, 0].tolist() == [[24, 7], [8, 5]] assert module.shared_req_big_page_ids.arr[:, 0].tolist() == [[40, 12, 3, -1], [6, -1, -1, -1]] @@ -106,7 +113,11 @@ def test_missing_source_page_fails_before_allocating(monkeypatch: pytest.MonkeyP @pytest.mark.parametrize("local_dp_rank", [0, 1]) @pytest.mark.parametrize("local_len,remote_len", [(24, 16), (24, 24), (31, 24), (0, 0)]) def test_local_longest_prefix_skips_copy_but_keeps_barriers( - monkeypatch: pytest.MonkeyPatch, hybrid: bool, local_dp_rank: int, local_len: int, remote_len: int, + monkeypatch: pytest.MonkeyPatch, + hybrid: bool, + local_dp_rank: int, + local_len: int, + remote_len: int, ): # 并列最长且本地为 DP 1 时，argmax 会优先选 DP 0；仍应在选择来源前跳过复制。 allocate_pages = Mock(side_effect=AssertionError("unexpected big page allocation")) @@ -121,9 +132,13 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( ) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = SimpleNamespace( - is_hybrid_att_model=hybrid, is_master_in_dp=True, node_nccl_group="node", + is_hybrid_att_model=hybrid, + is_master_in_dp=True, + node_nccl_group="node", model=SimpleNamespace(att_resource_manager=manager), - dp_att_resource_managers={}, dp_req_to_token_indices={}, _alloc_req_kv_mem=allocate_tokens, + dp_att_resource_managers={}, + dp_req_to_token_indices={}, + _alloc_req_kv_mem=allocate_tokens, ) module.dp_rank_in_node = local_dp_rank module.big_page_tokens = 8 if hybrid else 0 @@ -134,9 +149,13 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( ) owned_pages = {length: 70 + length // 8 for length in range(8, local_len + 1, 8)} if hybrid else {} req = SimpleNamespace( - req_idx=7, cur_kv_len=local_len, hold_kv_len=local_len, - shared_kv_node=None, hybrid_len_to_big_page_id=owned_pages.copy(), - get_cur_total_len=lambda: 32, _kv_cache_alloc_need=allocation_need, + req_idx=7, + cur_kv_len=local_len, + hold_kv_len=local_len, + shared_kv_node=None, + hybrid_len_to_big_page_id=owned_pages.copy(), + get_cur_total_len=lambda: 32, + _kv_cache_alloc_need=allocation_need, shm_req=SimpleNamespace(shm_cur_kv_len=local_len), ) monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=available_tokens)) @@ -160,7 +179,10 @@ def _build_fixture(monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: i allocator = StatePageAllocator(capacity) allocations = [] req = SimpleNamespace( - req_idx=0, cur_kv_len=8, hold_kv_len=8, hybrid_len_to_big_page_id={8: 77}, + req_idx=0, + cur_kv_len=8, + hold_kv_len=8, + hybrid_len_to_big_page_id={8: 77}, get_cur_total_len=lambda: 32, ) req._kv_cache_alloc_need = lambda length: ((length + 3) // 4) * 4 - req.hold_kv_len @@ -172,7 +194,10 @@ def allocate(cur_req: SimpleNamespace, count: int) -> torch.Tensor: module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = SimpleNamespace( - is_hybrid_att_model=True, dp_world_size=2, rank_in_dp=1, node_nccl_group=None, + is_hybrid_att_model=True, + dp_world_size=2, + rank_in_dp=1, + node_nccl_group=None, model=SimpleNamespace(att_resource_manager=SimpleNamespace(bounded_att_big_page_allocator=allocator)), dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, _alloc_req_kv_mem=allocate, diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py index 1bacdfda1d..135dd4fb4a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -33,7 +33,12 @@ def bind_before_listening(*args: object, **kwargs: object): try: with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: futures = [ - executor.submit(module._NcclControlChannel, "127.0.0.1", first_port, first_port + 4) for _ in range(2) + executor.submit( + module._NcclControlChannel, + "127.0.0.1", + first_port, + first_port + 4, + ) for _ in range(2) ] for future in futures: channels.append(future.result(timeout=15)) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py index bf5bd040a3..97f1a7e0fb 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -24,7 +24,13 @@ def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatc backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) backend.model = SimpleNamespace(req_manager=SimpleNamespace(req_to_token_indexs=table)) backend.is_master_in_dp = False - backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table(table, req, alloc_sizes, next_index, size) + backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table( + table=table, + req=req, + alloc_sizes=alloc_sizes, + next_index=next_index, + alloc_token_num=size, + ) def create_task(**kwargs): tasks.append( @@ -80,7 +86,12 @@ def test_pd_prefill_rejects_unaligned_transfer_page_size(): req = SimpleNamespace(cur_kv_len=4, shm_req=SimpleNamespace(input_len=10)) with pytest.raises(AssertionError, match="pd_kv_page_size must be divisible by page_size"): - backend._prefill_chuncked_handle_func(req, next_token_id=1, next_token_prob=1.0, output_len=0) + backend._prefill_chuncked_handle_func( + req_obj=req, + next_token_id=1, + next_token_prob=1.0, + output_len=0, + ) @pytest.mark.parametrize("tp_size,tp_rank", [(1, 0), (2, 0), (2, 1)]) @@ -108,7 +119,13 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat dp_req_to_token_indices={0: source_table, 1: source_table}, is_hybrid_att_model=False, ) - backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table(table, cur_req, alloc_sizes, next_index, size) + backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table( + table=table, + req=cur_req, + alloc_sizes=alloc_sizes, + next_index=next_index, + alloc_token_num=size, + ) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = backend module.dp_rank_in_node = 0 @@ -142,7 +159,13 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat (2, 4, 1, [1]), ], ) -def test_dp_imports_only_matching_tp_shards(monkeypatch, node_size, tp_size, rank_in_node, expected_ranks): +def test_dp_imports_only_matching_tp_shards( + monkeypatch, + node_size, + tp_size, + rank_in_node, + expected_ranks, +): from lightllm.server.router.model_infer.mode_backend import base_backend from lightllm.server.router.model_infer import shared_att_resources_utils @@ -156,7 +179,7 @@ def test_dp_imports_only_matching_tp_shards(monkeypatch, node_size, tp_size, ran ) tables[rank] = object() - def load_storage(rank, *, kind): + def load_storage(rank, kind): assert kind == "dp" imported_ranks.append(rank) return managers[rank], tables[rank] @@ -195,7 +218,13 @@ def _bind_alloc_need(req, page_size): return req -def _reserve_into_table(table, req, alloc_sizes, next_index, alloc_token_num): +def _reserve_into_table( + table, + req, + alloc_sizes, + next_index, + alloc_token_num, +): if alloc_token_num == 0: return None alloc_sizes.append(alloc_token_num) diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index a6cd2ce5e6..fa5686578d 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -76,7 +76,6 @@ def make_runtime( encoding: str, model_cls: type[TpPartBaseModel], tp: int, - *, mtp_layers: int = 0, ): hybrid = kind.startswith("qwen") or kind == "glm" @@ -119,7 +118,12 @@ def make_runtime( mtp_step=mtp_layers, ) if mtp_layers: - draft_config = dict(config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft_config = dict( + config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + ) draft_dir = tmp_path / "draft" draft_dir.mkdir() (draft_dir / "config.json").write_text(json.dumps(draft_config)) @@ -170,7 +174,12 @@ def make_runtime( draft.is_mtp_draft_model = True draft.main_model, draft.mtp_previous_draft_models = model, list(drafts) draft.args = args - draft.config = dict(config, num_hidden_layers=1, n_layer=1, full_attention_interval=1) + draft.config = dict( + config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + ) draft.max_total_token_num, draft.max_req_num = 16, 1 draft.tp_world_size_, draft.tp_rank_ = tp, rank prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) @@ -196,7 +205,15 @@ def test_cpu_service_describes_main_and_draft_as_distinct_resources( model_cls: type[TpPartBaseModel], tp: int, ): - meta, managers = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, tp, mtp_layers=2) + meta, managers = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=tp, + mtp_layers=2, + ) for resources in managers: (main,) = resources.get_full_attention_states("main") (draft,) = resources.get_full_attention_states("draft0") @@ -223,7 +240,14 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( model_cls: type[TpPartBaseModel], tp: int, ): - meta, managers = make_runtime(monkeypatch, tmp_path, kind, encoding, model_cls, tp) + meta, managers = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=tp, + ) # No byte-size formulas duplicated in the test: check the service's actual # meta-device descriptors against every real model's layout. for resources in managers: @@ -255,7 +279,13 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( if "token" in f.axes } ) - resources.write_req_to_cpu(torch.arange(8, dtype=torch.long, device="cuda"), buffer, [2], layouts=resources.get_cpu_transfer_layouts(), **kwargs) + resources.write_req_to_cpu( + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[2], + layouts=resources.get_cpu_transfer_layouts(), + **kwargs, + ) torch.cuda.synchronize() for resources, fields in zip(managers, expected): (full,) = resources.get_full_attention_states("main") @@ -268,7 +298,18 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( big_page_slot_ids=[slot], req_idx=1, ) - resources.read_req_from_cpu(buffer, [2], torch.arange(8, 16, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts(), **kwargs) + resources.read_req_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + layouts=resources.get_cpu_transfer_layouts(), + **kwargs, + ) if resources.bounded_att_big_page_allocator is not None: assert (main_linear_state(resources).conv_state[:, 1] == 3).all() assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() @@ -292,7 +333,14 @@ def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common kind, encoding, cls = ("qwen", "None", Qwen3NextTpPartModel) if hybrid else ("int4", "int4kv", LlamaTpPartModel) - meta, (resources,) = make_runtime(monkeypatch, tmp_path, kind, encoding, cls, 1) + meta, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=cls, + tp=1, + ) meta.page_num = 2 client = cpu_cache_factory(meta) pages, _ = client.allocate_pages([11, 22], disk_offload_enable=False) @@ -383,7 +431,11 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: destination = resources.bounded_att_big_page_allocator.alloc(1)[0] kwargs = dict(big_page_slot_ids=[destination], req_idx=1) resources.read_req_from_cpu( - client.cpu_kv_cache_tensor, [pages[1]], torch.arange(8, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts(), **kwargs + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=[pages[1]], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + layouts=resources.get_cpu_transfer_layouts(), + **kwargs, ) for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): if "token" in field.axes: @@ -404,14 +456,26 @@ def test_disk_preserves_resource_pages( pytest.importorskip("light_mem") from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker - meta, (resources,) = make_runtime(monkeypatch, tmp_path, "packed_fp8", "fp8kv_dsa", Deepseek3_2TpPartModel, 1) + meta, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind="packed_fp8", + encoding="fp8kv_dsa", + model_cls=Deepseek3_2TpPartModel, + tp=1, + ) meta.page_num = 2 client = cpu_cache_factory(meta) pages, _ = client.allocate_pages([11, 22], disk_offload_enable=True) (full,) = resources.get_full_attention_states("main") for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.random_(0, 256) - resources.write_req_to_cpu(torch.arange(16, dtype=torch.long, device="cuda"), client.cpu_kv_cache_tensor, pages, layouts=resources.get_cpu_transfer_layouts()) + resources.write_req_to_cpu( + token_indices=torch.arange(16, dtype=torch.long, device="cuda"), + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=pages, + layouts=resources.get_cpu_transfer_layouts(), + ) torch.cuda.synchronize() expected = client.cpu_kv_cache_tensor.clone() worker = DiskCacheWorker(0.125, client, str(tmp_path / "disk")) @@ -426,5 +490,10 @@ def test_disk_preserves_resource_pages( assert torch.equal(client.cpu_kv_cache_tensor, expected) for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.zero_() - resources.read_req_from_cpu(client.cpu_kv_cache_tensor, pages, torch.arange(16, dtype=torch.long, device="cuda"), layouts=resources.get_cpu_transfer_layouts()) + resources.read_req_from_cpu( + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=pages, + token_indices=torch.arange(16, dtype=torch.long, device="cuda"), + layouts=resources.get_cpu_transfer_layouts(), + ) assert any((tmp_path / "disk").glob("cache_file*")) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 34f5203713..27d63e1983 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -67,13 +67,20 @@ def linear_cache(monkeypatch, tmp_path, request): layer_types=("linear_attention", "full_attention"), ) full = GQAFullAttentionState( - FullAttentionConfig(tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), 1, 8, torch.bfloat16), + FullAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + global_head_num=1, + head_dim=8, + dtype=torch.bfloat16, + ), page_size=args.page_size, ) full.initialize_storage(384) requests = ReqManager(2, 512) resources = build_hybrid_resources( - config=config, request_capacity=3, mtp_step=0, + config=config, + request_capacity=3, + mtp_step=0, full_state=full, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), big_page_capacity=5, @@ -97,7 +104,10 @@ def linear_cache(monkeypatch, tmp_path, request): @pytest.mark.parametrize("token_num", [320, 384]) def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( - monkeypatch, linear_cache, cpu_cache_factory, token_num + monkeypatch, + linear_cache, + cpu_cache_factory, + token_num, ): from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common @@ -163,7 +173,9 @@ def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( module.cpu_cache_client = client module.gloo_group = None module.backend = SimpleNamespace( - is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, + is_master_in_dp=True, + is_hybrid_att_model=True, + dp_world_size=1, model=SimpleNamespace(req_manager=requests, att_resource_manager=resources), ) monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) @@ -215,10 +227,16 @@ def allocate(req, count): torch.testing.assert_close(linear.ssm_state[:, 0].cpu(), expected_states[1]) # Model data validity and multimodal resume boundaries are separate checks. - tree = HybridAttPagedRadixCache(384, 103, 64, 2, resources) + tree = HybridAttPagedRadixCache( + total_token_num=384, + rank_in_node=103, + hash_page_size=64, + big_page_num=2, + att_resource_manager=resources, + ) cache_len = max(req.hybrid_len_to_big_page_id) tree.insert( - torch.arange(cache_len), + key=torch.arange(cache_len), value=restored_indices[:cache_len].cpu(), block_hashs=list(range(1, cache_len // 64 + 1)), block_state_idxs=[None] * (cache_len // 64), @@ -226,7 +244,7 @@ def allocate(req, count): ) for length, expected in [(128, 128 if first_page_has_state else 0), (256, 256)]: node, matched, _ = tree.match_prefix( - torch.arange(length), + key=torch.arange(length), block_hashs=list(range(1, length // 64 + 1)), update_refs=True, mm_input_spans=req.mm_input_spans, @@ -234,7 +252,7 @@ def allocate(req, count): assert matched == expected tree.dec_node_ref_counter(node) req.shared_kv_node, _, _ = tree.match_prefix( - torch.arange(cache_len), + key=torch.arange(cache_len), block_hashs=list(range(1, cache_len // 64 + 1)), update_refs=True, mm_input_spans=req.mm_input_spans, @@ -245,7 +263,10 @@ def allocate(req, count): @pytest.mark.parametrize("shared_prefix", [False, True]) @pytest.mark.parametrize("allocated_pages", [0, 1, 2]) def test_offload_resolves_only_allocated_pages_and_requires_big_slots( - monkeypatch, linear_cache, shared_prefix, allocated_pages, + monkeypatch, + linear_cache, + shared_prefix, + allocated_pages, ): from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common @@ -255,8 +276,12 @@ def test_offload_resolves_only_allocated_pages_and_requires_big_slots( monkeypatch.setattr(g_infer_context, "radix_cache", SimpleNamespace(get_big_page_ids_by_node=lambda node: shared_slots)) big_slots = {128 * (i + 1): slot for i, slot in enumerate(slots) if not shared_prefix or i > 0} req = SimpleNamespace( - req_idx=0, shared_kv_node=None, hybrid_len_to_big_page_id=big_slots, - tail_small_page_buffer_id=None, cur_kv_len=256, cache_tiers=(CacheTier.CPU,), + req_idx=0, + shared_kv_node=None, + hybrid_len_to_big_page_id=big_slots, + tail_small_page_buffer_id=None, + cur_kv_len=256, + cache_tiers=(CacheTier.CPU,), shm_req=SimpleNamespace( token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), @@ -265,18 +290,25 @@ def test_offload_resolves_only_allocated_pages_and_requires_big_slots( requested_hashes = [] writes = [] - def allocate_pages(hashes, *, disk_offload_enable): + def allocate_pages(hashes, disk_offload_enable): requested_hashes.append(list(hashes)) return list(range(allocated_pages)) + [-1] * (len(hashes) - allocated_pages), [False] * len(hashes) - def write_to_cpu(token_indices, cpu_page_buffers, cpu_page_indices, **kwargs): + def write_to_cpu( + token_indices, + cpu_page_buffers, + cpu_page_indices, + **kwargs, + ): writes.append((token_indices.numel(), cpu_page_indices, kwargs["big_page_slot_ids"], kwargs["small_page_slot_ids"])) module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) module.args = get_env_start_args() module.gloo_group = None module.backend = SimpleNamespace( - is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, + is_master_in_dp=True, + is_hybrid_att_model=True, + dp_world_size=1, model=SimpleNamespace(att_resource_manager=resources, req_manager=requests), ) module.cpu_cache_client = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 699ff79ada..0284054b3f 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -51,7 +51,10 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): monkeypatch.setattr(g_infer_context, "att_resource_manager", resources) req = InferReq.__new__(InferReq) req.args = SimpleNamespace( - linear_att_hash_page_size=4, linear_att_page_block_num=2, max_req_total_len=16, page_size=4 + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + max_req_total_len=16, + page_size=4, ) req.req_idx = 1 req.cur_kv_len = req.cur_output_len = 0 @@ -93,7 +96,11 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): ], ) def test_hybrid_chunks_save_state_only_at_original_boundaries( - input_len, big_page_num, spans, expected_big, expected_small + input_len, + big_page_num, + spans, + expected_big, + expected_small, ): from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext @@ -143,7 +150,13 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( [(1, True, 1201, 1024), (1, False, 700, 512), (128, True, 0, 1024), (128, False, 640, 512)], ) def test_full_attention_gpu_restore_stays_outside_image( - monkeypatch, tmp_path, cpu_cache_factory, page_size, same_image, gpu_end, cpu_end + monkeypatch, + tmp_path, + cpu_cache_factory, + page_size, + same_image, + gpu_end, + cpu_end, ): from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache @@ -313,7 +326,9 @@ def allocate(req, count): att_resource_manager=SimpleNamespace( get_cpu_transfer_layouts=lambda: {}, bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), - read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: loaded.append(req.hold_kv_len), + read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded.append( + req.hold_kv_len + ), ), ), ) From fdcca5b19b1c111411950e3e87ab27a9440d7f12 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 08:39:12 +0000 Subject: [PATCH 13/56] refactor: simplify model attention resource configuration --- docs/cache_manager_design.md | 2 + lightllm/models/glm5_next/config.py | 19 +--- lightllm/models/glm5_next/model.py | 12 +-- lightllm/models/qwen3next/config.py | 94 +++---------------- lightllm/models/qwen3next/model.py | 6 +- lightllm/utils/backend_validator.py | 5 +- test/cpu_cache_kernel/test_speed.py | 5 - .../common/att_resources/model_helpers.py | 2 +- .../att_resources/test_model_resources.py | 10 +- .../basemodel/attention/linear/test_gdn.py | 5 +- unit_tests/models/glm5_next/test_cache.py | 17 ++-- unit_tests/models/glm5_next/test_pd_cache.py | 13 +-- .../test_hybrid_cpu_state_cache.py | 6 +- 13 files changed, 47 insertions(+), 149 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index ad81d08514..477b0aefaa 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1640,6 +1640,8 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 `model_logical_layer_ids` 表示当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；main/draft 保留现有全局编号机制。具体模型自行选择各类 State 的参与层和共享分组，该参数不代表物理存储层数，也不要求每个 State 使用全部层。 +`Qwen3NextResourceConfig` 及其 GLM 配置只描述当前模型，保留 `from_model_config()`、`load_from_args()` 和 `get_full_att_layer_indices()` 三个入口。层查询返回模型内部序号，由模型工厂转换为实际逻辑层 ID；不追加 draft 层，不计算物理层映射、Conv/SSM 存储形状或字节数。配置只保存全局 head 数，本地 TP 分片由 State 或后端校验调用方计算。full-attention 维度保存原始特征宽度，GLM 的 indexer 字段和 packed padding 由具体 State 的创建入口处理。 + 三个接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_att_resources()` 为 main 直接创建空的 `AttResourceManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttResourceManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 40d242f3f4..0e52e736a9 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING import dataclasses -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.torch_dtype_utils import get_torch_dtype from lightllm.models.qwen3next.config import Qwen3NextResourceConfig @@ -22,12 +21,8 @@ class Glm5NextResourceConfig(Qwen3NextResourceConfig): # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 index_head_dim: int = 128 - # 当前 packed MLA 行除 latent 外预留的字节数，不是 dtype 元素数。 - # 包含 128 字节量化 indexer key、4 字节 scale，并向上补齐到 16 字节边界。 - INDEX_PADDING_BYTES = 144 - @classmethod - def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict, include_draft: bool = True): + def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict) -> "Glm5NextResourceConfig": config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp @@ -37,29 +32,19 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) ] dtype = get_torch_dtype(args.data_type) - # Raw index keys and compression scores only live in request tails. - packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize return cls( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, - full_att_num_kv_heads=1, - full_att_head_dim=packed_dim, + full_att_head_dim=config["kv_lora_rank"], global_linear_k_heads=linear["num_heads"], global_linear_v_heads=linear["num_heads"], - num_linear_k_heads=linear["num_heads"] // tp, - num_linear_v_heads=linear["num_heads"] // tp, head_linear_k_dim=linear["head_dim"], head_linear_v_dim=linear["head_dim"], conv_kernel_size=linear["short_conv_kernel_size"], - linear_layer_num=len(linear["kda_layers"]), conv_state_dtype=dtype, ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=4, all_layer_num=layers, - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() - if include_draft and args.mtp_mode is not None - else 0, index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index f8ffa16bcb..a9ee2b15a2 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -51,7 +51,7 @@ def _create_bounded_att_states( config = get_resource_config(model_config, model_class=cls) assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" - cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) + cfg = Glm5NextResourceConfig.from_model_config(config, args) logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) if not logical_layer_ids: return [] @@ -122,15 +122,15 @@ def _create_full_att_states( config = get_resource_config(model_config, model_class=cls) assert args.llm_kv_type == "None", "GLM K-pool requires unquantized KV" - cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) - logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_logical_layer_ids()) + cfg = Glm5NextResourceConfig.from_model_config(config, args) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) if not logical_layer_groups: return [] state = create_mla_state( config=FullAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=1, - head_dim=config["kv_lora_rank"], + head_dim=cfg.full_att_head_dim, dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, @@ -159,8 +159,8 @@ def _create_att_tail_states( from lightllm.utils.torch_dtype_utils import get_torch_dtype config = get_resource_config(model_config, model_class=cls) - cfg = Glm5NextResourceConfig.from_model_config(config, args, include_draft=False) - logical_layer_ids = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_logical_layer_ids()) + cfg = Glm5NextResourceConfig.from_model_config(config, args) + logical_layer_ids = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) if not logical_layer_ids: return [] assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" diff --git a/lightllm/models/qwen3next/config.py b/lightllm/models/qwen3next/config.py index c5961e44e1..c40faa9e15 100644 --- a/lightllm/models/qwen3next/config.py +++ b/lightllm/models/qwen3next/config.py @@ -4,7 +4,7 @@ import torch import dataclasses -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.torch_dtype_utils import get_torch_dtype if TYPE_CHECKING: @@ -19,22 +19,14 @@ class Qwen3NextResourceConfig: # TP 切分前的全局 KV head 数，K/V 各有这么多个，不相加；GLM 的 MLA latent 记为 1。 full_att_all_num_kv_heads: int - # full-attention 布局使用的 dtype；普通 KV 为浮点精度，静态 FP8 为编码容器 uint8。 - # GLM 还用此浮点精度保存原始 indexer tail，具体量化字段由各 State 定义。 - full_att_dtype: torch.dtype - # 当前 TP rank 实际保存的 KV head 数；少于 TP 数时按组复制，MLA 为每个 rank 保存 1。 - full_att_num_kv_heads: int - # 普通 GQA 中为单个 K/V head 的维度；GLM packed MLA 中为包含 indexer/padding 的行元素数。 + # 未打包的 KV 特征维度：GQA 为单个 K/V head 的宽度，GLM MLA 为 latent 宽度。 + # 量化、indexer 字段及行 padding 由具体 full-attention State 确定，不计入本字段。 full_att_head_dim: int # TP 切分前的线性 attention K head 数；Q 使用相同的 head 数和维度。 global_linear_k_heads: int # TP 切分前的 V head 数；GDN 可与 K 不同，GLM KDA 与 K 相同。 global_linear_v_heads: int - # 当前 TP rank 的线性 K/Q head 数；Conv 的 Q/K 通道宽度由此与 head_linear_k_dim 相乘。 - num_linear_k_heads: int - # 当前 TP rank 的线性 V head 数，也是每层 SSM 矩阵的 head 轴长度。 - num_linear_v_heads: int # 每个 Q/K head 的特征维度，对应 SSM 的 key 轴，不随 TP 切分。 head_linear_k_dim: int # 每个 V head 的特征维度，对应 SSM 的 value 轴；KDA 中与 key 维度相等。 @@ -42,77 +34,26 @@ class Qwen3NextResourceConfig: # 原始卷积核宽度，单位为 token；基准 Conv 历史保存 kernel_size - 1 个位置。 # MTP 的额外候选位置由 State.mtp_step 决定，不计入本字段。 conv_kernel_size: int - # 主模型中线性 attention 层的数量，不含 full-attention 层或追加的 draft 层。 - linear_layer_num: int # Conv 运行态和大小页快照的元素 dtype，不代表算子的内部累加精度。 conv_state_dtype: torch.dtype # SSM 运行态和大小页快照的元素 dtype，可独立于 Conv 精度配置。 ssm_state_dtype: torch.dtype # 混合层排列周期：每隔这么多层放一个 full-attention 层，零起始层号为 interval - 1。 full_attention_interval: int - # 主模型总层数，包含线性与 full-attention 两类层，不包含追加的 draft 层。 + # 当前模型总层数，包含线性与 full-attention 两类层，不合并其他模型的层。 all_layer_num: int - # 需要追加到 full-attention 层布局中的 draft KV 层数；不包含 draft 的布局取 0。 - # 用于合并布局的层号/物理行换算，不代表 draft 请求数或 MTP 候选 token 数。 - draft_full_att_kv_layer_num: int = 0 - def get_conv_dim(self): - # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 - # 由于 k_dim = q_dim, k_heads = q_heads, 所以第一项和第二项的计算 - # 形式相同，但是实际内在含义是不同的。 - return ( - self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_v_dim * self.num_linear_v_heads - ) - - def get_main_model_full_att_layer_num(self): - full_att_layer_num = self.all_layer_num - self.linear_layer_num - assert full_att_layer_num == self.all_layer_num // self.full_attention_interval - return full_att_layer_num + def get_full_att_layer_indices(self) -> tuple[int, ...]: + """返回当前模型内部的 full-attention 层序号，不包含其他模型的层。 - def get_full_att_kv_layer_num_with_draft_model(self): - return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num - - def get_full_att_logical_layer_ids(self): - """Logical model layers stored in each physical full-attention row.""" - return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) + tuple( - range(self.all_layer_num, self.all_layer_num + self.draft_full_att_kv_layer_num) - ) - - def get_full_att_kv_layer_index(self, layer_index: int) -> int: - """Map a global target/draft layer index to the packed full-attention cache.""" - - layer_index = int(layer_index) - if layer_index >= self.all_layer_num: - kv_layer_index = layer_index - self.linear_layer_num - else: - kv_layer_index = layer_index // self.full_attention_interval - assert 0 <= kv_layer_index < self.get_full_att_kv_layer_num_with_draft_model(), ( - f"layer {layer_index} maps outside the packed full-attention KV cache: " - f"slot={kv_layer_index}, slots={self.get_full_att_kv_layer_num_with_draft_model()}" - ) - return kv_layer_index - - def get_conv_state_shape(self): - # Base committed sliding-window state, without speculative MTP tail. - return (self.get_conv_dim(), self.conv_kernel_size - 1) - - def get_mtp_conv_state_shape(self, mtp_step: int): - # Working state with room for S speculative tokens before acceptance. - return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step) - - def get_ssm_state_shape(self): - return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim) - - def get_conv_state_bytes_per_layer(self): - return self.get_conv_dim() * (self.conv_kernel_size - 1) * self.conv_state_dtype.itemsize - - def get_ssm_state_bytes_per_layer(self): - return self.num_linear_v_heads * self.head_linear_k_dim * self.head_linear_v_dim * self.ssm_state_dtype.itemsize + 调用方通过 model_logical_layer_ids[index] 转换为实际逻辑层 ID； + 物理存储层映射由 State 管理，draft 使用自身配置独立创建 State。 + """ + return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) @classmethod def load_from_args(cls) -> "Qwen3NextResourceConfig": + """从启动参数加载模型配置，供 attention 后端校验使用。""" args = get_env_start_args() model_path = args.model_dir from transformers.configuration_utils import PretrainedConfig @@ -125,32 +66,23 @@ def from_model_config( cls, model_cfg: dict[str, object], args: StartArgs | EasyDict, - include_draft: bool = True, - ): + ) -> "Qwen3NextResourceConfig": + """解析当前模型的全局几何参数，TP 分片与实际存储由 State 确定。""" llm_config = model_cfg.get("text_config", model_cfg) n_layer = llm_config["num_hidden_layers"] tp_world_size = args.tp // args.dp - full_att_dtype = ( - torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) - ) return cls( tp_world_size=tp_world_size, full_att_all_num_kv_heads=llm_config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=max(1, llm_config["num_key_value_heads"] // tp_world_size), full_att_head_dim=llm_config["head_dim"], global_linear_k_heads=llm_config["linear_num_key_heads"], global_linear_v_heads=llm_config["linear_num_value_heads"], - num_linear_k_heads=max(1, llm_config["linear_num_key_heads"] // tp_world_size), - num_linear_v_heads=max(1, llm_config["linear_num_value_heads"] // tp_world_size), head_linear_k_dim=llm_config["linear_key_head_dim"], head_linear_v_dim=llm_config["linear_value_head_dim"], conv_kernel_size=llm_config["linear_conv_kernel_dim"], - linear_layer_num=n_layer - (n_layer // llm_config["full_attention_interval"]), conv_state_dtype=get_torch_dtype(args.data_type), ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=llm_config["full_attention_interval"], all_layer_num=n_layer, - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if include_draft else 0, ) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 1de13e2d4d..6f6956f5c4 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -57,7 +57,7 @@ def _create_bounded_att_states( config = get_resource_config(model_config, model_class=cls) assert config["num_attention_heads"] % tp_world_size == 0 assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" - cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) + cfg = Qwen3NextResourceConfig.from_model_config(config, args) logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) if not logical_layer_ids: return [] @@ -113,8 +113,8 @@ def _create_full_att_states( config = get_resource_config(model_config, model_class=cls) assert args.llm_kv_type in ("None", "fp8kv_sph", "fp8kv_spt"), "unsupported Qwen3Next KV encoding" - cfg = Qwen3NextResourceConfig.from_model_config(config, args, include_draft=False) - logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_logical_layer_ids()) + cfg = Qwen3NextResourceConfig.from_model_config(config, args) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) if not logical_layer_groups: return [] state = create_gqa_state( diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 0d68854ebc..3458f10131 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -102,8 +102,9 @@ def _validate_flashqla(): from lightllm.models.qwen3next.config import Qwen3NextResourceConfig linear_config = Qwen3NextResourceConfig.load_from_args() - num_k_heads = linear_config.num_linear_k_heads - num_v_heads = linear_config.num_linear_v_heads + # 配置只保存全局 head 数，校验算子的输入按当前 TP 大小计算本地分片。 + num_k_heads = max(1, linear_config.global_linear_k_heads // linear_config.tp_world_size) + 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 diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 758b2dab9e..405a78e82f 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -37,17 +37,12 @@ def main(): cfg = Qwen3NextResourceConfig( tp_world_size=8, full_att_all_num_kv_heads=2, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, full_att_head_dim=256, global_linear_k_heads=16, global_linear_v_heads=64, - num_linear_k_heads=2, - num_linear_v_heads=8, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=36, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, full_attention_interval=4, diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 9f4625d462..81d8201b35 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -240,7 +240,7 @@ def build_hybrid_resources( logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), pool_size=config.index_kpool, head_dim=config.index_head_dim, - dtype=config.full_att_dtype, + dtype=state.config.dtype, tp_rank=tp_rank, tp_world_size=config.tp_world_size, ), diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 57a7127ed7..e37db5d1c8 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -40,28 +40,20 @@ def make_resources( draft_layers: int = 0, ): gk, gv = (4, 4) if glm else (2, 4) - local_k, local_v = max(1, gk // tp), max(1, gv // tp) - heads = 1 if glm else max(1, 2 // tp) cls = Glm5NextResourceConfig if glm else Qwen3NextResourceConfig config = cls( tp_world_size=tp, full_att_all_num_kv_heads=1 if glm else 2, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=heads, - full_att_head_dim=584 if glm else 16, + full_att_head_dim=512 if glm else 16, global_linear_k_heads=gk, global_linear_v_heads=gv, - num_linear_k_heads=local_k, - num_linear_v_heads=local_v, head_linear_k_dim=8, head_linear_v_dim=8, conv_kernel_size=3, - linear_layer_num=3 if glm else 2, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, full_attention_interval=4 if glm else 2, all_layer_num=4, - draft_full_att_kv_layer_num=draft_layers, ) if glm: full = PackedMLAFullAttentionState( diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 4fe82759ab..241e6cdc4f 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -179,8 +179,9 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( - num_linear_k_heads=2, - num_linear_v_heads=4, + tp_world_size=2, + global_linear_k_heads=4, + global_linear_v_heads=8, head_linear_k_dim=8, head_linear_v_dim=16, conv_state_dtype=torch.bfloat16, diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 421f552a9d..720c18909a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -40,8 +40,8 @@ def test_nope_cache_config_uses_native_mla_width(): }, StartArgs(tp=4, data_type="bfloat16"), ) - assert config.full_att_head_dim == 584 - assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 + # 模型配置保留原始 latent 宽度；packed 行宽由 State 创建入口计算。 + assert config.full_att_head_dim == 512 @pytest.mark.parametrize("small_page", [False, True]) @@ -71,26 +71,23 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( config = Glm5NextResourceConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=512, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, - num_linear_k_heads=2, - num_linear_v_heads=2, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=3, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, full_attention_interval=4, all_layer_num=4, - draft_full_att_kv_layer_num=int(mtp_step > 0), ) + full_layer_ids = config.get_full_att_layer_indices() + if mtp_step: + full_layer_ids += (config.all_layer_num,) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), global_head_num=1, head_dim=584, dtype=torch.bfloat16, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 28fd6541a8..326aa13fae 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -50,30 +50,27 @@ def create(tp, target_mtp_step=None): config = Glm5NextResourceConfig( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=512, global_linear_k_heads=8, global_linear_v_heads=8, - num_linear_k_heads=8 // tp, - num_linear_v_heads=8 // tp, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, - linear_layer_num=6, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, full_attention_interval=4, all_layer_num=8, - draft_full_att_kv_layer_num=int(mtp_step > 0), ) + full_layer_ids = config.get_full_att_layer_indices() + if mtp_step: + full_layer_ids += (config.all_layer_num,) mems = [] for rank in range(tp): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), global_head_num=1, head_dim=584, dtype=torch.bfloat16, diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 27d63e1983..6cd5713c57 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -52,13 +52,9 @@ def linear_cache(monkeypatch, tmp_path, request): config = Qwen3NextResourceConfig( tp_world_size=1, full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, full_att_head_dim=8, global_linear_k_heads=1, global_linear_v_heads=1, - num_linear_k_heads=1, - num_linear_v_heads=1, head_linear_k_dim=8, head_linear_v_dim=8, conv_kernel_size=3, @@ -68,7 +64,7 @@ def linear_cache(monkeypatch, tmp_path, request): ) full = GQAFullAttentionState( FullAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in config.get_full_att_logical_layer_ids()), + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_layer_indices()), global_head_num=1, head_dim=8, dtype=torch.bfloat16, From 7c7824906c84d2b2f30a09126f32bcba56da2f07 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 08:53:44 +0000 Subject: [PATCH 14/56] fix: upload PD token indices from CPU on the copy stream --- docs/cache_manager_design.md | 2 + .../decode_node_impl/decode_trans_process.py | 14 ++- .../prefill_trans_process.py | 14 ++- .../mode_backend/pd/resource_io.py | 3 +- .../mode_backend/test_pd_token_indices.py | 102 ++++++++++++++++++ 5 files changed, 126 insertions(+), 9 deletions(-) create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 477b0aefaa..0981bc90bb 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1661,6 +1661,8 @@ bounded/indexer 的运行态布局使用同一规则：按字段的实际最大 PD master 负责校验连接参数，长度通过传输任务的 `end_kv_index` 传递；资源层不附加 PD 页布局指纹及长度、版本字段。两个布局 getter 与页大小查询均使用 `lru_cache`，存储初始化后一起清除缓存。 +P/D 传输进程为每次 KV 任务创建普通 CPU token 索引 Tensor，在 `copy_cuda_stream` 中调用 `cuda(non_blocking=True)` 上传，后续 State 拷贝沿用同一流；预热使用相同路径。不预分配或保留 pinned 索引页缓冲，普通 CPU 内存也不保证 H2D 全程异步。`att_state` 任务按请求槽位搬运，不创建或上传 token 索引。 + CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_transfer_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 State 的全局存储字节数查询继续描述存储开销，不用于推导 CPU/PD 传输容量。CPU 联合页不预留页头、长度或版本字段。布局及其派生页大小均按需缓存，实际存储分配后一起清理，避免继续使用 meta 准备阶段的布局尺寸。 diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index a14b4210e6..b61c3bac6a 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -172,11 +172,13 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): + cpu_token_indices = torch.tensor([0], dtype=torch.long, device="cpu") + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, write_to_pd_buffer=False, - token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), + token_indices=token_indices, pd_page_index=0, dp_rank_in_node=dp_index, tp_world_size=self.dp_world_size, @@ -399,13 +401,17 @@ def read_page_to_mems_loop(self): copy_end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.stream(stream=self.copy_cuda_stream): copy_start_event.record(self.copy_cuda_stream) + # 只有 KV 页使用 token 索引；att_state 直接按请求槽位搬运运行态。 + token_indices = None + if trans_task.page_kind == "kv": + cpu_token_indices = torch.tensor(trans_task.mem_indexes, dtype=torch.long, device="cpu") + # 索引上传和后续页拷贝使用同一流，保证执行顺序。 + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, write_to_pd_buffer=False, - token_indices=torch.tensor( - trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device - ), + token_indices=token_indices, pd_page_index=trans_task.dst_page_index, dp_rank_in_node=trans_task.decode_dp_index, tp_world_size=self.dp_world_size, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index c28f5a9b06..e59a2201b6 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -147,11 +147,13 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): + cpu_token_indices = torch.tensor([0], dtype=torch.long, device="cpu") + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, write_to_pd_buffer=True, - token_indices=torch.tensor([0], dtype=torch.long, device=self.kv_move_buffer.device), + token_indices=token_indices, pd_page_index=0, dp_rank_in_node=dp_index, tp_world_size=self.dp_world_size, @@ -206,13 +208,17 @@ def local_copy_kv_loop(self): # 将kv 数据拷贝到 page 上，然后传输给 decode node，让其进行读取。 with torch.cuda.stream(stream=self.copy_cuda_stream): + # 只有 KV 页使用 token 索引；att_state 直接按请求槽位搬运运行态。 + token_indices = None + if trans_task.page_kind == "kv": + cpu_token_indices = torch.tensor(trans_task.mem_indexes, dtype=torch.long, device="cpu") + # 索引上传和后续页拷贝使用同一流，保证执行顺序。 + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, write_to_pd_buffer=True, - token_indices=torch.tensor( - trans_task.mem_indexes, dtype=torch.long, device=self.kv_move_buffer.device - ), + token_indices=token_indices, pd_page_index=trans_task.src_page_index, dp_rank_in_node=trans_task.prefill_dp_index, tp_world_size=self.dp_world_size, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index f805f61f9f..d26e56c099 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -23,7 +23,7 @@ def copy_page( pd_page_index: int, dp_rank_in_node: int, tp_world_size: int, - token_indices: torch.Tensor, + token_indices: torch.Tensor | None, write_to_pd_buffer: bool, task_kind: Literal["kv", "att_state"] = "kv", req_idx: int | None = None, @@ -33,6 +33,7 @@ def copy_page( node_att_managers 按节点内 DP 编号、组内 TP rank 顺序平铺。 pd_page_buffers 是 PD 页集合，pd_page_index 选择本次传输的页。 + token_indices 在 kv 任务中为 CUDA 索引；att_state 任务传 None，不进行索引上传。 write_to_pd_buffer=True 将 State 写入发送页；False 从接收页恢复 State。 """ tp_att_managers = node_att_managers[dp_rank_in_node * tp_world_size : (dp_rank_in_node + 1) * tp_world_size] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py new file mode 100644 index 0000000000..e73d8b0a2a --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -0,0 +1,102 @@ +"""PD 索引由普通 CPU Tensor 上传，使用实际拷贝流并跳过运行态任务。""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_trans_process +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_trans_process + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("prefill", [True, False]) +def test_pd_uploads_cpu_indices_on_copy_stream_and_skips_runtime_pages( + monkeypatch: pytest.MonkeyPatch, + prefill: bool, +) -> None: + implementation = prefill_trans_process if prefill else decode_trans_process + module_class = implementation._PrefillTransModule if prefill else implementation._DecodeTransModule + device_id = torch.cuda.current_device() + args = SimpleNamespace(tp=2, dp=2, nnodes=1, pd_kv_page_num=2, pd_kv_page_size=4, pd_node_id=1) + resource = SimpleNamespace(get_pd_page_bytes=lambda: 64) + uploads = [] + copies = [] + original_cuda = torch.Tensor.cuda + + def track_upload(tensor: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + uploads.append((tensor.device.type, tensor.is_pinned(), kwargs.get("non_blocking"), torch.cuda.current_stream())) + return original_cuda(tensor, *args, **kwargs) + + def capture_copy(**kwargs: object) -> None: + indices = kwargs["token_indices"] + copies.append((kwargs, None if indices is None else indices.clone())) + + monkeypatch.setattr(torch.Tensor, "cuda", track_upload) + monkeypatch.setattr(implementation, "copy_page", capture_copy) + monkeypatch.setattr(implementation, "create_kv_transporter", lambda **kwargs: SimpleNamespace()) + monkeypatch.setattr(implementation.threading.Thread, "start", lambda self: None) + if prefill: + module = module_class( + args=args, + device_id=device_id, + task_in_queue=None, + task_out_queue=None, + resource_managers=[resource, resource], + ) + else: + module = module_class( + args=args, + device_id=device_id, + task_in_queue=None, + task_out_queue=None, + resource_managers=[resource, resource], + up_status_in_queue=None, + ) + + # 构造时两个 DP 的预热均上传单个零索引。 + assert [indices.cpu().tolist() for _, indices in copies] == [[0], [0]] + assert [kwargs["dp_rank_in_node"] for kwargs, _ in copies] == [0, 1] + + tasks = [ + SimpleNamespace( + mem_indexes=indices, + page_kind=kind, + src_page_index=page, + dst_page_index=page, + prefill_dp_index=0, + decode_dp_index=0, + req_idx=1, + end_kv_index=7, + ) + for page, kind, indices in [(0, "kv", [5, 1, 7]), (1, "kv", [2]), (0, "att_state", None)] + ] + task_queue = module.local_copy_kv_queue if prefill else module.ready_page_task_queue + copy_loop = module.local_copy_kv_loop if prefill else module.read_page_to_mems_loop + monkeypatch.setattr(task_queue, "get", Mock(side_effect=[*tasks, _StopLoop()])) + with pytest.raises(_StopLoop): + copy_loop() + module.copy_cuda_stream.synchronize() + + copied_indices = [None if indices is None else indices.cpu().tolist() for _, indices in copies[2:]] + assert copied_indices == [[5, 1, 7], [2], None] + assert len(uploads) == 4 # 两次预热、两次 KV 上传；运行态任务不访问 mem_indexes。 + + # 下一批使用较短的索引，验证上传的数据与本次任务一致。 + tasks[0].mem_indexes = [9, 8] + monkeypatch.setattr(task_queue, "get", Mock(side_effect=[tasks[0], _StopLoop()])) + with pytest.raises(_StopLoop): + copy_loop() + module.copy_cuda_stream.synchronize() + assert copies[-1][1].cpu().tolist() == [9, 8] + assert all( + device == "cpu" and not pinned and non_blocking and stream == module.copy_cuda_stream + for device, pinned, non_blocking, stream in uploads + ) + assert all(kwargs["write_to_pd_buffer"] is prefill for kwargs, _ in copies) + assert all(indices is None or indices.is_cuda and indices.dtype == torch.long for _, indices in copies) + + +class _StopLoop(Exception): + pass From 446a2030d0caa32c13b34bca637605680a0252ee Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 09:01:23 +0000 Subject: [PATCH 15/56] fix --- lightllm/server/router/model_infer/infer_batch.py | 10 +++++++--- .../server/router/dynamic_prompt/test_radix_cache.py | 1 + .../mode_backend/test_multi_level_kv_cache.py | 9 ++++++--- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index eacf08242a..1899f69f49 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -176,11 +176,13 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 if req.tail_small_page_buffer_id is not None: - self.radix_cache._release_pages("small", [req.tail_small_page_buffer_id]) + self.att_resource_manager.bounded_att_small_page_allocator.free([req.tail_small_page_buffer_id]) req.tail_small_page_buffer_id = None # 释放请求执行期间申请、但不再插入 radix cache 的大页状态。 if req.hybrid_len_to_big_page_id: - self.radix_cache._release_pages("big", list(req.hybrid_len_to_big_page_id.values())) + self.att_resource_manager.bounded_att_big_page_allocator.free( + list(req.hybrid_len_to_big_page_id.values()) + ) req.hybrid_len_to_big_page_id.clear() # 解除请求对已命中 GPU radix cache 前缀节点的引用。 @@ -294,7 +296,9 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): # big page 模式下请求在 prefill 跨过 big page 边界后、到达末尾前被 pause / abort。 # 若不释放，会泄漏 big page state slot，并触发 free_a_req_mem 中 dict 为空的断言。 if req.hybrid_len_to_big_page_id: - self.radix_cache._release_pages("big", list(req.hybrid_len_to_big_page_id.values())) + self.att_resource_manager.bounded_att_big_page_allocator.free( + list(req.hybrid_len_to_big_page_id.values()) + ) req.hybrid_len_to_big_page_id.clear() req.cur_kv_len = shared_kv_len diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index 35c10802be..5bd3451fb9 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -661,6 +661,7 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) context.radix_cache = tree + context.att_resource_manager = tree.att_resource_manager context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(100, 132).view(1, -1)) req = InferReq.__new__(InferReq) req.req_idx = req.cur_output_len = 0 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index c97a54391d..dd7f7acd2a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -137,8 +137,9 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): context = InferenceContext() context.is_hybrid_att_model = True context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) - context.radix_cache = SimpleNamespace( - _release_pages=lambda kind, slots: (freed_small_pages if kind == "small" else freed_big_pages).extend(slots), + context.att_resource_manager = SimpleNamespace( + bounded_att_small_page_allocator=SimpleNamespace(free=freed_small_pages.extend), + bounded_att_big_page_allocator=SimpleNamespace(free=freed_big_pages.extend), ) req = SimpleNamespace( req_idx=0, @@ -182,7 +183,9 @@ def alloc_req_kv_mem(req, alloc_token_num): resources = SimpleNamespace( get_cpu_transfer_layouts=lambda: {}, - read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: loaded_indexes.extend(tokens.tolist()), + read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded_indexes.extend( + token_indices.tolist() + ), ) backend = SimpleNamespace( is_master_in_dp=False, From 010b81195151459e9fccb20e339c98c74f0c238f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 09:11:50 +0000 Subject: [PATCH 16/56] refactor: unify model configuration normalization --- docs/cache_manager_design.md | 15 +- lightllm/common/basemodel/basemodel.py | 18 ++- lightllm/common/basemodel/basemodel_utils.py | 15 +- lightllm/models/bloom/model.py | 3 +- lightllm/models/deepseek2/model.py | 3 +- lightllm/models/deepseek3_2/model.py | 3 +- lightllm/models/gemma3/model.py | 3 +- lightllm/models/gemma4/model.py | 3 +- lightllm/models/gemma_2b/model.py | 3 +- .../pre_and_post_layer_weight.py | 2 +- lightllm/models/glm4_moe_lite_mtp/model.py | 10 +- lightllm/models/glm5_next/model.py | 28 ++-- .../pre_and_post_layer_weight.py | 2 +- lightllm/models/glm5_next_mtp/model.py | 22 ++- lightllm/models/llama/model.py | 3 +- lightllm/models/mistral/model.py | 3 +- lightllm/models/mixtral/model.py | 3 +- lightllm/models/qwen3_5/model.py | 39 ++--- lightllm/models/qwen3_5_dflash/model.py | 32 ++-- lightllm/models/qwen3_5_dspark/model.py | 28 ++-- lightllm/models/qwen3_5_mtp/model.py | 16 +- lightllm/models/qwen3_dflash2/model.py | 33 ++-- lightllm/models/qwen3next/model.py | 6 +- lightllm/models/starcoder/model.py | 3 +- lightllm/models/starcoder2/model.py | 3 +- lightllm/utils/cpu_cache_layout.py | 4 +- .../test_model_config_normalization.py | 152 ++++++++++++++++++ .../models/test_qwen3_dspark_model_output.py | 76 ++++----- 28 files changed, 321 insertions(+), 210 deletions(-) create mode 100644 unit_tests/common/att_resources/test_model_config_normalization.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 0981bc90bb..ec7905c5d1 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1676,7 +1676,20 @@ State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_si full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例如复制的 MLA latent；否则按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 -`basemodel_utils.py` 仅保留 `get_resource_config(model_config, *, model_class=None)`，处理字段别名、text config 和模型专有的配置转换。原生 MTP、DFlash 等模型可通过 `_resource_config` 将原始配置转为自己的有效层配置；该工具不选择 attention 算法、不构造 State，也不通过 options 字典分派布局。 +模型类统一暴露 `normalize_model_config(model_config)`：不创建实例、不加载权重或初始化 CUDA，返回新的有效配置，重复调用保持结果不变。模型启动读取配置后、CPU 服务创建 meta State 前，以及各模型 State 工厂，都使用这个入口。DFlash 的嵌套覆盖、原生 MTP 的自有层数等差异由子类实现，不再保留 `_resource_config` 这条单独的转换路径。`basemodel_utils.py` 只处理通用 text_config 展开和字段别名，不反射调用模型方法，也不选择 attention 算法。 + +DFlash 先合并 `dflash_config`，再同步 `num_hidden_layers` / `n_layer` 等别名，避免运行时与 CPU 页布局使用不同层数。GLM 原生 MTP 的有效层数描述 draft 自身；`mtp_checkpoint_layer_start` 单独保留 checkpoint 中接在 main 层后的权重起点，用于 transformer 和 pre/post 权重加载，重复规范化不会改变这个起点。 + +```mermaid +flowchart LR + A[模型启动读取配置] --> C[模型类 normalize_model_config] + B[CPU cache 读取配置] --> C + C --> D[模型专有覆盖与通用别名同步] + D --> E[模型初始化及权重加载] + D --> F[模型 State 工厂] + F --> G[实际资源分配或 CPU meta 布局] +``` + GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `_create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 015cc71abe..823403240b 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -15,7 +15,6 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.req_manager import ReqManager from lightllm.common.att_resources import AttResourceManager -from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, @@ -201,13 +200,20 @@ def init_runtime(self) -> None: if not self.is_mtp_draft_model: self._check_max_len_infer() + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """返回当前模型的有效配置，供模型初始化和 CPU cache 共用。 + + 不依赖模型实例、权重或 CUDA；不修改输入配置，重复调用应保持结果不变。 + 子类负责模型专有的配置覆盖，通用工具仅展开 text_config 并同步字段别名。 + """ + from lightllm.common.basemodel.basemodel_utils import normalize_model_config + + return normalize_model_config(model_config) + def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: - self.config = json.load(json_file) - # rename keys - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + self.config = self.normalize_model_config(json.load(json_file)) if self.finetune_config: self.config["vocab_size"] = self.finetune_config.vocab_size return diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py index 34aefaa408..2aa5d02d37 100644 --- a/lightllm/common/basemodel/basemodel_utils.py +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -1,21 +1,10 @@ """模型启动与 CPU cache 共用的配置展开及字段别名处理。""" -from __future__ import annotations - -from typing import TYPE_CHECKING from lightllm.common.build_utils import repair_config -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - -def get_resource_config( - model_config: dict[str, object], model_class: type[TpPartBaseModel] | None = None, -) -> dict[str, object]: - """优先采用模型专有配置转换；默认展开 text_config 并补齐常见字段别名。""" - build_config = getattr(model_class, "_resource_config", None) - if build_config is not None: - return build_config(model_config) +def normalize_model_config(model_config: dict[str, object]) -> dict[str, object]: + """复制文本配置并补齐通用别名；模型专有覆盖由模型类在调用前完成。""" config = model_config.get("text_config", model_config).copy() repair_config(config, same_names=["num_attention_heads", "n_head"]) repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index b0f52c182b..fb9da7b995 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -38,10 +38,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 39e763bb56..7f90e2cddb 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -70,10 +70,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index e9baff5b26..beaf390d9a 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -26,10 +26,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index fb013cd29b..caa0563f0e 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -150,9 +150,8 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 3ab394548e..81f22d6218 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -97,10 +97,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) # 每层的 KV 槽位须容纳 sliding/full 两类 attention 的最大宽度。 head_dim = config["head_dim"] num_global_kv = config.get("num_global_key_value_heads") or config["num_key_value_heads"] diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index c417982e2f..4ba92d2ed9 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -48,10 +48,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py index 205b8d6026..5d9150fea2 100644 --- a/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py @@ -14,7 +14,7 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg, layer_idx=Non self.quant_cfg: Quantcfg = quant_cfg if layer_idx is None: - layer_idx = network_config["num_hidden_layers"] + layer_idx = network_config["mtp_checkpoint_layer_start"] hidden_size = network_config["hidden_size"] self.eh_proj_weight_ = ROWMMWeight( diff --git a/lightllm/models/glm4_moe_lite_mtp/model.py b/lightllm/models/glm4_moe_lite_mtp/model.py index 3e5c34e1df..18ceb80e78 100644 --- a/lightllm/models/glm4_moe_lite_mtp/model.py +++ b/lightllm/models/glm4_moe_lite_mtp/model.py @@ -25,10 +25,10 @@ def _pre_init(self, kvargs: dict): self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models") @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config - - config = get_resource_config(model_config) + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # checkpoint 中 MTP 权重接在 main 层之后；重复规范化不能覆盖这个起点。 + config.setdefault("mtp_checkpoint_layer_start", config["num_hidden_layers"]) count = config.get("num_nextn_predict_layers", 1) config.update(num_hidden_layers=count, n_layer=count) return config @@ -40,7 +40,7 @@ def _init_custom(self): def _init_weights(self, start_layer_index=None): assert start_layer_index is None - mtp_layer_start = self.config["num_hidden_layers"] + mtp_layer_start = self.config["mtp_checkpoint_layer_start"] num_mtp_layers = self.config.get("num_nextn_predict_layers", 1) self.pre_post_weight = self.pre_and_post_weight_class( diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index a9ee2b15a2..c652d0c5ef 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -5,7 +5,6 @@ import torch import triton -from lightllm.common.build_utils import repair_config from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend @@ -47,9 +46,8 @@ def _create_bounded_att_states( ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig - from lightllm.common.basemodel.basemodel_utils import get_resource_config - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Glm5NextResourceConfig.from_model_config(config, args) logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) @@ -74,16 +72,18 @@ def _create_bounded_att_states( ) return [state] + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """保留外层量化配置，并统一 GLM 文本配置的别名及 autotune 层数。""" + config = super().normalize_model_config(model_config) + if "quantization_config" in model_config: + config["quantization_config"] = dict(model_config["quantization_config"]) + config["autotune_layer_num"] = 4 + return config + def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: - outer_config = json.load(f) - self.config = dict(outer_config.get("text_config", outer_config)) - if "quantization_config" in outer_config: - self.config["quantization_config"] = dict(outer_config["quantization_config"]) - self.config["autotune_layer_num"] = 4 - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + self.config = self.normalize_model_config(json.load(f)) def _verify_params(self): super()._verify_params() @@ -117,10 +117,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) assert args.llm_kv_type == "None", "GLM K-pool requires unquantized KV" cfg = Glm5NextResourceConfig.from_model_config(config, args) logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) @@ -155,10 +154,9 @@ def _create_att_tail_states( ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) cfg = Glm5NextResourceConfig.from_model_config(config, args) logical_layer_ids = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) if not logical_layer_ids: diff --git a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py index 06e5576f16..40260a3be6 100644 --- a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py @@ -7,7 +7,7 @@ class Glm5NextMTPPreAndPostLayerWeight(PreAndPostLayerWeight): def __init__(self, data_type, network_config, quant_cfg): super().__init__(data_type, network_config) hidden_size = network_config["hidden_size"] - prefix = f"model.layers.{network_config['num_hidden_layers']}" + prefix = f"model.layers.{network_config['mtp_checkpoint_layer_start']}" self.eh_proj_weight_ = ROWMMWeight( in_dim=2 * hidden_size, out_dims=[hidden_size], diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index eb5bd01fd1..742c636e64 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -33,9 +33,8 @@ def __init__(self, kvargs): def _init_config(self): # Reuse the validated target configuration without changing its mHC layout. - self.config = deepcopy(self.main_model.config) + self.config = self.normalize_model_config(deepcopy(self.main_model.config)) assert self.config.get("num_nextn_predict_layers") == 1, "GLM NextN requires one native MTP block" - self.config["mhc"] = False def _init_weights(self, start_layer_index=None): assert start_layer_index is None @@ -45,7 +44,7 @@ def _init_weights(self, start_layer_index=None): self.pre_post_weight.main_norm_weight_ = self.main_model.pre_post_weight.final_norm_weight_ self.trans_layers_weight = [ self.transformer_weight_class( - layer_num=self.config["num_hidden_layers"], + layer_num=self.config["mtp_checkpoint_layer_start"], data_type=self.data_type, network_config=self.config, quant_cfg=self.quant_cfg, @@ -83,10 +82,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] @@ -119,10 +117,9 @@ def _create_att_tail_states( ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_ids = model_logical_layer_ids if not logical_layer_ids: return [] @@ -143,11 +140,12 @@ def _create_att_tail_states( return [state] @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config - - config = get_resource_config(model_config) - config.update(num_hidden_layers=1, n_layer=1) + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """区分 draft 自有层数与 checkpoint 中的 MTP 权重层起点。""" + config = super().normalize_model_config(model_config) + # checkpoint 中 MTP 权重接在 main 层之后；重复规范化不能覆盖这个起点。 + config.setdefault("mtp_checkpoint_layer_start", config["num_hidden_layers"]) + config.update(num_hidden_layers=1, n_layer=1, mhc=False) return config def _init_att_backend(self): diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index d80873e196..0d059c5516 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -46,10 +46,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index a0ca1c6faf..6118971c1c 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -40,10 +40,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 0b86d866d2..54373ba665 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -48,10 +48,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/qwen3_5/model.py b/lightllm/models/qwen3_5/model.py index 3315831f53..1dac4c1f23 100644 --- a/lightllm/models/qwen3_5/model.py +++ b/lightllm/models/qwen3_5/model.py @@ -15,7 +15,6 @@ Qwen35TransformerLayerInfer, ) from lightllm.models.qwen3_5.infer_struct import Qwen35InferStateInfo -from lightllm.common.build_utils import repair_config from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -57,35 +56,37 @@ class Qwen3_5TpPartModel(Qwen3NextTpPartModel): infer_state_class = Qwen35InferStateInfo + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """展开文本配置并补齐 Qwen3.5 的 rotary 与 MoE 默认值。""" + config = super().normalize_model_config(model_config) + rope_parameters = config.get("rope_parameters") + if isinstance(rope_parameters, dict): + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters + + # MoE routing parameters - set defaults for Qwen3.5 compatibility + if "norm_topk_prob" not in config: + config["norm_topk_prob"] = True + + return config + def _init_config(self): config_path = os.path.join(self.weight_dir_, "config.json") with open(config_path, "r") as json_file: all_config = json.load(json_file) - self.config = all_config["text_config"] + self.config = self.normalize_model_config(all_config) self.vision_config = all_config.get("vision_config", None) if self.vision_config is None: logger.warning("No vision_config found in checkpoint. " "Multimodal features may not work correctly.") - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) - - rope_parameters = self.config.get("rope_parameters") - if isinstance(rope_parameters, dict): - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters - - # MoE routing parameters - set defaults for Qwen3.5 compatibility - if "norm_topk_prob" not in self.config: - self.config["norm_topk_prob"] = True - if self.finetune_config: self.config["vocab_size"] = self.finetune_config.vocab_size self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 7a0d917f49..c51c3f4d90 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -11,26 +11,22 @@ class Qwen3_5DFlashModel(Qwen3DFlashModel): pre_and_post_weight_class = Qwen35DFlashPreAndPostLayerWeight @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config - - config = get_resource_config(model_config) - config.update(config.get("dflash_config", {})) - config["n_layer"] = config["num_hidden_layers"] + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) + config.update(dflash_config) + config = super().normalize_model_config(config) + + rope_parameters = config["rope_parameters"] + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters return config - def _init_config(self): - super()._init_config() - self.config.update(self.config.get("dflash_config", {})) - - rope_parameters = self.config["rope_parameters"] - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters - def _init_custom(self): # Draft and target use different rotary shapes, so the draft owns its rotary cache. LlamaTpPartModel._init_custom(self) diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 841653e30b..d81d5a9b2a 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -6,27 +6,23 @@ class Qwen3_5DSparkModel(Qwen3DSparkModel): """Adapter for the current Qwen3 DSpark checkpoint with a Qwen3.5 target.""" @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) + config.update(dflash_config) + config = super().normalize_model_config(config) - config = get_resource_config(model_config) - config.update(config.get("dflash_config", {})) - config["n_layer"] = config["num_hidden_layers"] - return config - - def _init_config(self): - super()._init_config() - self.config.update(self.config.get("dflash_config", {})) - - rope_parameters = self.config["rope_parameters"] - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] + rope_parameters = config["rope_parameters"] + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] # The draft is Qwen3-style and owns a 1D rotary cache. Released # checkpoints rotate the full head, while target-shaped custom # checkpoints can retain Qwen3.5's partial rotary layout. - self.config["rope_scaling"] = rope_parameters - self.config["partial_rotary_factor"] = rope_parameters.get("partial_rotary_factor", 1.0) + config["rope_scaling"] = rope_parameters + config["partial_rotary_factor"] = rope_parameters.get("partial_rotary_factor", 1.0) + return config def _init_custom(self): # Draft and target use different rotary shapes, so the draft owns its rotary cache. diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py index 0c31b7c853..27c1d455d6 100644 --- a/lightllm/models/qwen3_5_mtp/model.py +++ b/lightllm/models/qwen3_5_mtp/model.py @@ -30,22 +30,12 @@ def _init_custom(self): return @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config - - config = get_resource_config(model_config) + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """原生 MTP 只包含一个 full-attention 层，与 main 共用 checkpoint。""" + config = super().normalize_model_config(model_config) config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1) return config - def _init_config(self): - super()._init_config() - # MTP draft model: reuses the main model's config, but overrides the following: - # 因为 qwen3.5 的 mtp 和 main 是存储在一起的，所以需要进行修复。 - self.config["full_attention_interval"] = 1 - self.config["num_hidden_layers"] = 1 - self.config["n_layer"] = 1 - return - def _init_some_value(self): super()._init_some_value() self.layers_num = 1 diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py index e4485995b1..e43a534b1f 100644 --- a/lightllm/models/qwen3_dflash2/model.py +++ b/lightllm/models/qwen3_dflash2/model.py @@ -19,28 +19,23 @@ class Qwen3DFlash2Model(Qwen3DFlashModel): transformer_layer_infer_class = Qwen3DFlash2TransformerLayerInfer @classmethod - def _resource_config(cls, model_config): - from lightllm.common.basemodel.basemodel_utils import get_resource_config - - config = get_resource_config(model_config) - config.update(config.get("dflash_config", {})) - config["n_layer"] = config["num_hidden_layers"] - return config - - def _init_config(self): - super()._init_config() - dflash_config = self.config.get("dflash_config", {}) + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) if not isinstance(dflash_config, dict): raise ValueError("dflash_config must be an object in the DFlash2 checkpoint config") - self.config.update(dflash_config) + config.update(dflash_config) + config = super().normalize_model_config(config) - rope_parameters = self.config.get("rope_parameters", {}) - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if rope_parameters and "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters + rope_parameters = config.get("rope_parameters", {}) + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if rope_parameters and "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters + return config def _verify_params(self): LlamaTpPartModel._verify_params(self) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 6f6956f5c4..db1d26b740 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -52,9 +52,8 @@ def _create_bounded_att_states( ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig - from lightllm.common.basemodel.basemodel_utils import get_resource_config - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) assert config["num_attention_heads"] % tp_world_size == 0 assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Qwen3NextResourceConfig.from_model_config(config, args) @@ -108,10 +107,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) assert args.llm_kv_type in ("None", "fp8kv_sph", "fp8kv_spt"), "unsupported Qwen3Next KV encoding" cfg = Qwen3NextResourceConfig.from_model_config(config, args) logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index f78760b858..b97af1d1bd 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -36,10 +36,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 9044c78c2b..34ddb866aa 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -39,10 +39,9 @@ def _create_full_att_states( ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.torch_dtype_utils import get_torch_dtype - config = get_resource_config(model_config, model_class=cls) + config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index 5e2b996441..7ac0dfeead 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -16,7 +16,6 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: from transformers.configuration_utils import PretrainedConfig from lightllm.models import get_draft_model_class from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS - from lightllm.common.basemodel.basemodel_utils import get_resource_config from lightllm.utils.llm_utils import get_llm_model_class # 阶段一：读取 main 和启用 attention 的 draft 配置，不构造实际模型。 @@ -35,7 +34,8 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 logical_layer_start = 0 for role, model_class, config in models: - count = get_resource_config(config, model_class=model_class)["num_hidden_layers"] + config = model_class.normalize_model_config(config) + count = config["num_hidden_layers"] logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) full_states = model_class._create_full_att_states( model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, diff --git a/unit_tests/common/att_resources/test_model_config_normalization.py b/unit_tests/common/att_resources/test_model_config_normalization.py new file mode 100644 index 0000000000..4519763c26 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_config_normalization.py @@ -0,0 +1,152 @@ +"""模型启动和 CPU cache 必须使用同一份有效配置。""" + +from copy import deepcopy +import importlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3_5_dflash.model import Qwen3_5DFlashModel +from lightllm.models.qwen3_5_dspark.model import Qwen3_5DSparkModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3_dflash2.model import Qwen3DFlash2Model +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("model_class", [Qwen3DFlash2Model, Qwen3_5DFlashModel, Qwen3_5DSparkModel]) +@pytest.mark.parametrize("nested", [False, True]) +def test_dflash_runtime_and_cpu_state_use_overridden_config( + model_class: type[TpPartBaseModel], tmp_path: Path, nested: bool, +) -> None: + text_config = { + "num_hidden_layers": 32, + "n_layer": 32, + "num_attention_heads": 8, + "n_head": 8, + "hidden_size": 1024, + "num_key_value_heads": 4, + "rope_parameters": {"rope_theta": 10000, "partial_rotary_factor": 0.5}, + "dflash_config": {"num_hidden_layers": 4, "num_attention_heads": 4, "hidden_size": 512}, + } + raw_config = {"text_config": text_config} if nested else text_config + original = deepcopy(raw_config) + (tmp_path / "config.json").write_text(json.dumps(raw_config)) + + # CPU 服务只调用类接口；模型启动从同一文件走真实 _init_config 调用链。 + cpu_config = model_class.normalize_model_config(raw_config) + model = model_class.__new__(model_class) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model._init_config() + assert model.config == cpu_config + assert cpu_config["num_hidden_layers"] == cpu_config["n_layer"] == 4 + assert cpu_config["num_attention_heads"] == cpu_config["n_head"] == 4 + assert cpu_config["hidden_size"] == cpu_config["n_embd"] == cpu_config["n_embed"] == 512 + assert model_class.normalize_model_config(cpu_config) == cpu_config + assert raw_config == original + + args = StartArgs(data_type="bfloat16", llm_kv_type="None", page_size=4) + (state,) = model_class._create_full_att_states( + model_config=cpu_config, + args=args, + model_logical_layer_ids=tuple(range(cpu_config["num_hidden_layers"])), + ) + assert state.config.logical_layer_groups == ((0,), (1,), (2,), (3,)) + assert state.config.head_dim == 128 + + +@pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel, Qwen3_5MTPModel]) +def test_native_mtp_runtime_matches_normalized_config( + model_class: type[TpPartBaseModel], tmp_path: Path, +) -> None: + text_config = { + "num_hidden_layers": 32, + "num_nextn_predict_layers": 1, + "hidden_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "full_attention_interval": 4, + "mhc": True, + "scoring_func": "sigmoid", + } + original = deepcopy(text_config) + raw_config = {"text_config": text_config, "vision_config": {}} + (tmp_path / "config.json").write_text(json.dumps(raw_config)) + config = model_class.normalize_model_config(raw_config) + model = model_class.__new__(model_class) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model.tp_world_size_ = 1 + model.main_model = SimpleNamespace(config=text_config) + model._init_config() + + assert model.config == config + assert config["num_hidden_layers"] == config["n_layer"] == 1 + assert model_class.normalize_model_config(config) == config + assert text_config == original + if model_class is Qwen3_5MTPModel: + assert config["full_attention_interval"] == 1 + else: + assert config["mtp_checkpoint_layer_start"] == 32 + + +@pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel]) +def test_native_mtp_weights_keep_checkpoint_layer_start( + model_class: type[TpPartBaseModel], monkeypatch: pytest.MonkeyPatch, +) -> None: + config = model_class.normalize_model_config({ + "num_hidden_layers": 32, "num_nextn_predict_layers": 1, "hidden_size": 1024, + }) + module = importlib.import_module(model_class.pre_and_post_weight_class.__module__) + names = [] + + def init_base(self: object, data_type: torch.dtype, network_config: dict[str, object]) -> None: + self.data_type_ = data_type + + def record_weight(**kwargs: object) -> SimpleNamespace: + names.append(kwargs.get("weight_names", kwargs.get("weight_name"))) + return SimpleNamespace() + + monkeypatch.setattr(module.PreAndPostLayerWeight, "__init__", init_base) + monkeypatch.setattr(module, "ROWMMWeight", record_weight) + monkeypatch.setattr(module, "RMSNormWeight", record_weight) + quant_cfg = SimpleNamespace(get_quant_method=lambda layer, name: None) + pre_post_weight = model_class.pre_and_post_weight_class( + data_type=torch.bfloat16, network_config=config, quant_cfg=quant_cfg, + ) + assert len(names) == 4 + assert all(name.startswith("model.layers.32.") for name in names) + + model = model_class.__new__(model_class) + model.config = config + model.data_type = torch.bfloat16 + model.quant_cfg = quant_cfg + model.main_model = SimpleNamespace(pre_post_weight=SimpleNamespace( + wte_weight_=object(), lm_head_weight_=object(), final_norm_weight_=object(), + )) + model.pre_and_post_weight_class = lambda *args, **kwargs: pre_post_weight + layers = [] + + def record_layer( + layer_num: int, + data_type: torch.dtype, + network_config: dict[str, object], + quant_cfg: object, + ) -> object: + layers.append(layer_num) + return SimpleNamespace(verify_load=lambda: None) + + model.transformer_weight_class = record_layer + if model_class is Glm4MoeLiteMTPModel: + model.weight_dir_ = "unused" + model.weight_dict = None + pre_post_weight.verify_load = lambda: None + monkeypatch.setattr(importlib.import_module(model_class.__module__), "load_hf_weights", lambda *a, **kw: None) + model._init_weights() + assert layers == [32] diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py index 37fb311fbc..124fe208b4 100644 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -299,49 +299,39 @@ def test_fixed_dspark_does_not_require_confidence_head(monkeypatch): model._verify_params() -def test_qwen35_dspark_adapter_uses_current_checkpoint_rope_layout(monkeypatch): - def init_dspark_config(self): - self.config = { - "dflash_config": {"mask_token_id": 1}, - "rope_parameters": { - "rope_theta": 1_000_000, - "factor": 32.0, - "original_max_position_embeddings": 8192, - "rope_type": "yarn", - }, - } - - monkeypatch.setattr(Qwen3DSparkModel, "_init_config", init_dspark_config) - model = Qwen3_5DSparkModel.__new__(Qwen3_5DSparkModel) - - model._init_config() - - assert model.config["rope_scaling"] == model.config["rope_parameters"] - assert model.config["rope_theta"] == 1_000_000 - assert model.config["partial_rotary_factor"] == 1.0 - assert model.config["mask_token_id"] == 1 - - -def test_qwen35_dspark_adapter_preserves_custom_partial_rotary_layout(monkeypatch): - def init_dspark_config(self): - self.config = { - "dflash_config": {"mask_token_id": 1}, - "rope_parameters": { - "mrope_interleaved": True, - "mrope_section": [11, 11, 10], - "partial_rotary_factor": 0.25, - "rope_theta": 10_000_000, - "rope_type": "default", - }, - } - - monkeypatch.setattr(Qwen3DSparkModel, "_init_config", init_dspark_config) - model = Qwen3_5DSparkModel.__new__(Qwen3_5DSparkModel) - - model._init_config() - - assert model.config["partial_rotary_factor"] == 0.25 - assert model.config["rope_scaling"] == model.config["rope_parameters"] +def test_qwen35_dspark_adapter_uses_current_checkpoint_rope_layout(): + raw_config = { + "dflash_config": {"mask_token_id": 1}, + "rope_parameters": { + "rope_theta": 1_000_000, + "factor": 32.0, + "original_max_position_embeddings": 8192, + "rope_type": "yarn", + }, + } + config = Qwen3_5DSparkModel.normalize_model_config(raw_config) + + assert config["rope_scaling"] == config["rope_parameters"] + assert config["rope_theta"] == 1_000_000 + assert config["partial_rotary_factor"] == 1.0 + assert config["mask_token_id"] == 1 + + +def test_qwen35_dspark_adapter_preserves_custom_partial_rotary_layout(): + raw_config = { + "dflash_config": {"mask_token_id": 1}, + "rope_parameters": { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "rope_type": "default", + }, + } + config = Qwen3_5DSparkModel.normalize_model_config(raw_config) + + assert config["partial_rotary_factor"] == 0.25 + assert config["rope_scaling"] == config["rope_parameters"] def test_qwen3_parallel_block_draft_applies_partial_rotary_factor(monkeypatch): From 20e517d3180c28b12341b4a1ec7caa3489d0a693 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 10:49:45 +0000 Subject: [PATCH 17/56] refactor: expose state factories and optimize bounded attention copies --- docs/CN/source/models/add_new_model.md | 2 +- docs/EN/source/models/add_new_model.md | 2 +- docs/cache_manager_design.md | 41 +- .../att_resources/bounded_att/copy_kernel.py | 571 ++++++++++++++++++ .../att_resources/bounded_att/gdn_state.py | 130 ++-- .../att_resources/bounded_att/kda_state.py | 130 ++-- .../bounded_att/transfer_utils.py | 46 +- lightllm/common/basemodel/basemodel.py | 12 +- lightllm/models/bloom/model.py | 2 +- lightllm/models/deepseek2/model.py | 2 +- lightllm/models/deepseek3_2/model.py | 2 +- lightllm/models/gemma3/model.py | 2 +- .../layer_infer/transformer_layer_infer.py | 2 +- lightllm/models/gemma4/model.py | 2 +- lightllm/models/gemma_2b/model.py | 2 +- lightllm/models/glm5_next/model.py | 6 +- lightllm/models/glm5_next_mtp/model.py | 4 +- lightllm/models/llama/model.py | 2 +- lightllm/models/mistral/model.py | 2 +- lightllm/models/mixtral/model.py | 2 +- lightllm/models/qwen3next/model.py | 4 +- lightllm/models/starcoder/model.py | 2 +- lightllm/models/starcoder2/model.py | 2 +- lightllm/utils/cpu_cache_layout.py | 4 +- .../att_resources/test_bounded_copy_kernel.py | 264 ++++++++ .../att_resources/test_bounded_state.py | 31 +- .../common/att_resources/test_capacity.py | 20 +- .../common/att_resources/test_layer_groups.py | 2 +- .../test_layer_state_selection.py | 2 +- .../att_resources/test_manager_helpers.py | 13 +- .../att_resources/test_model_calibration.py | 2 +- .../test_model_config_normalization.py | 2 +- .../test_model_storage_definition.py | 4 +- .../att_resources/test_runtime_transport.py | 10 +- 34 files changed, 1094 insertions(+), 232 deletions(-) create mode 100644 lightllm/common/att_resources/bounded_att/copy_kernel.py create mode 100644 unit_tests/common/att_resources/test_bounded_copy_kernel.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index abddb2009b..1f0412cddc 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -124,7 +124,7 @@ class TpPartBaseModel: | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | 读取初始化模型的 config.json, 并进行一些 key 名的同名合法化操作 | | def _verify_params(self)： | 校验参数 | -| `_create_full_att_states(...)` / `_create_bounded_att_states(...)` / `_create_att_tail_states(...)` | 创建并返回当前模型的 full / bounded / tail State 列表；默认返回空列表，公共流程统一注册 | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` / `create_att_tail_states(...)` | 创建并返回当前模型的 full / bounded / tail State 列表；默认返回空列表，公共流程统一注册 | | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index fbb3776a80..dc18d4da91 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -124,7 +124,7 @@ Commonly used interfaces that need to be inherited and implemented | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | Read the config.json for initializing the model and perform some key name legalization operations | | def _verify_params(self)： | Validate parameters | -| `_create_full_att_states(...)` / `_create_bounded_att_states(...)` / `_create_att_tail_states(...)` | Return this model's full / bounded / tail State lists; defaults return empty lists and the shared startup stage registers the results | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` / `create_att_tail_states(...)` | Return this model's full / bounded / tail State lists; defaults return empty lists and the shared startup stage registers the results | | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index ec7905c5d1..9a3244952e 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -310,7 +310,8 @@ lightllm/common/ │ │ ├── __init__.py │ │ ├── base_state.py # BoundedAttentionState：有界状态的抽象操作契约 │ │ ├── utils.py # 页池分配策略：DP fetch 使用共享 pinned 存储 -│ │ ├── transfer_utils.py # bounded_att_buffer_copies：复制完整状态快照 +│ │ ├── transfer_utils.py # bounded_att_buffer_copies：生成快照字段的复制视图 +│ │ ├── copy_kernel.py # copy_bounded_att_tensors：聚合多页、多字段的 Triton 拷贝 │ │ ├── conv_ssm_types.py # KDA/GDN 的 Conv/SSM 布局配置 │ │ ├── create_kda_state.py # KDA 统一创建入口 │ │ ├── kda_state.py # KDAState：Conv / SSM / MTP 与页操作 @@ -340,9 +341,29 @@ lightllm/utils/ `KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_pages`、`restore_bounded_att_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 -`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `_create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 +`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；其余复制视图函数返回源、目标视图列表，由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个私有流程函数顺序执行：`_view_min_dimensions` 合并连续尾维，`_view_max_dtype` 整理最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧不连续就停止，保留前面的维度及单元素轴，再按行宽、stride、storage_offset 及整批地址对齐条件创建最大 8 字节的整数 view，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 + +默认参数在 H200 上固定 grid=24 联合调优，覆盖 Conv BF16 / SSM FP32、三种规模和 CUDA↔CUDA / CUDA↔pinned CPU / pinned CPU↔pinned CPU，共 24 个场景。第一轮比较 BLOCK_ELEMENTS=512/1024/2048/4096 与 warps=1/2/4/8，并保留 256/8/1 对照；第二轮为入选组合比较 num_stages=1/2/3/4，两轮共测试 29 个不同组合。每个候选均校验数据，预先创建指针表，以随机顺序进行 CUDA Graph 计时并取中位数。按场景等权的速度几何平均选择 1024/8/1；相比旧默认 grid=16、BLOCK_ELEMENTS=256、warps=8、stages=1，第二轮综合速度约为 1.58 倍。该组合并非每项最快，小块场景最多增加约 1.35 μs（Conv CPU→GPU 约 2.42→3.77 μs）。`grid_num`、`num_warps`、`num_stages` 可由调用方覆盖；流水深度用于 `tl.range` 的循环属性。 + +另外使用不同地址的独立块连续发射复测：Conv 为 64×3 MiB，SSM 为 16×64 MiB，块间不做同步，九轮取平均。下表为有效载荷带宽（GB/s），不含 Python 准备、指针上传和网络传输；旧默认与新默认在同一轮实验对照。未测量与推理计算并发时的性能影响。 + +| 数据类型 | 方向 | 旧默认 16/256/8/1 | 新默认 24/1024/8/1 | +| --- | --- | ---: | ---: | +| Conv | GPU→GPU | 16.54 | 59.78 | +| Conv | GPU→CPU | 16.74 | 40.73 | +| Conv | CPU→GPU | 6.22 | 26.70 | +| Conv | CPU→CPU | 19.51 | 36.32 | +| SSM | GPU→GPU | 76.83 | 303.01 | +| SSM | GPU→CPU | 50.61 | 51.99 | +| SSM | CPU→GPU | 22.51 | 50.11 | +| SSM | CPU→CPU | 19.91 | 36.88 | + +CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 + + +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；tail 的复制视图由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttResourceManager.get_cpu_transfer_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 @@ -1624,7 +1645,7 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 -2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `_create_full_att_states()`、`_create_bounded_att_states()`、`_create_att_tail_states()`，再通过 `set_att_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 +2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `create_full_att_states()`、`create_bounded_att_states()`、`create_att_tail_states()`，再通过 `set_att_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_creation()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 仍必须有 K/V 和 Q 标定。随后 backend 调用 `initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 @@ -1634,20 +1655,20 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 | 创建接口 | 返回值与职责 | | --- | --- | -| `_create_full_att_states(...)` | 返回 `list[FullAttentionState]`；只准备布局，不加载静态标定或分配 KV 容量；可返回多种量化或存储方案 | -| `_create_bounded_att_states(...)` | 返回 `list[BoundedAttentionState]`；按请求容量创建运行态，大小页延后分配；CPU 查询使用 `device="meta"` | -| `_create_att_tail_states(...)` | 返回 `list[AttentionTailState]`；按请求容量创建本模型 tail，CPU 联合页不调用此接口 | +| `create_full_att_states(...)` | 返回 `list[FullAttentionState]`；只准备布局，不加载静态标定或分配 KV 容量；可返回多种量化或存储方案 | +| `create_bounded_att_states(...)` | 返回 `list[BoundedAttentionState]`；按请求容量创建运行态，大小页延后分配；CPU 查询使用 `device="meta"` | +| `create_att_tail_states(...)` | 返回 `list[AttentionTailState]`；按请求容量创建本模型 tail，CPU 联合页不调用此接口 | `model_logical_layer_ids` 表示当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；main/draft 保留现有全局编号机制。具体模型自行选择各类 State 的参与层和共享分组，该参数不代表物理存储层数，也不要求每个 State 使用全部层。 `Qwen3NextResourceConfig` 及其 GLM 配置只描述当前模型，保留 `from_model_config()`、`load_from_args()` 和 `get_full_att_layer_indices()` 三个入口。层查询返回模型内部序号，由模型工厂转换为实际逻辑层 ID；不追加 draft 层，不计算物理层映射、Conv/SSM 存储形状或字节数。配置只保存全局 head 数，本地 TP 分片由 State 或后端校验调用方计算。full-attention 维度保存原始特征宽度，GLM 的 indexer 字段和 packed padding 由具体 State 的创建入口处理。 -三个接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_att_resources()` 为 main 直接创建空的 `AttResourceManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttResourceManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 +三个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_att_resources()` 为 main 直接创建空的 `AttResourceManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttResourceManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -`_create_full_att_states()` 由具体模型直接构造 `FullAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 +`create_full_att_states()` 由具体模型直接构造 `FullAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 `AttResourceManager.initialize_resources()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 @@ -1691,7 +1712,7 @@ flowchart LR ``` -GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `_create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 +GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 diff --git a/lightllm/common/att_resources/bounded_att/copy_kernel.py b/lightllm/common/att_resources/bounded_att/copy_kernel.py new file mode 100644 index 0000000000..580c526f5a --- /dev/null +++ b/lightllm/common/att_resources/bounded_att/copy_kernel.py @@ -0,0 +1,571 @@ +"""将 bounded attention 的同一形状的多页视图聚合到 CUDA stream 拷贝。""" + +import torch +import triton +import triton.language as tl + + +# ======================================================================== +# 一、聚合入口：只上传指针表，shape/stride 作为常量 +# ======================================================================== + + +def copy_bounded_att_tensors( + copies: list[tuple[torch.Tensor, torch.Tensor]], + grid_num: int = 24, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """在当前 CUDA stream 上聚合复制同一形状的视图，默认 24 个 program；warp 数与循环流水深度可按搬运场景覆盖。 + + num_stages 控制 tl.range 的循环流水深度，默认 1；提高该值不保证所有方向的带宽都增加。 + 支持 CUDA/IPC 存储和本进程已 pinned/registered 的 CPU 存储，两端均为 + CPU 页池时也走 GPU。整批所有源、目标的形状和 dtype 必须相同，支持最多六维的 stride。 + 同一侧的全部视图共用 stride，源与目标之间可不同；入口只核对首尾项。 + 只合并两端共同连续的尾部维度，并按对齐条件创建最大 8 字节的整数视图。 + 这些 view 不复制数据；不满足合并或宽类型条件时保留原布局，不转换数值。 + 每项表示独立复制，目标区域不得相互覆盖，也不得覆盖本批仍需读取的源数据。 + 直接使用 data_ptr()；调用方保证地址可被当前 GPU 访问，并持有存储直到拷贝完成。 + 源数据就绪，以及 CPU 读取或复用页面前的完成事件，也由调用方管理。 + """ + if not copies: + return + assert grid_num > 0, "copy grid must be positive" + # 阶段一：只合并两端共同连续的尾部维度，保留前面的布局。 + copies = _view_min_dimensions(copies=copies) + # 阶段二：按整批地址和行布局选择最宽的搬运单元，不转换数值。 + copies = _view_max_dtype(copies=copies) + # 阶段三：上传两张指针表，按最终视图维数发射聚合拷贝算子。 + _launch_copy_kernel( + copies=copies, + grid_num=grid_num, + num_warps=num_warps, + num_stages=num_stages, + ) + + +# ======================================================================== +# 二、视图整理与算子发射：降维 → 扩大搬运单元 → kernel +# ======================================================================== + + +def _view_min_dimensions( + copies: list[tuple[torch.Tensor, torch.Tensor]], +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """按首对布局合并连续尾维，遇到任一侧不连续就停止。""" + # 每批复制使用同一种 shape；由调用方分别收集 Conv、SSM 或同一传输区域。 + source, destination = copies[0] + last_source, last_destination = copies[-1] + # 整批布局由调用方保证一致，只检查首尾两对 Tensor。 + assert source.shape == destination.shape == last_source.shape == last_destination.shape, "copy shapes must match" + assert source.dtype == destination.dtype == last_source.dtype == last_destination.dtype, "copy dtypes must match" + ndim = source.ndim + assert 1 <= ndim <= 6, "copy tensors must have one to six dimensions" + assert source.numel() > 0, "copy tensors must be nonempty" + shape = tuple(source.shape) + # 同一侧所有视图共用 stride；只核对首尾，源和目标之间允许不同布局。 + src_stride = source.stride() + dst_stride = destination.stride() + assert src_stride == last_source.stride(), "source strides must match across the batch" + assert dst_stride == last_destination.stride(), "destination strides must match across the batch" + # 从最后一维向前扩展连续区间；任一侧有间隙就停止，不再处理前面的轴。 + # 例如 (2, 3, 4, 5) 的末轴带 padding 时保持原形状，不额外合并前面的 (2, 3, 4)。 + first_tail_axis = ndim - 1 + tail_elements = shape[-1] + while first_tail_axis > 0: + outer_axis = first_tail_axis - 1 + if ( + src_stride[outer_axis] != src_stride[first_tail_axis] * shape[first_tail_axis] + or dst_stride[outer_axis] != dst_stride[first_tail_axis] * shape[first_tail_axis] + ): + break + tail_elements *= shape[outer_axis] + first_tail_axis = outer_axis + if first_tail_axis == ndim - 1: + return copies + transfer_shape = shape[:first_tail_axis] + (tail_elements,) + return [(source.view(transfer_shape), destination.view(transfer_shape)) for source, destination in copies] + + +def _view_max_dtype( + copies: list[tuple[torch.Tensor, torch.Tensor]], +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """将整批 view 为最大可用整数宽度，保留原始位模式和存储间隙。""" + source, destination = copies[0] + if source.stride(-1) != 1 or destination.stride(-1) != 1: + return copies + + # 每个切片的起点可能不同；地址和 storage_offset 都须支持相同的宽类型 view。 + # 累计低位即可检查整批对齐，不能仅检查首尾，否则会漏掉中间错位的视图。 + pointer_alignment_bits = 0 + offset_alignment_bits = 0 + for src_tensor, dst_tensor in copies: + pointer_alignment_bits |= src_tensor.data_ptr() | dst_tensor.data_ptr() + offset_alignment_bits |= src_tensor.storage_offset() | dst_tensor.storage_offset() + + # 最内维连续时，将相邻元素作为一个更宽的整数搬运；优先 8 字节，其次 4/2 字节。 + # 行宽、外层 stride、storage_offset 及整批地址都须对齐，避免跨行 padding 或错位读取。 + element_bytes = source.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + pointer_alignment_bits % transfer_bytes != 0 + or source.shape[-1] % elements_per_unit != 0 + or offset_alignment_bits % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in source.stride()[:-1]) + or any(stride % elements_per_unit for stride in destination.stride()[:-1]) + ): + continue + return [(src.view(transfer_dtype), dst.view(transfer_dtype)) for src, dst in copies] + return copies + + +def _launch_copy_kernel( + copies: list[tuple[torch.Tensor, torch.Tensor]], + grid_num: int, + num_warps: int, + num_stages: int, +) -> None: + """使用整理后的视图地址和共享布局，在当前流上发射一次聚合拷贝。""" + # 确保当前设备上下文已经建立；纯 CPU 页池之间的复制也使用这个设备和流。 + stream = torch.cuda.current_stream() + source, destination = copies[0] + source_pointers, destination_pointers = [], [] + for src_tensor, dst_tensor in copies: + source_pointers.append(src_tensor.data_ptr()) + destination_pointers.append(dst_tensor.data_ptr()) + + # 整批共享 shape/dtype 和同侧 stride，首对最终视图提供全部编译期布局参数。 + ndim = source.ndim + shape = tuple(source.shape) + src_stride = source.stride() + dst_stride = destination.stride() + unsigned_dtypes = {1: tl.uint8, 2: tl.uint16, 4: tl.uint32, 8: tl.uint64} + + kernels = { + 1: _copy_bounded_att_1d, + 2: _copy_bounded_att_2d, + 3: _copy_bounded_att_3d, + 4: _copy_bounded_att_4d, + 5: _copy_bounded_att_5d, + 6: _copy_bounded_att_6d, + } + # 源、目标地址分别上传为独立 Tensor，同一位置对应一对复制视图。 + source_pointer_tensor = torch.tensor( + data=source_pointers, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(device=stream.device, non_blocking=True) + destination_pointer_tensor = torch.tensor( + data=destination_pointers, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(device=stream.device, non_blocking=True) + # 固定 grid=24 的跨场景调优使用每轮 1024 个搬运单元，兼顾小块延迟和大块吞吐。 + kernels[ndim][(grid_num,)]( + SRC_POINTERS=source_pointer_tensor, + DST_POINTERS=destination_pointer_tensor, + total_element_count=source.numel() * len(copies), + SHAPE=shape, + SRC_STRIDES=src_stride, + DST_STRIDES=dst_stride, + DTYPE=unsigned_dtypes[source.element_size()], + BLOCK_ELEMENTS=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 三、分维数聚合算子：计算索引与偏移 → 读取指针 → load 数据 → store 数据 +# ======================================================================== + + +@triton.jit +def _copy_bounded_att_1d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 一维直接按元素序号应用两端 stride。 + axis0 = element_indices + source_offsets = axis0 * src_stride_0 + destination_offsets = axis0 * dst_stride_0 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_2d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 2 维坐标；两端可以具有不同的物理 stride。 + axis1 = element_indices % SHAPE[1] + remaining_indices = element_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_3d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 3 维坐标；两端可以具有不同的物理 stride。 + axis2 = element_indices % SHAPE[2] + remaining_indices = element_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_4d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 4 维坐标；两端可以具有不同的物理 stride。 + axis3 = element_indices % SHAPE[3] + remaining_indices = element_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + + axis3 * src_stride_3 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + + axis3 * dst_stride_3 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_5d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] * SHAPE[4] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + src_stride_4 = tl.cast(input=SRC_STRIDES[4], dtype=tl.int64) + dst_stride_4 = tl.cast(input=DST_STRIDES[4], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 5 维坐标；两端可以具有不同的物理 stride。 + axis4 = element_indices % SHAPE[4] + remaining_indices = element_indices // SHAPE[4] + axis3 = remaining_indices % SHAPE[3] + remaining_indices = remaining_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + + axis3 * src_stride_3 + + axis4 * src_stride_4 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + + axis3 * dst_stride_3 + + axis4 * dst_stride_4 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_6d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] * SHAPE[4] * SHAPE[5] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + src_stride_4 = tl.cast(input=SRC_STRIDES[4], dtype=tl.int64) + dst_stride_4 = tl.cast(input=DST_STRIDES[4], dtype=tl.int64) + src_stride_5 = tl.cast(input=SRC_STRIDES[5], dtype=tl.int64) + dst_stride_5 = tl.cast(input=DST_STRIDES[5], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 6 维坐标；两端可以具有不同的物理 stride。 + axis5 = element_indices % SHAPE[5] + remaining_indices = element_indices // SHAPE[5] + axis4 = remaining_indices % SHAPE[4] + remaining_indices = remaining_indices // SHAPE[4] + axis3 = remaining_indices % SHAPE[3] + remaining_indices = remaining_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + + axis3 * src_stride_3 + + axis4 * src_stride_4 + + axis5 * src_stride_5 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + + axis3 * dst_stride_3 + + axis4 * dst_stride_4 + + axis5 * dst_stride_5 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py index 1cedb23e59..dd57007b6b 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -13,6 +13,7 @@ from .conv_ssm_types import ConvSSMStateConfig from .gdn_pages_pool import GDNPagesPool from .transfer_utils import bounded_att_buffer_copies +from .copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: from lightllm.common.att_resources.types import PageKind @@ -129,27 +130,28 @@ def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, sl assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" pool = self._pool(page_kind) - # TODO：将多请求 Conv/SSM 保存融合为按 req/slot 索引搬运的 kernel； - # GPU 直访 CPU 页池需保证页池为可访问的 pinned/registered 内存，完成通知仍由调用方负责。 # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): - pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) - pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) + conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" - # TODO：融合多请求 Conv/SSM 恢复及目标游标清零，直接按 slot/req 索引写入基准位置； - # CPU 页池须满足 GPU 可访问条件，并保持未参与请求和其他候选区不变。 pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): - self.conv_state[:, req_idx, :, : self.config.conv_width].copy_( - pool.conv_state[slot_id], non_blocking=True, - ) - self.ssm_state[:, req_idx * (self.mtp_step + 1)].copy_(pool.ssm_state[slot_id], non_blocking=True) - if self.mtp_state_index is not None: + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state( @@ -192,23 +194,21 @@ def write_page_state_to_cpu( assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" if cpu_page_readies is not None: assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" - copies = [] - # TODO：批量融合 Conv/SSM 多页写出及 TP 切片，减少逐槽位、逐字段 copy_； - # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + page_fields, page_indices = [], [] for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): - # ready 页的完整联合数据保持不变，也无需构造快照复制视图。 + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 if cpu_page_readies is not None and cpu_page_readies[index]: continue - fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies( - fields=fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_index=page, - write=True, - )) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) def read_page_state_from_cpu( self, @@ -225,20 +225,15 @@ def read_page_state_from_cpu( assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - copies = [] - # TODO：批量融合 Conv/SSM 多页读入及 TP 切片，减少逐槽位、逐字段 copy_； - # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 - for page, slot in zip(cpu_page_indices, big_page_slot_ids): - fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies( - fields=fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_index=page, - write=False, - )) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) # ======================================================================== # 四、DP 迁移：大页快照与基准运行态 @@ -260,13 +255,13 @@ def copy_big_page_state_from_other_dp( dst_pool = self.big_pages_pool assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # TODO：改用 Triton kernel 按页索引批量复制 Conv/SSM，在 CUDA stream 上异步执行。 - # 两端为 pinned CPU Tensor 时，当前 copy_ 仍执行 CPU 复制，non_blocking=True 不生效； - # kernel 需直接访问本进程已注册且 GPU 可寻址的两端页池，调用方须建立源页就绪依赖， - # 并在复制完成前保留源、目标页的引用；当前 CPU 路径须先确认源页 GPU 写入已完成。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] for src_page, dst_page in zip(src_page_ids, dst_page_ids): - dst_pool.conv_state[dst_page].copy_(src_pool.conv_state[src_page], non_blocking=True) - dst_pool.ssm_state[dst_page].copy_(src_pool.ssm_state[src_page], non_blocking=True) + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def copy_run_time_state_from_other_dp( self, @@ -282,16 +277,21 @@ def copy_run_time_state_from_other_dp( ) # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # TODO：融合多请求 Conv/SSM 的 DP 复制与目标游标清零；跨设备源存储需可通过 peer/IPC 访问。 + # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 + conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.conv_state[:, dst_req, :, : self.config.conv_width].copy_( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], non_blocking=True, - ) - self.ssm_state[:, dst_req * (self.mtp_step + 1)].copy_( + conv_copies.append(( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + )) + ssm_copies.append(( remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - non_blocking=True, - ) - if self.mtp_state_index is not None: + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + )) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) # ======================================================================== @@ -313,16 +313,14 @@ def write_to_pd_buffer( assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) - # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 - copies = bounded_att_buffer_copies( - fields=fields, + for copies in bounded_att_buffer_copies( + page_fields=[fields], page_buffers=pd_page_buffers, layout=layout, - page_index=pd_page_index, + page_indices=[pd_page_index], write=True, - ) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + ): + copy_bounded_att_tensors(copies=copies) def read_from_pd_buffer( self, @@ -335,16 +333,14 @@ def read_from_pd_buffer( assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) - # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 - copies = bounded_att_buffer_copies( - fields=fields, + for copies in bounded_att_buffer_copies( + page_fields=[fields], page_buffers=pd_page_buffers, layout=layout, - page_index=pd_page_index, + page_indices=[pd_page_index], write=False, - ) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + ): + copy_bounded_att_tensors(copies=copies) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py index d92c8bf7d0..f5628878a9 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -13,6 +13,7 @@ from .conv_ssm_types import ConvSSMStateConfig from .kda_pages_pool import KDAPagesPool from .transfer_utils import bounded_att_buffer_copies +from .copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: from lightllm.common.att_resources.types import PageKind @@ -129,27 +130,28 @@ def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, sl assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" pool = self._pool(page_kind) - # TODO：将多请求 Conv/SSM 保存融合为按 req/slot 索引搬运的 kernel； - # GPU 直访 CPU 页池需保证页池为可访问的 pinned/registered 内存，完成通知仍由调用方负责。 # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): - pool.conv_state[slot_id].copy_(self.conv_state[:, req_idx, :, : self.config.conv_width], non_blocking=True) - pool.ssm_state[slot_id].copy_(self.ssm_state[:, req_idx * (self.mtp_step + 1)], non_blocking=True) + conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert len(req_indices) == len(slot_ids), "request and page counts differ" - # TODO：融合多请求 Conv/SSM 恢复及目标游标清零，直接按 slot/req 索引写入基准位置； - # CPU 页池须满足 GPU 可访问条件，并保持未参与请求和其他候选区不变。 pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): - self.conv_state[:, req_idx, :, : self.config.conv_width].copy_( - pool.conv_state[slot_id], non_blocking=True, - ) - self.ssm_state[:, req_idx * (self.mtp_step + 1)].copy_(pool.ssm_state[slot_id], non_blocking=True) - if self.mtp_state_index is not None: + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) def commit_bounded_att_state( @@ -192,23 +194,21 @@ def write_page_state_to_cpu( assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" if cpu_page_readies is not None: assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" - copies = [] - # TODO：批量融合 Conv/SSM 多页写出及 TP 切片，减少逐槽位、逐字段 copy_； - # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 + page_fields, page_indices = [], [] for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): - # ready 页的完整联合数据保持不变，也无需构造快照复制视图。 + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 if cpu_page_readies is not None and cpu_page_readies[index]: continue - fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies( - fields=fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_index=page, - write=True, - )) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) def read_page_state_from_cpu( self, @@ -225,20 +225,15 @@ def read_page_state_from_cpu( assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - copies = [] - # TODO：批量融合 Conv/SSM 多页读入及 TP 切片，减少逐槽位、逐字段 copy_； - # 页池与缓存页均在 CPU 时应使用 CPU 批量复制，仅有 GPU 端且主机内存可访问时使用传输 kernel。 - for page, slot in zip(cpu_page_indices, big_page_slot_ids): - fields = self._fields(pool.conv_state[slot], pool.ssm_state[slot]) - copies.extend(bounded_att_buffer_copies( - fields=fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_index=page, - write=False, - )) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) # ======================================================================== # 四、DP 迁移：大页快照与基准运行态 @@ -260,13 +255,13 @@ def copy_big_page_state_from_other_dp( dst_pool = self.big_pages_pool assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # TODO：改用 Triton kernel 按页索引批量复制 Conv/SSM，在 CUDA stream 上异步执行。 - # 两端为 pinned CPU Tensor 时，当前 copy_ 仍执行 CPU 复制，non_blocking=True 不生效； - # kernel 需直接访问本进程已注册且 GPU 可寻址的两端页池，调用方须建立源页就绪依赖， - # 并在复制完成前保留源、目标页的引用；当前 CPU 路径须先确认源页 GPU 写入已完成。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] for src_page, dst_page in zip(src_page_ids, dst_page_ids): - dst_pool.conv_state[dst_page].copy_(src_pool.conv_state[src_page], non_blocking=True) - dst_pool.ssm_state[dst_page].copy_(src_pool.ssm_state[src_page], non_blocking=True) + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) def copy_run_time_state_from_other_dp( self, @@ -282,16 +277,21 @@ def copy_run_time_state_from_other_dp( ) # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # TODO：融合多请求 Conv/SSM 的 DP 复制与目标游标清零；跨设备源存储需可通过 peer/IPC 访问。 + # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 + conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - self.conv_state[:, dst_req, :, : self.config.conv_width].copy_( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], non_blocking=True, - ) - self.ssm_state[:, dst_req * (self.mtp_step + 1)].copy_( + conv_copies.append(( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + )) + ssm_copies.append(( remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - non_blocking=True, - ) - if self.mtp_state_index is not None: + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + )) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) # ======================================================================== @@ -313,16 +313,14 @@ def write_to_pd_buffer( assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) - # TODO：将 Conv/SSM 的 PD 打包融合到一次 kernel，直接写指定页的字段偏移并保留 TP 写入规则。 - copies = bounded_att_buffer_copies( - fields=fields, + for copies in bounded_att_buffer_copies( + page_fields=[fields], page_buffers=pd_page_buffers, layout=layout, - page_index=pd_page_index, + page_indices=[pd_page_index], write=True, - ) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + ): + copy_bounded_att_tensors(copies=copies) def read_from_pd_buffer( self, @@ -335,16 +333,14 @@ def read_from_pd_buffer( assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) - # TODO：融合 PD 页中 Conv/SSM 的恢复与目标游标清零，按布局切片直接写入基准运行态。 - copies = bounded_att_buffer_copies( - fields=fields, + for copies in bounded_att_buffer_copies( + page_fields=[fields], page_buffers=pd_page_buffers, layout=layout, - page_index=pd_page_index, + page_indices=[pd_page_index], write=False, - ) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) + ): + copy_bounded_att_tensors(copies=copies) if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) diff --git a/lightllm/common/att_resources/bounded_att/transfer_utils.py b/lightllm/common/att_resources/bounded_att/transfer_utils.py index 0736008dcb..7f9cb8bb9c 100644 --- a/lightllm/common/att_resources/bounded_att/transfer_utils.py +++ b/lightllm/common/att_resources/bounded_att/transfer_utils.py @@ -12,41 +12,35 @@ # ======================================================================== def bounded_att_buffer_copies( - fields: Mapping[str, TensorField], + page_fields: list[Mapping[str, TensorField]], page_buffers: torch.Tensor, layout: StateTransferLayout, - page_index: int, + page_indices: list[int], write: bool, -) -> list[tuple[torch.Tensor, torch.Tensor]]: - """生成一份 bounded 状态与外部页之间的复制计划，不处理 token 数量。 +) -> list[list[tuple[torch.Tensor, torch.Tensor]]]: + """按传输区域收集所有页面的复制视图，每个返回批次只包含一种 shape。 - fields 由 State 选取运行态或大小页快照，例如 KDA/GDN 的 Conv、SSM。 - 每份快照完整复制；Conv 的 Q/K/V 区间可以分别映射到同一全局字段。 - page_buffers 至少为二维，第一维统一表示页；不接受没有页维度的单页 Tensor。 - page_index 必须显式传入；CPU/PD 调用方均传入页集合和实际物理页索引。 - 返回 (源视图, 目标视图) 列表，由调用方遍历执行 dst.copy_(src, non_blocking=True)。 + page_fields 与 page_indices 一一对应，调用方已排除 ready 页。 + Conv 的 Q/K/V 切片与 SSM 分别成批，不将不同字段或不同 TP 切片混入同一次拷贝。 + page_buffers 的第一维为物理页；这里只创建视图，由调用方逐批提交 Triton 算子。 """ - # CPU/PD 统一沿第一维选页，再按字节解释页内存储，不复制数据。 assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" - page_buffer = page_buffers[page_index] - page_bytes = page_buffer.view(torch.uint8).view(-1) - copies = [] + batches = [] for region in layout.regions: - # 阶段一：定位快照字段和外部页中的全局状态区域。 - assert region.field in fields, f"unknown transfer field {region.field}" - field = fields[region.field] - # region 起止位置相对整页，按字段 dtype 和 shape 直接构造外部视图。 - external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) if write and not region.write_enabled: continue - - # 阶段二：按 channel/head 等 TP 切片生成视图，Conv 窗口及 SSM 状态完整保留。 - local_view, external_view = _copy_pair(field, external, region) - # 比较切片后的实际复制范围，避免 copy_ 隐式广播掩盖布局错误。 - assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" - # 写出只由选定的 TP 副本执行；读回覆盖每个副本自己的分片。 - copies.append((local_view, external_view) if write else (external_view, local_view)) - return copies + copies = [] + for fields, page_index in zip(page_fields, page_indices): + # 阶段一：按字节区间定位当前页的字段，构造外部存储视图。 + field = fields[region.field] + page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) + external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + # 阶段二：同一区域使用相同的 TP 切片，所有页面生成的视图 shape 一致。 + local_view, external_view = _copy_pair(field, external, region) + assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" + copies.append((local_view, external_view) if write else (external_view, local_view)) + batches.append(copies) + return batches # ======================================================================== diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 823403240b..043973a622 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -158,16 +158,16 @@ def prepare_att_resources(self) -> None: # 当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；各工厂自行选择参与层及共享分组。 logical_layer_ids = tuple(layer.layer_num_ for layer in self.layers_infer) - full_states = self._create_full_att_states( + full_states = self.create_full_att_states( model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, ) - bounded_states = self._create_bounded_att_states( + bounded_states = self.create_bounded_att_states( model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, device="cuda", ) - tail_states = self._create_att_tail_states( + tail_states = self.create_att_tail_states( model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, device="cuda", @@ -259,7 +259,7 @@ def load_weights(self, weight_dict: dict): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: @@ -271,7 +271,7 @@ def _create_full_att_states( return [] @classmethod - def _create_bounded_att_states( + def create_bounded_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", @@ -283,7 +283,7 @@ def _create_bounded_att_states( return [] @classmethod - def _create_att_tail_states( + def create_att_tail_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index fb9da7b995..121174d513 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -32,7 +32,7 @@ def __init__(self, kvargs): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 7f90e2cddb..76c760dc21 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -64,7 +64,7 @@ def _verify_params(self): return super()._verify_params() @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index beaf390d9a..2fba860e64 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -20,7 +20,7 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): transformer_layer_infer_class = Deepseek3_2TransformerLayerInfer @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index caa0563f0e..5114544b5f 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -144,7 +144,7 @@ def _init_custom(self): self._init_to_get_rotary() @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2c73be470c..5a3e427360 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -22,7 +22,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): time. The KV cache layout is uniform (sliding shape: num_kv_heads=16, head_dim=256); full-attention layers pack their (4, 512) tensor into the first 8 heads of the 16-head slot at cache-write time, then reshape on - read. See Gemma4TpPartModel._create_full_att_states for context. + read. See Gemma4TpPartModel.create_full_att_states for context. """ def __init__(self, layer_num, network_config): diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 81f22d6218..322ccf112b 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -87,7 +87,7 @@ def _verify_params(self): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index 4ba92d2ed9..94c315507e 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -42,7 +42,7 @@ def _verify_params(self): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index c652d0c5ef..658334f123 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -34,7 +34,7 @@ class Glm5NextTpPartModel(TpPartBaseModel): transformer_layer_infer_class = Glm5NextTransformerLayerInfer @classmethod - def _create_bounded_att_states( + def create_bounded_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", @@ -107,7 +107,7 @@ def _init_some_value(self): self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", @@ -142,7 +142,7 @@ def _create_full_att_states( return [state] @classmethod - def _create_att_tail_states( + def create_att_tail_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 742c636e64..9a94863a23 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -72,7 +72,7 @@ def _init_some_value(self): self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", @@ -105,7 +105,7 @@ def _create_full_att_states( return [state] @classmethod - def _create_att_tail_states( + def create_att_tail_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 0d059c5516..be2c8448b2 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -40,7 +40,7 @@ def __init__(self, kvargs): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 6118971c1c..0de2177240 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -34,7 +34,7 @@ def __init__(self, kvargs): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 54373ba665..51a5b0b1f2 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -42,7 +42,7 @@ def _init_config(self): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index db1d26b740..79415952fe 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -40,7 +40,7 @@ def __init__(self, kvargs) -> None: super().__init__(kvargs) @classmethod - def _create_bounded_att_states( + def create_bounded_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", @@ -97,7 +97,7 @@ def _init_config(self): self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index b97af1d1bd..ce97049e02 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -30,7 +30,7 @@ def __init__(self, kvargs): super().__init__(kvargs) @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 34ddb866aa..24a21cf1e0 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -33,7 +33,7 @@ def __init__(self, kvargs): return @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index 7ac0dfeead..ecc24a1a9a 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -37,10 +37,10 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: config = model_class.normalize_model_config(config) count = config["num_hidden_layers"] logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) - full_states = model_class._create_full_att_states( + full_states = model_class.create_full_att_states( model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, ) - bounded_states = model_class._create_bounded_att_states( + bounded_states = model_class.create_bounded_att_states( model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, diff --git a/unit_tests/common/att_resources/test_bounded_copy_kernel.py b/unit_tests/common/att_resources/test_bounded_copy_kernel.py new file mode 100644 index 0000000000..d64b42ce15 --- /dev/null +++ b/unit_tests/common/att_resources/test_bounded_copy_kernel.py @@ -0,0 +1,264 @@ +"""聚合复制必须保持位模式、stride、流顺序和未参与区域。""" + +import pytest +import torch +import triton.language as tl + +from lightllm.common.att_resources.bounded_att.copy_kernel import copy_bounded_att_tensors +from lightllm.common.att_resources.bounded_att.utils import register_page_pool_tensor + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("direction", ["cuda_to_host", "host_to_cuda", "host_to_host", "cuda_to_cuda"]) +@pytest.mark.parametrize("grid_num", [1, 5, 16]) +@pytest.mark.parametrize("dtype", [torch.int8, torch.bfloat16, torch.float32, torch.float16, torch.int64]) +def test_aggregate_same_dtype_and_strides(direction: str, grid_num: int, dtype: torch.dtype) -> None: + src_device = "cuda" if direction.startswith("cuda") else "cpu" + dst_device = "cuda" if direction.endswith("cuda") else "cpu" + sources, targets, copies, expected_targets = [], [], [], [] + # Conv 窗口裁剪和 SSM 轴交换均保留间隙；使用非 2 次幂的边界覆盖最后一个 block。 + # 同批视图保持相同 shape，验证 grid 跳步和跨视图 block 分配。 + for _ in range(41): + shape = (3, 5, 7, 131) + source = torch.empty( + size=shape, + dtype=dtype, + device=src_device, + pin_memory=src_device == "cpu", + ) + source.copy_(torch.arange(source.numel()).reshape(shape) % 127) + target_shape = (shape[1], shape[0], shape[2] + 2, shape[3] + 4) + target = torch.full( + size=target_shape, + fill_value=-17, + dtype=dtype, + device=dst_device, + pin_memory=dst_device == "cpu", + ) + source_view = source[..., :shape[3] - 1] + target_view = target.permute(1, 0, 2, 3)[:, :, 1:shape[2] + 1, 2:shape[3] + 1] + expected = target.cpu().clone() + expected.permute(1, 0, 2, 3)[:, :, 1:shape[2] + 1, 2:shape[3] + 1].copy_(source_view.cpu()) + sources.append(source) + targets.append(target) + copies.append((source_view, target_view)) + expected_targets.append(expected) + copy_bounded_att_tensors(copies=copies, grid_num=grid_num) + torch.cuda.current_stream().synchronize() + for target, expected in zip(targets, expected_targets): + assert torch.equal(target.cpu(), expected) + + +def test_registered_shared_pages_and_stream_dependency() -> None: + source = torch.empty(size=(4099,), dtype=torch.float32).share_memory_() + target = torch.empty(size=(4099,), dtype=torch.float32).share_memory_() + register_page_pool_tensor(tensor=source) + register_page_pool_tensor(tensor=target) + producer = torch.cuda.Stream() + consumer = torch.cuda.Stream() + values = torch.arange(source.numel(), dtype=source.dtype, device="cuda") + with torch.cuda.stream(producer): + copy_bounded_att_tensors(copies=[(values, source)]) + with torch.cuda.stream(consumer): + consumer.wait_stream(producer) + # 这里不经过主机 synchronize；CPU→CPU 传输必须在 GPU stream 上等待上游。 + copy_bounded_att_tensors(copies=[(source, target)]) + restored = torch.empty_like(values) + copy_bounded_att_tensors(copies=[(target, restored)]) + consumer.synchronize() + assert torch.equal(restored, values) + assert torch.equal(target, values.cpu()) + + +@pytest.mark.parametrize("float_dtype,int_dtype,patterns", [ + (torch.float16, torch.int16, [0x7E01, -32768, 0x7C00, 1]), + (torch.bfloat16, torch.int16, [0x7FC1, -32768, 0x7F80, 1]), + (torch.float32, torch.int32, [0x7FC00001, -2147483648, 0x7F800000, 1]), + (torch.float64, torch.int64, [0x7FF8000000000001, -9223372036854775808, 0x7FF0000000000000, 1]), +]) +def test_exact_float_bit_patterns(float_dtype: torch.dtype, int_dtype: torch.dtype, patterns: list[int]) -> None: + # 包含 NaN payload、负零和无穷，等宽整数复制不能触发精度转换或规范化。 + bits = torch.tensor(data=patterns, dtype=int_dtype, device="cuda") + target = torch.empty(size=(4,), dtype=float_dtype, pin_memory=True) + copy_bounded_att_tensors(copies=[(bits.view(float_dtype), target)]) + torch.cuda.synchronize() + assert torch.equal(target.view(int_dtype), bits.cpu()) + + +def test_one_kernel_for_same_dimension_batch(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.att_resources.bounded_att import copy_kernel + + actual_kernel = copy_kernel._copy_bounded_att_1d + launches = [] + + class CountLaunches: + def __getitem__(self, grid: tuple[int]): + launches.append(grid) + return actual_kernel[grid] + + monkeypatch.setattr(copy_kernel, "_copy_bounded_att_1d", CountLaunches()) + source = torch.ones(size=(32, 67), device="cuda") + target = torch.zeros_like(source) + copy_bounded_att_tensors(copies=[(source[i], target[i]) for i in range(32)]) + assert launches == [(24,)] + assert torch.equal(source, target) + copy_bounded_att_tensors(copies=[]) + with pytest.raises(AssertionError, match="copy tensors must be nonempty"): + copy_bounded_att_tensors(copies=[(source[:0], target[:0])]) + assert len(launches) == 1 + + +@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5, 6]) +def test_up_to_six_dimensions(ndim: int) -> None: + # 非整 block 的大小，覆盖每个维度算子的跳步循环与尾部 mask。 + shape = (2,) * (ndim - 1) + (2051,) + source = torch.arange(2 ** (ndim - 1) * 2051, dtype=torch.float32, device="cuda") + source = source.reshape(shape) + if ndim >= 2: + source = source.transpose(0, ndim - 1) + target = torch.full( + size=source.shape, + fill_value=-1, + dtype=source.dtype, + pin_memory=True, + ) + copy_bounded_att_tensors(copies=[(source, target)]) + torch.cuda.synchronize() + assert torch.equal(source.cpu(), target) + + +@pytest.mark.parametrize("ndim", range(1, 7)) +@pytest.mark.parametrize("num_stages", [1, 2, 4, 8]) +def test_same_shape_batch_uses_one_kernel( + ndim: int, + num_stages: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from lightllm.common.att_resources.bounded_att import copy_kernel + + launches = [] + kernel_name = f"_copy_bounded_att_{ndim}d" + actual_kernel = getattr(copy_kernel, kernel_name) + + class CountLaunches: + def __getitem__(self, grid: tuple[int]): + launches.append(grid) + return actual_kernel[grid] + + monkeypatch.setattr(copy_kernel, kernel_name, CountLaunches()) + copies, targets, expected_targets = [], [], [] + for _ in range(3): + shape = (2,) * (ndim - 1) + (517,) + storage_shape = (4,) * (ndim - 1) + (1034,) + source = torch.arange(torch.Size(storage_shape).numel(), dtype=torch.float32, device="cuda").reshape(storage_shape) + # 每个轴都保留间隙，保证降维后仍需调用对应的 1～6 维算子。 + source_view = source[(slice(None, None, 2),) * ndim] + target = torch.full(size=(*shape[:-1], 1034), fill_value=-17, dtype=source.dtype, pin_memory=True) + target_view = target[..., ::2] + copies.append((source_view, target_view)) + expected = target.clone() + expected[..., ::2].copy_(source_view.cpu()) + targets.append(target) + expected_targets.append(expected) + + copy_bounded_att_tensors(copies=copies, grid_num=5, num_stages=num_stages) + torch.cuda.synchronize() + assert launches == [(5,)] + for target, expected in zip(targets, expected_targets): + assert torch.equal(target, expected) + + +@pytest.mark.parametrize("layout", ["ssm", "conv_mtp", "dense", "width4", "middle_unaligned", "prefix", "singletons", "wide_singleton", "wide_singleton_prefix", "all_singletons"]) +def test_transfer_views_preserve_gaps_and_alignment(layout: str, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.att_resources.bounded_att import copy_kernel + + launches = [] + + class CaptureLaunch: + def __init__(self, kernel, ndim: int) -> None: + self.kernel = kernel + self.ndim = ndim + + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((self.ndim, kwargs["SHAPE"], kwargs["DTYPE"])) + return self.kernel[grid](**kwargs) + return launch + + for ndim in range(1, 7): + name = f"_copy_bounded_att_{ndim}d" + monkeypatch.setattr(copy_kernel, name, CaptureLaunch(getattr(copy_kernel, name), ndim)) + + if layout == "ssm": + source = torch.arange(2 * 4 * 3 * 128, dtype=torch.float32, device="cuda").view(2, 4, 3, 128) + target = torch.full(size=(4, 2, 3, 128), fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[:, index], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[:, index].cpu()) + expected_launch = (2, (2, 192), tl.uint64) + elif layout == "conv_mtp": + source = torch.arange(2 * 4 * 16 * 7, dtype=torch.float32, device="cuda").to(torch.bfloat16).view(2, 4, 16, 7) + target = torch.full(size=(4, 2, 16, 4), fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[:, index, :, :4], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[:, index, :, :4].cpu()) + expected_launch = (3, (2, 16, 4), tl.uint16) + elif layout in ("prefix", "singletons", "wide_singleton", "wide_singleton_prefix", "all_singletons"): + if layout == "prefix": + shape = (2, 3, 4, 7) + target_shape = (2, 3, 4, 9) + width = 5 + expected_launch = (4, (2, 3, 4, 5), tl.uint32) + elif layout == "singletons": + shape = (1, 2, 1, 3, 1, 7) + target_shape = (1, 2, 1, 3, 1, 9) + width = 5 + expected_launch = (6, (1, 2, 1, 3, 1, 5), tl.uint32) + elif layout == "wide_singleton_prefix": + # dtype 扩宽只改变末轴长度，前后的单元素轴均保持不变。 + shape = (1, 3, 4) + target_shape = (1, 3, 6) + width = 2 + expected_launch = (3, (1, 3, 1), tl.uint64) + elif layout == "wide_singleton": + shape = (3, 4) + target_shape = (3, 6) + width = 2 + expected_launch = (2, (3, 1), tl.uint64) + else: + shape = target_shape = (1, 1, 1) + width = 1 + expected_launch = (1, (1,), tl.uint32) + source = torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").view(shape) + target = torch.full(size=target_shape, fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[..., :width], target[..., :width])] + expected = target.clone() + expected[..., :width].copy_(source[..., :width].cpu()) + elif layout == "middle_unaligned": + source = torch.arange(64, dtype=torch.bfloat16, device="cuda") + target = torch.full(size=(64,), fill_value=-1, dtype=source.dtype, pin_memory=True) + expected = target.clone() + copies = [] + # 首尾地址都满足 8 字节对齐，中间请求偏移不对齐；整批必须保留 2 字节搬运。 + for src_offset, dst_offset in ((0, 0), (9, 16), (24, 32)): + copies.append((source[src_offset:src_offset + 8], target[dst_offset:dst_offset + 8])) + expected[dst_offset:dst_offset + 8].copy_(source[src_offset:src_offset + 8].cpu()) + expected_launch = (1, (8,), tl.uint16) + else: + shape = (4, 2, 3, 128) if layout == "dense" else (4, 6) + source = torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").to(torch.bfloat16).view(shape) + target = torch.full(size=shape, fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[index], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[index].cpu()) + expected_launch = (1, (192,), tl.uint64) if layout == "dense" else (1, (3,), tl.uint32) + + copy_bounded_att_tensors(copies=copies) + torch.cuda.synchronize() + assert torch.equal(target, expected) + assert launches == [expected_launch] diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 3b43e74763..4c4290449f 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -22,6 +22,8 @@ from lightllm.common.att_resources.types import PageKind +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for aggregated copies") + def make_state( cls: type[KDAState] | type[GDNState] = KDAState, device: str | torch.device | int = "cpu", @@ -118,7 +120,7 @@ def page_layout(state: KDAState | GDNState): @pytest.mark.parametrize("cls,pool_cls", [(KDAState, KDAPagesPool), (GDNState, GDNPagesPool)]) -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cuda"]) @pytest.mark.parametrize("page_kind", ["big", "small"]) def test_prefill_state_is_saved_and_restored( cls: type[KDAState] | type[GDNState], @@ -188,9 +190,9 @@ def test_gdn_dimensions_and_page_restore(): config=config, request_capacity=3, mtp_step=2, - device="cpu", + device="cuda", ) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2, pin_memory=False) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2, pin_memory=True) conv = torch.arange(40, dtype=config.conv_dtype).reshape(2, 10, 2) ssm = torch.arange(24, dtype=config.ssm_dtype).reshape( shape=(2, 2, 2, 3), @@ -202,15 +204,15 @@ def test_gdn_dimensions_and_page_restore(): state.reset_bounded_att_state([0]) for req, page_kind in enumerate(("big", "small"), start=1): state.restore_bounded_att_pages([req], page_kind=page_kind, slot_ids=[0]) - assert torch.equal(state.conv_state[:, req, :, :2], conv) - assert torch.equal(state.ssm_state[:, req * 3], ssm) + assert torch.equal(state.conv_state[:, req, :, :2].cpu(), conv) + assert torch.equal(state.ssm_state[:, req * 3].cpu(), ssm) with pytest.raises(AssertionError, match="equal Q/K/V"): KDAPagesPool(2, config, pin_memory=False) @pytest.mark.parametrize("cls", [KDAState, GDNState]) def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): - state = make_state(cls) + state = make_state(cls, device="cuda") state.conv_state[:, 0].fill_(7) state.ssm_state[:, 0:3].fill_(11) state.save_bounded_att_pages([0], page_kind="small", slot_ids=[0]) @@ -229,12 +231,12 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): @pytest.mark.parametrize("cls", [KDAState, GDNState]) def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): - source = make_state(cls) + source = make_state(cls, device="cuda") source.conv_state[:, 0].fill_(13) source.ssm_state[:, 0].fill_(29) source.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) layout = page_layout(source) - buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8) + buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) source.write_page_state_to_cpu( page_kind="big", slot_ids=[0], @@ -242,8 +244,9 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): cpu_page_indices=[2], layout=layout, ) + torch.cuda.synchronize() assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() - target = make_state(cls) + target = make_state(cls, device="cuda") target.read_page_state_from_cpu( cpu_page_buffers=buffer, cpu_page_indices=[2], @@ -267,10 +270,11 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): cpu_page_indices=[0], layout=layout, ) + torch.cuda.synchronize() assert torch.equal(buffer[0], buffer[2]) -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cuda"]) @pytest.mark.parametrize("cls", [KDAState, GDNState]) def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( cls: type[KDAState] | type[GDNState], @@ -313,7 +317,7 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( @pytest.mark.parametrize("cls", [KDAState, GDNState]) -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("device", ["cuda"]) @pytest.mark.parametrize("src_mtp_step,dst_mtp_step", [(0, 2), (2, 1)]) def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp_step, dst_mtp_step): if device == "cuda" and not torch.cuda.is_available(): @@ -355,8 +359,8 @@ def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp @pytest.mark.parametrize("cls", [KDAState, GDNState]) def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): - source = make_state(cls) - target = make_state(cls) + source = make_state(cls, device="cuda") + target = make_state(cls, device="cuda") for page in range(3): source.big_pages_pool.conv_state[page].fill_(10 + page) source.big_pages_pool.ssm_state[page].fill_(20 + page) @@ -374,6 +378,7 @@ def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): src_page_ids=[2, 0], dst_page_ids=[0, 1], remote_dp_bounded_state=source, ) + torch.cuda.synchronize() for src_page, dst_page in ((2, 0), (0, 1)): assert torch.equal(target.big_pages_pool.conv_state[dst_page], src_conv[src_page]) assert torch.equal(target.big_pages_pool.ssm_state[dst_page], src_ssm[src_page]) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index 4aea99aa7b..a8b2d31039 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -199,6 +199,7 @@ def test_shared_allocator_accounts_for_distinct_state_stores(): @pytest.mark.parametrize("state_cls", [KDAState, GDNState]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( monkeypatch: pytest.MonkeyPatch, state_cls: type[KDAState] | type[GDNState], @@ -223,13 +224,13 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( ), request_capacity=3, mtp_step=3, - device="cpu", + device="cuda", ) reference = state_cls( config=config, request_capacity=3, mtp_step=3, - device="cpu", + device="cuda", ) conv = reference.conv_state[:, 0, :, : config.conv_width] ssm = reference.ssm_state[:, 0] @@ -244,7 +245,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( reference.conv_state.fill_(2) reference.ssm_state.fill_(3) tail.tail_buffer.fill_(4) - pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8) + pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "tail"} manager.write_to_pd_buffer( @@ -446,6 +447,7 @@ def _get_transfer_regions(self, page_tokens: int): @pytest.mark.parametrize("full_page_tokens", [None, 8, 10, 16]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_pd_runtime_capacity_includes_field_padding( monkeypatch: pytest.MonkeyPatch, pd_args: EasyDict, @@ -479,7 +481,7 @@ def get_pd_transfer_regions(self): conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ), - request_capacity=1, device="cpu", + request_capacity=1, device="cuda", ) manager.set_att_state("main", state) manager.finish_state_creation() @@ -498,7 +500,7 @@ def get_pd_transfer_regions(self): assert manager.get_pd_page_bytes() == expected assert manager.get_pd_bounded_att_or_tail_transfer_layouts() is layouts # 实际传输页可大于运行态布局，选中第二页后仍从该页起点搬运。 - pages = torch.full((3, expected), 91, dtype=torch.uint8) + pages = torch.full((3, expected), 91, dtype=torch.uint8, device="cuda") state.conv_state.fill_(3) state.ssm_state.fill_(5) manager.write_to_pd_buffer( @@ -783,6 +785,7 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea @pytest.mark.parametrize("page_tokens", [1, 9]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: EasyDict, page_tokens: int): cpu_args.cpu_cache_token_page_size = page_tokens manager = AttResourceManager(request_capacity=1) @@ -803,9 +806,9 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas config=config, request_capacity=1, mtp_step=2, - device="cpu", + device="cuda", ) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=False) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=True) manager.set_att_state(role, state) manager.finish_state_creation() states = manager.get_bounded_attention_states() @@ -819,7 +822,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas state.ssm_state.fill_(i + 7) source_slots = manager.bounded_att_big_page_allocator.alloc(1) manager.save_bounded_att_pages([0], page_kind="big", slot_ids=source_slots) - pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8) + pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) manager.write_req_to_cpu( token_indices=torch.tensor([], dtype=torch.long), cpu_page_buffers=pages, @@ -827,6 +830,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas layouts=layouts, big_page_slot_ids=source_slots, ) + torch.cuda.synchronize() assert (pages[0] == 87).all() assert (pages[1, 12:16] == 87).all() manager.reset_bounded_att_state([0]) diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 0570effb52..543ca01046 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -80,7 +80,7 @@ def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "E num_kv_shared_layers=2, layer_types=["sliding_attention", "full_attention"] * 3, ) - (state,) = Gemma4TpPartModel._create_full_att_states( + (state,) = Gemma4TpPartModel.create_full_att_states( model_config=config, args=pd_args, model_logical_layer_ids=(10, 11, 12, 13, 14, 15), ) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 5f5522e328..35c266c7a1 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -33,7 +33,7 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: ), device="meta", ) monkeypatch.setattr( - LlamaTpPartModel, "_create_full_att_states", classmethod(lambda cls, **kwargs: [state]), + LlamaTpPartModel, "create_full_att_states", classmethod(lambda cls, **kwargs: [state]), ) model.prepare_att_resources() diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 0c8d653627..72a2e794c9 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -72,8 +72,9 @@ def init_source(manager: AttResourceManager): return slots +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.MonkeyPatch): - manager = composition() + manager = composition(device="cuda") slots = init_source(manager) manager.restore_bounded_att_pages([2], page_kind="big", slot_ids=[slots[0]]) for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): @@ -106,6 +107,7 @@ def test_disabled_domains(): empty_manager.copy_full_att_tokens(torch.tensor([0], dtype=torch.long), torch.tensor([1], dtype=torch.long)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_reset_tail_clears_all_tail_states_and_preserves_linear(): first = make_tail() second = KPoolAttentionTailState( @@ -123,7 +125,7 @@ def test_reset_tail_clears_all_tail_states_and_preserves_linear(): tail_manager.reset_attention_tail_state([0, 1]) assert not first.tail_buffer[:, :2].any() and not second.tail_buffer[:, :2].any() assert (first.tail_buffer[:, 2] == 5).all() and (second.tail_buffer[:, 2] == 7).all() - manager = composition() + manager = composition(device="cuda") init_source(manager) # Preserve accepted-state cursors as well as both algorithms' runtime and # cached pages. A full reset here would silently corrupt a prefix restore. @@ -346,8 +348,9 @@ def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.d assert (source.get_attention_tail_states()[0].tail_buffer[:, 1] == 9).all() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): - source, target = composition(), composition() + source, target = composition(device="cuda"), composition(device="cuda") init_source(source) tails = target.get_attention_tail_states() for state in tails: @@ -367,9 +370,9 @@ def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): @pytest.mark.parametrize("copy_tokens,copy_pages", [(True, True), (True, False), (False, True)]) def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: bool, copy_pages: bool): - if copy_tokens and not torch.cuda.is_available(): + if not torch.cuda.is_available(): pytest.skip("CUDA unavailable") - device = "cuda" if copy_tokens else "cpu" + device = "cuda" source, target = composition(device), composition(device) init_source(source) for state in target.get_full_attention_states(): diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index c91fe76fd0..f96e96fca6 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -102,7 +102,7 @@ def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encod model.att_resource_manager = manager states = [] if not (model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att")): - states = model._create_full_att_states( + states = model.create_full_att_states( model_config=model.config, args=model.args, model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), diff --git a/unit_tests/common/att_resources/test_model_config_normalization.py b/unit_tests/common/att_resources/test_model_config_normalization.py index 4519763c26..f7dfa70a73 100644 --- a/unit_tests/common/att_resources/test_model_config_normalization.py +++ b/unit_tests/common/att_resources/test_model_config_normalization.py @@ -52,7 +52,7 @@ def test_dflash_runtime_and_cpu_state_use_overridden_config( assert raw_config == original args = StartArgs(data_type="bfloat16", llm_kv_type="None", page_size=4) - (state,) = model_class._create_full_att_states( + (state,) = model_class.create_full_att_states( model_config=cpu_config, args=args, model_logical_layer_ids=tuple(range(cpu_config["num_hidden_layers"])), diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index 241b6c3291..d8c57ed244 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -209,7 +209,7 @@ def test_cpu_layout_uses_all_model_factory_states( class MixedStorageModel(LlamaTpPartModel): @classmethod - def _create_full_att_states( + def create_full_att_states( cls, model_config: dict[str, object], args: StartArgs | EasyDict, @@ -217,7 +217,7 @@ def _create_full_att_states( tp_rank: int = 0, tp_world_size: int = 1, ) -> list[FullAttentionState]: - (template,) = super()._create_full_att_states( + (template,) = super().create_full_att_states( model_config=model_config, args=args, model_logical_layer_ids=model_logical_layer_ids, diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 60f6c0bdbc..0279153752 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -186,6 +186,13 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( rank=0, mtp=2, ) + if source_device != 0: + # 此测试直接传入另一张卡的 Tensor，未走 IPC 导入；装配阶段先建立 peer 映射。 + # 生产拷贝入口只使用已映射的资源，不负责打开设备间访问。 + from cuda.bindings import runtime + + error, = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) for i, field in enumerate(full_state_test_fields(source_full, full_state_test_buffers(source_full)).values()): if "token" in field.axes: field.tensor.fill_(i + 5) @@ -317,7 +324,8 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( pool.conv_state.zero_() pool.ssm_state.zero_() slot = resources.bounded_att_big_page_allocator.alloc(1) - resources.save_bounded_att_pages([0], page_kind="big", slot_ids=slot) + with torch.cuda.device(source_device): + resources.save_bounded_att_pages([0], page_kind="big", slot_ids=slot) mapping = torch.zeros((3, 16), dtype=torch.int32, device=f"cuda:{source_device}") mapping[1, 0] = 7 torch.cuda.synchronize(source_device) From c92e5f8625086d4ae065d08956a76ef66a16b687 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 11:02:14 +0000 Subject: [PATCH 18/56] perf: optimize full attention page transfers --- docs/cache_manager_design.md | 22 +++- .../att_resources/full_att/gqa_transfer.py | 42 ++++++-- .../full_att/group_quantized_transfer.py | 42 ++++++-- .../att_resources/full_att/mla_transfer.py | 42 ++++++-- .../full_att/packed_mla_transfer.py | 42 ++++++-- .../full_att/static_fp8_transfer.py | 42 ++++++-- .../common/att_resources/test_gqa_transfer.py | 100 ++++++++++++++++++ 7 files changed, 296 insertions(+), 36 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 9a3244952e..3498017d9a 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1212,7 +1212,27 @@ DP 副本之间使用相同的 State 成员名。`copy_from_other_dp` 接收远 full-attention 使用 `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_page_state_from_other_dp` 与 `copy_run_time_state_from_other_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源索引可为 int32/int64，目标索引须为 int64（当前 `index_copy_` 的要求）。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 -CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=36)` 控制每个字段发射的 program 数；kernel 从 `program_id * BLOCK` 开始，按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 保持各自的数据操作路径。 +CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=24, num_warps=8, num_stages=1)` 控制每个字段的发射配置，默认 `BLOCK=1024`；kernel 从 `program_id * BLOCK` 开始，以 `tl.range(..., num_stages=NUM_STAGES)` 按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 保持各自的数据操作路径。 + +full-attention 的宽类型搬运只调整字段最内轴，不合并具有独立寻址语义的 layer/token/head/page 轴。两端末轴连续，且行宽、起点、storage_offset、外层 stride 都满足对齐时，依次尝试 8/4/2 字节的整数 view；不满足时保留原元素宽度。view 不复制数据、不转换数值，WIDTH、ELEMENT_BYTES、各轴 stride 和搬运单元总数均由转换后的视图得到，浮点 NaN、量化 payload 和 scale 的原始位模式保持不变。页起点、页跨度不满足宽对齐或最内维不连续时须正确回退；TP 切片、乱序 token/page 索引、ready 页跳过、尾页及 packed padding 的语义不变。五类方案继续使用各自的专用入口和 kernel，不新增通用 State 或合并字段 launch。 + +full-attention 参数在 H200 上实测：先对 GQA 的小/中/大规模、CPU/PD 双向搬运比较 12 个组合，再在 INT4、INT8、静态 FP8、MLA 和三种 packed 布局上比较入选组合，共 68 个场景。候选包含 grid=24/36、BLOCK=256/512/1024/2048、warps=4/8、循环 stages=1/2/3/4；固定配置 24/1024/8/1 的速度几何平均约为旧实现的 2.12 倍。小块不保证更快，本轮最多增加约 1.74 μs。旧实现为未扩宽 view 的 grid=36、BLOCK=256、默认 4 warps、普通 range 循环。计时预先准备页索引并预热编译，随机顺序执行 CUDA Graph，三轮取中位数；每个候选均与旧实现逐字节对照。 + +进一步使用 16 个独立地址数据块连续发射，覆盖五类专用实现的 CPU/PD 双向操作，共 20 个场景。每块 16 层、1025 个有效 token、64 token/页，使用乱序物理 token 与页索引；GQA/量化 GQA 为 8 个 KV head、head_dim=128，MLA 使用单 latent head。核函数七轮取平均，完整 API 五轮取平均；块间不做同步。下面是代表性结果（GB/s，按有效载荷计），完整 API 包含视图准备、页索引上传、发射及末尾同步，不含 buffer 分配或 PD 网络收发。新配置未测试与推理计算并发的影响。 + +| 布局及方向 | 旧核带宽 | 新核带宽 | 旧完整 API | 新完整 API | +| --- | ---: | ---: | ---: | ---: | +| GQA CPU→GPU | 11.72 | 43.83 | 11.63 | 43.26 | +| GQA GPU→CPU | 29.59 | 50.07 | 29.22 | 49.56 | +| GQA PD 页→KV | 30.44 | 205.40 | 30.08 | 198.41 | +| INT4 CPU→GPU | 7.34 | 43.91 | 7.03 | 42.47 | +| 静态 FP8 CPU→GPU | 6.10 | 43.73 | 6.08 | 42.13 | +| MLA CPU→GPU | 11.87 | 45.58 | 11.50 | 43.81 | +| packed FP8 MLA CPU→GPU | 7.02 | 43.55 | 6.90 | 39.12 | +| packed FP8 MLA PD 页→KV | 16.93 | 169.67 | 16.54 | 48.24 | + +packed MLA 等多字段布局的 Python 准备及逐字段发射开销仍会限制完整调用吞吐；因此核带宽不能直接当作完整 CPU cache 或 PD 吞吐。字段融合仍保留为后续优化，不在本轮引入。 + diff --git a/lightllm/common/att_resources/full_att/gqa_transfer.py b/lightllm/common/att_resources/full_att/gqa_transfer.py index 920ef23715..335ecc3d1f 100644 --- a/lightllm/common/att_resources/full_att/gqa_transfer.py +++ b/lightllm/common/att_resources/full_att/gqa_transfer.py @@ -20,14 +20,17 @@ def transfer_gqa_pages( page_indices: list[int], layout: StateTransferLayout, write: bool, - grid_num: int = 36, + grid_num: int = 24, page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 - grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 """ assert grid_num > 0, "grid_num must be positive" assert k_storage.is_cuda, "page transfers require CUDA State storage" @@ -106,12 +109,34 @@ def transfer_gqa_pages( copy_views.append((local, external)) # 阶段三：每个 K/V 字段一次 launch，覆盖全部层、有效 token 和外部页。 - # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: + # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 + # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 + # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 + if local.stride(-1) == external.stride(-1) == 1: + element_bytes = local.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + local.shape[-1] % elements_per_unit != 0 + or local.data_ptr() % transfer_bytes != 0 + or external.data_ptr() % transfer_bytes != 0 + or local.storage_offset() % elements_per_unit != 0 + or external.storage_offset() % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in local.stride()[:-1]) + or any(stride % elements_per_unit for stride in external.stride()[:-1]) + ): + continue + local = local.view(transfer_dtype) + external = external.view(transfer_dtype) + break elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_gqa_pages[(grid_num,)]( LOCAL=local, @@ -130,7 +155,9 @@ def transfer_gqa_pages( WIDTH=local.shape[3], ELEMENT_BYTES=local.element_size(), WRITE=write, - BLOCK=256, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, ) @@ -150,7 +177,7 @@ def _copy_gqa_pages( PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 token_count, # 本次有效 token 数；调用方保证大于 0。 - element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 # 编译期寻址配置：stride 的单位是元素数，不是字节数。 TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 @@ -162,7 +189,8 @@ def _copy_gqa_pages( ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 - BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 ): # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 @@ -187,7 +215,7 @@ def _copy_gqa_pages( grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK block_offsets = tl.arange(start=0, end=BLOCK) # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 - for block_start in range(program_start, element_count, grid_stride): + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): offsets = block_start + block_offsets valid = offsets < element_count token = offsets // (WIDTH * HEADS) % token_count diff --git a/lightllm/common/att_resources/full_att/group_quantized_transfer.py b/lightllm/common/att_resources/full_att/group_quantized_transfer.py index cb4ec6b4f1..402c254e86 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_transfer.py +++ b/lightllm/common/att_resources/full_att/group_quantized_transfer.py @@ -22,14 +22,17 @@ def transfer_group_quantized_pages( page_indices: list[int], layout: StateTransferLayout, write: bool, - grid_num: int = 36, + grid_num: int = 24, page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 - grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 """ assert grid_num > 0, "grid_num must be positive" assert k_storage.is_cuda, "page transfers require CUDA State storage" @@ -116,12 +119,34 @@ def transfer_group_quantized_pages( copy_views.append((local, external)) # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 - # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 及动态 scale 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: + # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 + # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 + # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 + if local.stride(-1) == external.stride(-1) == 1: + element_bytes = local.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + local.shape[-1] % elements_per_unit != 0 + or local.data_ptr() % transfer_bytes != 0 + or external.data_ptr() % transfer_bytes != 0 + or local.storage_offset() % elements_per_unit != 0 + or external.storage_offset() % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in local.stride()[:-1]) + or any(stride % elements_per_unit for stride in external.stride()[:-1]) + ): + continue + local = local.view(transfer_dtype) + external = external.view(transfer_dtype) + break elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_group_quantized_pages[(grid_num,)]( LOCAL=local, @@ -140,7 +165,9 @@ def transfer_group_quantized_pages( WIDTH=local.shape[3], ELEMENT_BYTES=local.element_size(), WRITE=write, - BLOCK=256, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, ) @@ -160,7 +187,7 @@ def _copy_group_quantized_pages( PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 token_count, # 本次有效 token 数；调用方保证大于 0。 - element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 # 编译期寻址配置：stride 的单位是元素数，不是字节数。 TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 @@ -172,7 +199,8 @@ def _copy_group_quantized_pages( ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 - BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 ): # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 @@ -197,7 +225,7 @@ def _copy_group_quantized_pages( grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK block_offsets = tl.arange(start=0, end=BLOCK) # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 - for block_start in range(program_start, element_count, grid_stride): + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): offsets = block_start + block_offsets valid = offsets < element_count token = offsets // (WIDTH * HEADS) % token_count diff --git a/lightllm/common/att_resources/full_att/mla_transfer.py b/lightllm/common/att_resources/full_att/mla_transfer.py index 45f68e56e7..6b666ec0e3 100644 --- a/lightllm/common/att_resources/full_att/mla_transfer.py +++ b/lightllm/common/att_resources/full_att/mla_transfer.py @@ -19,14 +19,17 @@ def transfer_mla_pages( page_indices: list[int], layout: StateTransferLayout, write: bool, - grid_num: int = 36, + grid_num: int = 24, page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, ) -> None: """直接搬运 latent；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 - grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 """ assert grid_num > 0, "grid_num must be positive" assert kv_storage.is_cuda, "page transfers require CUDA State storage" @@ -103,10 +106,32 @@ def transfer_mla_pages( copy_views.append((local, external)) # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 - # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 for local, external in copy_views: + # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 + # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 + # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 + if local.stride(-1) == external.stride(-1) == 1: + element_bytes = local.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + local.shape[-1] % elements_per_unit != 0 + or local.data_ptr() % transfer_bytes != 0 + or external.data_ptr() % transfer_bytes != 0 + or local.storage_offset() % elements_per_unit != 0 + or external.storage_offset() % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in local.stride()[:-1]) + or any(stride % elements_per_unit for stride in external.stride()[:-1]) + ): + continue + local = local.view(transfer_dtype) + external = external.view(transfer_dtype) + break elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_mla_pages[(grid_num,)]( LOCAL=local, @@ -125,7 +150,9 @@ def transfer_mla_pages( WIDTH=local.shape[3], ELEMENT_BYTES=local.element_size(), WRITE=write, - BLOCK=256, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, ) @@ -145,7 +172,7 @@ def _copy_mla_pages( PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 token_count, # 本次有效 token 数；调用方保证大于 0。 - element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 # 编译期寻址配置：stride 的单位是元素数，不是字节数。 TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 @@ -157,7 +184,8 @@ def _copy_mla_pages( ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 - BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 ): # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 @@ -182,7 +210,7 @@ def _copy_mla_pages( grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK block_offsets = tl.arange(start=0, end=BLOCK) # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 - for block_start in range(program_start, element_count, grid_stride): + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): offsets = block_start + block_offsets valid = offsets < element_count token = offsets // (WIDTH * HEADS) % token_count diff --git a/lightllm/common/att_resources/full_att/packed_mla_transfer.py b/lightllm/common/att_resources/full_att/packed_mla_transfer.py index 08e757899b..b26854c91e 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_transfer.py +++ b/lightllm/common/att_resources/full_att/packed_mla_transfer.py @@ -22,14 +22,17 @@ def transfer_packed_mla_pages( indexer_offset: int, indexer_dim: int, write: bool, - grid_num: int = 36, + grid_num: int = 24, page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, ) -> None: """直接搬运 latent 及配套字段；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 - grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 """ assert grid_num > 0, "grid_num must be positive" assert kv_storage.is_cuda, "page transfers require CUDA State storage" @@ -125,12 +128,34 @@ def transfer_packed_mla_pages( copy_views.append((local, external)) # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 - # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 KV、indexer 及各量化字段 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: + # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 + # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 + # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 + if local.stride(-1) == external.stride(-1) == 1: + element_bytes = local.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + local.shape[-1] % elements_per_unit != 0 + or local.data_ptr() % transfer_bytes != 0 + or external.data_ptr() % transfer_bytes != 0 + or local.storage_offset() % elements_per_unit != 0 + or external.storage_offset() % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in local.stride()[:-1]) + or any(stride % elements_per_unit for stride in external.stride()[:-1]) + ): + continue + local = local.view(transfer_dtype) + external = external.view(transfer_dtype) + break elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_packed_mla_pages[(grid_num,)]( LOCAL=local, @@ -149,7 +174,9 @@ def transfer_packed_mla_pages( WIDTH=local.shape[3], ELEMENT_BYTES=local.element_size(), WRITE=write, - BLOCK=256, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, ) @@ -169,7 +196,7 @@ def _copy_packed_mla_pages( PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 token_count, # 本次有效 token 数；调用方保证大于 0。 - element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 # 编译期寻址配置：stride 的单位是元素数，不是字节数。 TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 @@ -181,7 +208,8 @@ def _copy_packed_mla_pages( ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 - BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 ): # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 @@ -206,7 +234,7 @@ def _copy_packed_mla_pages( grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK block_offsets = tl.arange(start=0, end=BLOCK) # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 - for block_start in range(program_start, element_count, grid_stride): + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): offsets = block_start + block_offsets valid = offsets < element_count token = offsets // (WIDTH * HEADS) % token_count diff --git a/lightllm/common/att_resources/full_att/static_fp8_transfer.py b/lightllm/common/att_resources/full_att/static_fp8_transfer.py index ef2db76f17..bcbf55f871 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_transfer.py +++ b/lightllm/common/att_resources/full_att/static_fp8_transfer.py @@ -20,14 +20,17 @@ def transfer_static_fp8_pages( page_indices: list[int], layout: StateTransferLayout, write: bool, - grid_num: int = 36, + grid_num: int = 24, page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 - grid_num 控制每个字段发射的 program 数，默认 36；kernel 跳步循环覆盖全部元素。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 """ assert grid_num > 0, "grid_num must be positive" assert k_storage.is_cuda, "page transfers require CUDA State storage" @@ -108,12 +111,34 @@ def transfer_static_fp8_pages( copy_views.append((local, external)) # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 - # elements 是标量元素数，不是字节数；每个字段固定发射 grid_num 个 program。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 FP8 K/V 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: + # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 + # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 + # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 + if local.stride(-1) == external.stride(-1) == 1: + element_bytes = local.element_size() + for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if ( + local.shape[-1] % elements_per_unit != 0 + or local.data_ptr() % transfer_bytes != 0 + or external.data_ptr() % transfer_bytes != 0 + or local.storage_offset() % elements_per_unit != 0 + or external.storage_offset() % elements_per_unit != 0 + or any(stride % elements_per_unit for stride in local.stride()[:-1]) + or any(stride % elements_per_unit for stride in external.stride()[:-1]) + ): + continue + local = local.view(transfer_dtype) + external = external.view(transfer_dtype) + break elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_static_fp8_pages[(grid_num,)]( LOCAL=local, @@ -132,7 +157,9 @@ def transfer_static_fp8_pages( WIDTH=local.shape[3], ELEMENT_BYTES=local.element_size(), WRITE=write, - BLOCK=256, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, ) @@ -152,7 +179,7 @@ def _copy_static_fp8_pages( PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 token_count, # 本次有效 token 数；调用方保证大于 0。 - element_count, # layer 数 × token_count × HEADS × WIDTH，单位为标量元素。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 # 编译期寻址配置：stride 的单位是元素数，不是字节数。 TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 @@ -164,7 +191,8 @@ def _copy_static_fp8_pages( ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 - BLOCK: tl.constexpr, # 每个 program 处理的标量元素数，当前调用方传入 256。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 ): # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 @@ -189,7 +217,7 @@ def _copy_static_fp8_pages( grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK block_offsets = tl.arange(start=0, end=BLOCK) # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 - for block_start in range(program_start, element_count, grid_stride): + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): offsets = block_start + block_offsets valid = offsets < element_count token = offsets // (WIDTH * HEADS) % token_count diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index d86a671d50..089f81d8c3 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -84,3 +84,103 @@ def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_a for region in regions: last_page = page_view[1, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) assert (last_page[:, 1:].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +@pytest.mark.parametrize("num_stages", [1, 3]) +@pytest.mark.parametrize("byte_start, extra_bytes, dim_stride, element_bytes", [ + (8, 32, 1, 8), + (2, 32, 1, 2), + (4, 32, 1, 4), + (8, 34, 1, 2), + (8, 36, 1, 4), + (8, 32, 2, 2), +]) +def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( + page_device: str, + num_stages: int, + byte_start: int, + extra_bytes: int, + dim_stride: int, + element_bytes: int, + pd_args: "EasyDict", + monkeypatch: pytest.MonkeyPatch, +) -> None: + from lightllm.common.att_resources.full_att import gqa_transfer + + pd_args.pd_kv_page_size = 2 + config = FullAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=128, + dtype=torch.bfloat16, + ) + source = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") + source.initialize_storage(token_capacity=8) + target.initialize_storage(token_capacity=8) + for name in ("k_storage", "v_storage"): + storage = getattr(source, name) + if dim_stride != 1: + storage = torch.empty(size=(*storage.shape[:-1], 256), dtype=storage.dtype, device="cuda") + storage.view(torch.uint8).random_(0, 256) + setattr(source, name, storage[..., ::dim_stride]) + getattr(target, name).view(torch.uint8).fill_(173) + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=source.resource_id, + tokens_per_page=2, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + # 页起点、页跨度和最内维 stride 分别破坏宽类型对齐，回退必须保持完整字段及保护区。 + pages = torch.full( + size=(3, layout.byte_end + extra_bytes), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, byte_start:byte_start + layout.byte_end] + src_ids = torch.tensor(data=[3, 3, 1, 1, 7, 7], dtype=torch.int32, device="cuda")[::2] + dst_ids = torch.tensor(data=[6, 2, 4], dtype=torch.int64, device="cuda") + expected = {} + for name in ("k_storage", "v_storage"): + expected[name] = getattr(target, name).clone() + expected[name][:, dst_ids] = getattr(source, name)[:, src_ids] + + actual_kernel = gqa_transfer._copy_gqa_pages + write_units = [] + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + if kwargs["WRITE"]: + write_units.append(kwargs["ELEMENT_BYTES"]) + return actual_kernel[grid](**kwargs) + return launch + + monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) + for state, ids, write in ((source, src_ids, True), (target, dst_ids, False)): + transfer_gqa_pages( + k_storage=state.k_storage, + v_storage=state.v_storage, + token_indices=ids, + page_buffers=page_view, + page_indices=[2, 0], + layout=layout, + write=write, + num_stages=num_stages, + ) + torch.cuda.synchronize() + assert write_units == [element_bytes, element_bytes] + for name in ("k_storage", "v_storage"): + assert torch.equal(getattr(target, name).view(torch.uint8), expected[name].view(torch.uint8)) + assert (pages[1] == 91).all() + assert (pages[:, :byte_start] == 91).all() + assert (pages[:, byte_start + layout.byte_end:] == 91).all() + for region in regions: + tail_page = page_view[0, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + assert (tail_page[:, 1:].contiguous().view(torch.uint8) == 91).all() From 4138da890b7b350bfa271be96610b8f75fd81367 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 11:13:55 +0000 Subject: [PATCH 19/56] perf: fuse full attention DP copies with Triton --- docs/cache_manager_design.md | 19 +- .../att_resources/full_att/base_state.py | 3 +- .../att_resources/full_att/dp_transfer.py | 189 ++++++++++++++++++ .../att_resources/full_att/gqa_state.py | 18 +- .../full_att/group_quantized_state.py | 28 ++- .../att_resources/full_att/mla_state.py | 13 +- .../full_att/packed_mla_state.py | 14 +- .../full_att/static_fp8_state.py | 18 +- .../test_full_att_dp_transfer.py | 183 +++++++++++++++++ 9 files changed, 435 insertions(+), 50 deletions(-) create mode 100644 lightllm/common/att_resources/full_att/dp_transfer.py create mode 100644 unit_tests/common/att_resources/test_full_att_dp_transfer.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 3498017d9a..2c938ca743 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -297,6 +297,7 @@ lightllm/common/ │ │ ├── create_gqa_state.py # GQA 统一创建入口：选择普通精度、INT4/INT8 或静态 FP8 │ │ ├── gqa_state.py # GQAFullAttentionState：兼容 MHA / GQA / MQA │ │ ├── gqa_transfer.py # GQA CPU/PD 共用的 K/V 页传输入口与 kernel +│ │ ├── dp_transfer.py # 各 full-attention 方案的直接索引 DP 聚合复制 │ │ ├── group_quantized_state.py # GroupQuantizedFullAttentionState：INT4 / INT8 等 │ │ ├── group_quantized_transfer.py # INT4/INT8 KV 与动态 scale 的定制传输 │ │ ├── static_fp8_state.py # StaticFP8FullAttentionState @@ -1210,11 +1211,23 @@ DP 请求共享信息分为两个数组，均按请求在当前批次中的位 DP 副本之间使用相同的 State 成员名。`copy_from_other_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.resource_id)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_run_time_state_from_other_dp`，再清零目标 tail。调用方保证源 MTP 游标为 0、前缀处于 tail 可重置边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 -full-attention 使用 `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_page_state_from_other_dp` 与 `copy_run_time_state_from_other_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源索引可为 int32/int64，目标索引须为 int64（当前 `index_copy_` 的要求）。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 +full-attention 使用 `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_page_state_from_other_dp` 与 `copy_run_time_state_from_other_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 -CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=24, num_warps=8, num_stages=1)` 控制每个字段的发射配置，默认 `BLOCK=1024`；kernel 从 `program_id * BLOCK` 开始，以 `tl.range(..., num_stages=NUM_STAGES)` 按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 保持各自的数据操作路径。 +full-attention DP 的 `dp_transfer.py` 提供显式存储入口：`copy_gqa_tokens` 接收 K/V，供 GQA 和静态 FP8 使用；`copy_group_quantized_tokens` 另接收 K/V 动态 scale；`copy_mla_tokens` 接收完整 latent/packed 存储。三种入口共用四维 token 布局的 Triton kernel，一次 launch 搬运当前 State 的全部字段；编译时展开 Tensor tuple 为指针参数，不创建或上传指针表。按两端 token 索引及各自真实 stride 直接寻址，无 `index_select`、跨设备 payload `.to()` 或 `index_copy_` 中间存储。每个字段独立选择对齐的整数搬运宽度，支持两端不同 token 容量、非连续索引和存储视图；默认 grid=24、BLOCK=4096、warps=8、stages=1，通过跳步循环覆盖全部元素，地址计算使用 int64。静态 FP8 标定参数不搬运；packed MLA 按原 DP 语义复制完整行，包括 padding，区别于 CPU/PD 的有效字段页格式。调用方保证目标索引唯一、源目标区域不重叠，并负责远端 peer/IPC 映射、源数据就绪及存储生命周期；算子在目标设备当前流执行，不在内部建立 peer 映射或同步。 -full-attention 的宽类型搬运只调整字段最内轴，不合并具有独立寻址语义的 layer/token/head/page 轴。两端末轴连续，且行宽、起点、storage_offset、外层 stride 都满足对齐时，依次尝试 8/4/2 字节的整数 view；不满足时保留原元素宽度。view 不复制数据、不转换数值，WIDTH、ELEMENT_BYTES、各轴 stride 和搬运单元总数均由转换后的视图得到，浮点 NaN、量化 payload 和 scale 的原始位模式保持不变。页起点、页跨度不满足宽对齐或最内维不连续时须正确回退；TP 切片、乱序 token/page 索引、ready 页跳过、尾页及 packed padding 的语义不变。五类方案继续使用各自的专用入口和 kernel，不新增通用 State 或合并字段 launch。 +DP 拷贝在两张 H200 上对照旧 gather → `.to()` → scatter 路径测量：16 层、1025 个乱序 token，源/目标容量为 2048/3072；GQA/INT4/INT8 使用 8 个 KV head、head_dim=128。每种方案和方向预热后连续调用 20 次，三轮取中位数，计时包含 Python 入口及末尾设备同步，不含分配、初始化和编译；重复使用同一组存储，未测试与推理并发。 + +| DP 布局 | 同卡加速比 | 跨卡加速比 | 跨卡旧→新有效载荷带宽（GB/s） | +| --- | ---: | ---: | ---: | +| GQA BF16 | 2.46× | 1.54× | 202.83 → 311.41 | +| INT4 + scale | 3.33× | 2.82× | 107.44 → 302.90 | +| INT8 + scale | 4.14× | 2.65× | 117.73 → 312.54 | +| MLA BF16 | 1.70× | 1.29× | 178.72 → 230.66 | +| Packed MLA FP8 | 2.97× | 1.68× | 129.99 → 217.73 | + +CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=24, num_warps=8, num_stages=1)` 控制每个字段的发射配置，默认 `BLOCK=1024`；kernel 从 `program_id * BLOCK` 开始，以 `tl.range(..., num_stages=NUM_STAGES)` 按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 直接索引复制使用下述专用入口。 + +full-attention 的宽类型搬运只调整字段最内轴，不合并具有独立寻址语义的 layer/token/head/page 轴。两端末轴连续，且行宽、起点、storage_offset、外层 stride 都满足对齐时，依次尝试 8/4/2 字节的整数 view；不满足时保留原元素宽度。view 不复制数据、不转换数值，WIDTH、ELEMENT_BYTES、各轴 stride 和搬运单元总数均由转换后的视图得到，浮点 NaN、量化 payload 和 scale 的原始位模式保持不变。页起点、页跨度不满足宽对齐或最内维不连续时须正确回退；TP 切片、乱序 token/page 索引、ready 页跳过、尾页及 packed padding 的语义不变。五类 CPU/PD 方案继续使用各自的专用入口和 kernel，不新增通用 State 或合并字段 launch。 full-attention 参数在 H200 上实测：先对 GQA 的小/中/大规模、CPU/PD 双向搬运比较 12 个组合，再在 INT4、INT8、静态 FP8、MLA 和三种 packed 布局上比较入选组合，共 68 个场景。候选包含 grid=24/36、BLOCK=256/512/1024/2048、warps=4/8、循环 stages=1/2/3/4；固定配置 24/1024/8/1 的速度几何平均约为旧实现的 2.12 倍。小块不保证更快，本轮最多增加约 1.74 μs。旧实现为未扩宽 view 的 grid=36、BLOCK=256、默认 4 warps、普通 range 循环。计时预先准备页索引并预热编译，随机顺序执行 CUDA Graph，三轮取中位数；每个候选均与旧实现逐字节对照。 diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index 5e89139b44..941e7d65af 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -205,7 +205,8 @@ def copy_tokens_from_other_dp( ) -> None: """复制其他 DP 同一 TP 分片的完整 token 数据，包含量化参数及 indexer 字段。 - 源、目标索引须已位于各自存储的 CUDA 设备；目标索引为 int64，接口不转换设备或 dtype。 + 源、目标索引须已位于各自存储的 CUDA 设备，可为 int32/int64；接口不转换设备或 dtype。 + 调用方保证 peer/IPC 可访问、源数据就绪、目标索引唯一，且源目标数据区域不重叠。 """ raise NotImplementedError() diff --git a/lightllm/common/att_resources/full_att/dp_transfer.py b/lightllm/common/att_resources/full_att/dp_transfer.py new file mode 100644 index 0000000000..c9fc88d752 --- /dev/null +++ b/lightllm/common/att_resources/full_att/dp_transfer.py @@ -0,0 +1,189 @@ +"""Full-attention DP 复制：显式传入各方案的存储，一次发射直接搬运完整 token。""" + +import torch +import triton +import triton.language as tl + + +def copy_gqa_tokens( + src_k_storage: torch.Tensor, + src_v_storage: torch.Tensor, + dst_k_storage: torch.Tensor, + dst_v_storage: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """融合 GQA/静态 FP8 的 K/V 复制；静态标定参数不参与搬运。""" + _copy_token_storages( + src_storages=(src_k_storage, src_v_storage), + dst_storages=(dst_k_storage, dst_v_storage), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def copy_group_quantized_tokens( + src_k_storage: torch.Tensor, + src_v_storage: torch.Tensor, + src_k_scales: torch.Tensor, + src_v_scales: torch.Tensor, + dst_k_storage: torch.Tensor, + dst_v_storage: torch.Tensor, + dst_k_scales: torch.Tensor, + dst_v_scales: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """融合 INT4/INT8 的 K/V payload 和动态 scale，分别保留字段 dtype。""" + _copy_token_storages( + src_storages=(src_k_storage, src_v_storage, src_k_scales, src_v_scales), + dst_storages=(dst_k_storage, dst_v_storage, dst_k_scales, dst_v_scales), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def copy_mla_tokens( + src_storage: torch.Tensor, + dst_storage: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """复制 MLA 完整行；packed 方案同时保留 RoPE、scale、indexer 和行 padding。""" + _copy_token_storages( + src_storages=(src_storage,), + dst_storages=(dst_storage,), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def _copy_token_storages( + src_storages: tuple[torch.Tensor, ...], + dst_storages: tuple[torch.Tensor, ...], + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """共享的四维 token 布局搬运，不创建 payload 或索引的中间 Tensor。 + + 在目标设备当前流发射；调用方保证远端存储及索引已可通过 peer/IPC 访问， + 并负责源数据就绪、存储生命周期和完成事件。目标 token 不重复，源目标区域不重叠。 + """ + token_count = src_token_indices.numel() + if token_count == 0: + return + + # 阶段一：只构造整数 view。字段仍为 [layer, token_slot, head, width]； + # 各字段可有不同 dtype/宽度，两侧 token 容量和实际 stride 也可以不同。 + # 连续末轴优先按 8/4/2 字节搬运，保留 FP8、NaN、动态 scale 等原始位模式。 + sources, destinations = [], [] + for source, destination in zip(src_storages, dst_storages): + if source.stride(-1) == destination.stride(-1) == 1: + element_bytes = source.element_size() + for unit_bytes, unit_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if unit_bytes <= element_bytes: + break + ratio = unit_bytes // element_bytes + if ( + source.shape[-1] % ratio + or source.data_ptr() % unit_bytes or destination.data_ptr() % unit_bytes + or source.storage_offset() % ratio or destination.storage_offset() % ratio + or any(stride % ratio for stride in source.stride()[:-1]) + or any(stride % ratio for stride in destination.stride()[:-1]) + ): + continue + source = source.view(unit_dtype) + destination = destination.view(unit_dtype) + break + sources.append(source) + destinations.append(destination) + + # 阶段二：Tensor tuple 在编译时展开为独立指针参数，不上传指针表。 + # 一次 launch 完成当前方案全部字段；每个 program 用跳步循环覆盖任意 token 数。 + # 保持 24 个 program，使用 4096 个搬运单元提高跨卡读取吞吐，循环流水深度为 1。 + _copy_indexed_token_storages[(24,)]( + SOURCES=tuple(sources), + DESTINATIONS=tuple(destinations), + SRC_TOKEN_IDS=src_token_indices, + DST_TOKEN_IDS=dst_token_indices, + token_count=token_count, + SRC_INDEX_STRIDE=src_token_indices.stride(0), + DST_INDEX_STRIDE=dst_token_indices.stride(0), + SHAPES=tuple((tensor.shape[0], tensor.shape[2], tensor.shape[3]) for tensor in sources), + SRC_STRIDES=tuple(tensor.stride() for tensor in sources), + DST_STRIDES=tuple(tensor.stride() for tensor in destinations), + ELEMENT_BYTES=tuple(tensor.element_size() for tensor in sources), + BLOCK_ELEMENTS=4096, + NUM_STAGES=1, + num_warps=8, + ) + + +@triton.jit +def _copy_indexed_token_storages( + # 各字段的四维存储指针；K/V、scale 或 packed 行在一次 launch 中处理。 + SOURCES, + DESTINATIONS, + SRC_TOKEN_IDS, + DST_TOKEN_IDS, + token_count, + # 索引 Tensor 可以非连续，直接使用调用方提供的 int32/int64 索引。 + SRC_INDEX_STRIDE: tl.constexpr, + DST_INDEX_STRIDE: tl.constexpr, + SHAPES: tl.constexpr, # 各字段的 (layer_count, head_count, width)。 + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + ELEMENT_BYTES: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + token_count = tl.cast(token_count, tl.int64) + program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS + lane = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) + + for field in tl.static_range(len(SHAPES)): + # 每个字段独立计算地址；不同容量下 layer stride 不同，不能复用源端 stride。 + layers = SHAPES[field][0] + heads = SHAPES[field][1] + width = SHAPES[field][2] + if ELEMENT_BYTES[field] == 8: + copy_dtype = tl.uint64 + elif ELEMENT_BYTES[field] == 4: + copy_dtype = tl.uint32 + elif ELEMENT_BYTES[field] == 2: + copy_dtype = tl.uint16 + else: + copy_dtype = tl.uint8 + source = SOURCES[field].to(tl.pointer_type(copy_dtype)) + destination = DESTINATIONS[field].to(tl.pointer_type(copy_dtype)) + element_count = token_count * layers * heads * width + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段一：按 [token, layer, head, width] 展开；除数均为布局常量， + # 避免按动态 token_count 做逐元素的 int64 除法。 + element = block_start + lane + valid = element < element_count + column = element % width + head = element // width % heads + layer = element // (width * heads) % layers + token = element // (width * heads * layers) + + # 阶段二：读取两端物理槽位；索引值及后续地址乘法使用 int64 防止溢出。 + src_token = tl.load(SRC_TOKEN_IDS + token * SRC_INDEX_STRIDE, mask=valid, other=0).to(tl.int64) + dst_token = tl.load(DST_TOKEN_IDS + token * DST_INDEX_STRIDE, mask=valid, other=0).to(tl.int64) + src_offset = ( + layer * tl.cast(SRC_STRIDES[field][0], tl.int64) + + src_token * tl.cast(SRC_STRIDES[field][1], tl.int64) + + head * tl.cast(SRC_STRIDES[field][2], tl.int64) + + column * tl.cast(SRC_STRIDES[field][3], tl.int64) + ) + dst_offset = ( + layer * tl.cast(DST_STRIDES[field][0], tl.int64) + + dst_token * tl.cast(DST_STRIDES[field][1], tl.int64) + + head * tl.cast(DST_STRIDES[field][2], tl.int64) + + column * tl.cast(DST_STRIDES[field][3], tl.int64) + ) + + # 阶段三：直接从远端存储读出并写入目标槽位，无 gather 或跨设备中转存储。 + value = tl.load(source + src_offset, mask=valid, other=0) + tl.store(destination + dst_offset, value, mask=valid) diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index ff7ba301d4..172133380f 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -12,6 +12,7 @@ from ..types import AxisSlice, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .gqa_transfer import transfer_gqa_pages +from .dp_transfer import copy_gqa_tokens if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -207,15 +208,14 @@ def copy_tokens_from_other_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； - # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 - # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 - src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) - dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) - self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) - src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) - dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) - self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) + copy_gqa_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) # ======================================================================== # 五、PD 分离传输 diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index d71c7e0c67..948440b310 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -13,6 +13,7 @@ from ..types import AxisSlice, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .group_quantized_transfer import transfer_group_quantized_pages +from .dp_transfer import copy_group_quantized_tokens if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -266,21 +267,18 @@ def copy_tokens_from_other_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； - # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 - # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 - src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) - dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) - self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) - src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) - dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) - self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) - src_k_scale = remote_dp_full_att_state.k_scales.index_select(dim=1, index=src_token_indices) - dst_k_scale = src_k_scale.to(device=self.k_scales.device, non_blocking=True) - self.k_scales.index_copy_(dim=1, index=dst_token_indices, source=dst_k_scale) - src_v_scale = remote_dp_full_att_state.v_scales.index_select(dim=1, index=src_token_indices) - dst_v_scale = src_v_scale.to(device=self.v_scales.device, non_blocking=True) - self.v_scales.index_copy_(dim=1, index=dst_token_indices, source=dst_v_scale) + copy_group_quantized_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + src_k_scales=remote_dp_full_att_state.k_scales, + src_v_scales=remote_dp_full_att_state.v_scales, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + dst_k_scales=self.k_scales, + dst_v_scales=self.v_scales, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) # ======================================================================== # 五、PD 分离传输 diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index aa49881630..90844abd9f 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -12,6 +12,7 @@ from ..types import AxisSlice, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .mla_transfer import transfer_mla_pages +from .dp_transfer import copy_mla_tokens if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -183,12 +184,12 @@ def copy_tokens_from_other_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； - # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 - # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 - src_kv = remote_dp_full_att_state.storage.index_select(dim=1, index=src_token_indices) - dst_kv = src_kv.to(device=self.storage.device, non_blocking=True) - self.storage.index_copy_(dim=1, index=dst_token_indices, source=dst_kv) + copy_mla_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) # ======================================================================== # 五、PD 分离传输 diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index 0c86ff5784..ee2f5d5502 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -13,6 +13,7 @@ from ..types import AxisSlice, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .packed_mla_transfer import transfer_packed_mla_pages +from .dp_transfer import copy_mla_tokens if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -271,13 +272,12 @@ def copy_tokens_from_other_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； - # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 - # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 - # 整行包含 KV、动态 scale、RoPE、indexer 及 padding，直接复制可保留全部有效字段。 - src_kv = remote_dp_full_att_state.storage.index_select(dim=1, index=src_token_indices) - dst_kv = src_kv.to(device=self.storage.device, non_blocking=True) - self.storage.index_copy_(dim=1, index=dst_token_indices, source=dst_kv) + copy_mla_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) # ======================================================================== # 五、PD 分离传输 diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index 0e61d53bae..9e0d38f2fb 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -13,6 +13,7 @@ from ..types import AxisSlice, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .static_fp8_transfer import transfer_static_fp8_pages +from .dp_transfer import copy_gqa_tokens if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -292,15 +293,14 @@ def copy_tokens_from_other_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - # TODO：改用 Triton kernel 按源/目标 token 索引直接复制，融合字段搬运，去掉中间 gather 存储； - # 跨设备时源存储需可通过 peer/IPC 访问，并遵循调用方的 stream/event 依赖。 - # 同一 TP 分片复制完整 token；调用方提供各自设备上的 CUDA 索引，目标索引须为 int64。 - src_k = remote_dp_full_att_state.k_storage.index_select(dim=1, index=src_token_indices) - dst_k = src_k.to(device=self.k_storage.device, non_blocking=True) - self.k_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_k) - src_v = remote_dp_full_att_state.v_storage.index_select(dim=1, index=src_token_indices) - dst_v = src_v.to(device=self.v_storage.device, non_blocking=True) - self.v_storage.index_copy_(dim=1, index=dst_token_indices, source=dst_v) + copy_gqa_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) # ======================================================================== # 五、PD 分离传输 diff --git a/unit_tests/common/att_resources/test_full_att_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py new file mode 100644 index 0000000000..29859af637 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -0,0 +1,183 @@ +"""DP 直接索引复制：完整编码、真实 stride、非默认流和跨设备 IPC。""" + +from __future__ import annotations + +from multiprocessing.connection import Connection +from multiprocessing.reduction import ForkingPickler +import traceback + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState +from lightllm.common.att_resources.full_att.dp_transfer import copy_gqa_tokens +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, rebuild_full_state +from unit_tests.common.att_resources.test_capacity import definition + + +SCHEMES = ("gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", SCHEMES) +@pytest.mark.parametrize("source_device", [0, 1]) +@pytest.mark.parametrize("token_count", [0, 1, 513]) +def test_dp_copy_preserves_complete_rows_without_temporary_payload( + kind: str, + source_device: int, + token_count: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if source_device >= torch.cuda.device_count(): + pytest.skip("two CUDA devices required") + torch.cuda.set_device(0) + if source_device != 0: + if not torch.cuda.can_device_access_peer(0, source_device): + pytest.skip("CUDA peer access required") + # 仅测试直接跨卡 Tensor 的场景需要预先开 peer；生产通过 IPC 装配，不在复制时建映射。 + from cuda.bindings import runtime + error, = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) + + template = definition(kind) + source = rebuild_full_state(template, token_capacity=1024, device=f"cuda:{source_device}") + target = rebuild_full_state(template, token_capacity=1200, device="cuda:0") + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + expected, original = {}, {} + for name, tensor in source_buffers.items(): + tensor.view(torch.uint8).random_(0, 256) + target_buffers[name].view(torch.uint8).fill_(173) + original[name] = tensor.cpu().view(torch.uint8) + expected[name] = target_buffers[name].cpu().view(torch.uint8) + + # 源索引允许重复，目标唯一；两端都带 storage offset 和非单位 stride。 + src_list = [(i * 29) % 257 for i in range(token_count)] + dst_list = [(i * 7) % 1200 for i in range(token_count)] + src_dtype = torch.int32 if source_device == 0 else torch.int64 + dst_dtype = torch.int64 if source_device == 0 else torch.int32 + src_ids = torch.empty(size=(token_count * 2 + 1,), dtype=src_dtype, device=f"cuda:{source_device}")[1::2] + dst_ids = torch.empty(size=(token_count * 3 + 2,), dtype=dst_dtype, device="cuda:0")[2::3] + src_ids.copy_(torch.tensor(src_list, dtype=src_dtype, device=f"cuda:{source_device}")) + dst_ids.copy_(torch.tensor(dst_list, dtype=dst_dtype, device="cuda:0")) + for name in expected: + expected[name][:, dst_list] = original[name][:, src_list] + calibration = [] + if isinstance(target, StaticFP8FullAttentionState): + for tensor in (target.scales, target.q_scales): + if tensor is not None: + calibration.append((tensor, tensor.clone())) + + # 显式依赖事件，不让测试偶然依赖默认流；跨设备源索引也必须等生产完成。 + producer_done = torch.cuda.Event() + producer_done.record(torch.cuda.current_stream(device=source_device)) + stream = torch.cuda.Stream(device=0) + stream.wait_event(producer_done) + stream.wait_stream(torch.cuda.current_stream(device=0)) + + def reject_intermediate(*args: object, **kwargs: object) -> None: + pytest.fail("DP copies must use supplied CUDA indices and storage directly") + + with monkeypatch.context() as patch: + for method in ("index_select", "index_copy_", "to", "tolist", "cpu", "contiguous"): + patch.setattr(torch.Tensor, method, reject_intermediate) + with torch.cuda.stream(stream): + target.copy_tokens_from_other_dp( + src_token_indices=src_ids, + dst_token_indices=dst_ids, + remote_dp_full_att_state=source, + ) + stream.synchronize() + for name, tensor in target_buffers.items(): + # 比较全部存储：包含未选槽位、HOLD 页、packed padding 和任意浮点位模式。 + assert torch.equal(tensor.cpu().view(torch.uint8), expected[name]) + assert torch.equal(source_buffers[name].cpu().view(torch.uint8), original[name]) + for tensor, unchanged in calibration: + assert torch.equal(tensor, unchanged) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("column_stride", [1, 2]) +@pytest.mark.parametrize("offset", [0, 1, 2, 4]) +def test_dp_copy_uses_real_storage_strides(column_stride: int, offset: int) -> None: + # 宽类型 view 的 8/4/2 字节对齐及不连续末轴回退；源目标层/token/head stride 不同。 + source_base = torch.empty(size=(2, 19, 3, 32), dtype=torch.bfloat16, device="cuda") + target_base = torch.empty(size=(2, 23, 3, 48), dtype=torch.bfloat16, device="cuda") + source_base.view(torch.uint8).random_(0, 256) + target_base.view(torch.uint8).fill_(173) + source = source_base[:, ::2, :, offset:offset + 8 * column_stride:column_stride] + target = target_base[:, ::3, :, offset:offset + 8 * column_stride:column_stride] + expected = target_base.clone() + expected[:, ::3, :, offset:offset + 8 * column_stride:column_stride][:, [4, 0, 6]] = source[:, [7, 2, 7]] + copy_gqa_tokens( + src_k_storage=source, + src_v_storage=source, + dst_k_storage=target, + dst_v_storage=torch.empty_like(target), + src_token_indices=torch.tensor([7, 2, 7], dtype=torch.int32, device="cuda"), + dst_token_indices=torch.tensor([4, 0, 6], dtype=torch.int32, device="cuda"), + ) + torch.cuda.synchronize() + assert torch.equal(target_base.view(torch.uint8), expected.view(torch.uint8)) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_dp_copy_through_production_peer_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + states = [rebuild_full_state(definition(kind), token_capacity=16, device="cuda:1") for kind in SCHEMES] + for state in states: + for tensor in full_state_test_buffers(state).values(): + tensor.view(torch.uint8).random_(0, 256) + source_ids = torch.tensor([3, 99, 1, 99, 3], dtype=torch.int32, device="cuda:1")[::2] + torch.cuda.synchronize(device=1) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps((states, source_ids))) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_ipc_copy_worker, args=(child, payload)) + process.start() + child.close() + try: + assert parent.poll(60), "DP peer worker did not finish" + error = parent.recv() + assert error is None, error + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _ipc_copy_worker(connection: Connection, payload: bytes) -> None: + try: + torch.cuda.set_device(0) + sources, src_ids = ForkingPickler.loads(payload) + for kind, source in zip(SCHEMES, sources): + target = rebuild_full_state(definition(kind), token_capacity=24, device="cuda:0") + target_buffers = full_state_test_buffers(target) + for tensor in target_buffers.values(): + tensor.view(torch.uint8).fill_(173) + expected = {name: tensor.cpu().view(torch.uint8) for name, tensor in target_buffers.items()} + for name, tensor in full_state_test_buffers(source).items(): + expected[name][:, [8, 2, 5]] = tensor.cpu().view(torch.uint8)[:, [3, 1, 3]] + target.copy_tokens_from_other_dp( + src_token_indices=src_ids, + dst_token_indices=torch.tensor([8, 2, 5], dtype=torch.int64, device="cuda:0"), + remote_dp_full_att_state=source, + ) + torch.cuda.synchronize() + for name, tensor in target_buffers.items(): + assert torch.equal(tensor.cpu().view(torch.uint8), expected[name]) + connection.send(None) + connection.recv() + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() From 1f7007fa47c27f19140e34d4cf0dd323c69c99df Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 30 Oct 2026 11:24:54 +0000 Subject: [PATCH 20/56] refactor: simplify PD page capacity descriptors --- .../mode_backend/pd/nccl_kv_transporter.py | 18 +++++++++--------- .../mode_backend/pd/nixl_kv_transporter.py | 8 ++++---- .../model_infer/mode_backend/test_pd_notify.py | 16 +++++++++------- 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index a665dca08d..54b06bab00 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -91,8 +91,8 @@ def local_page_mem_desc(self) -> bytes: return pickle.dumps( { "num_pages": self.num_pages, - "page_shape": tuple(self.kv_move_buffer.shape[1:]), - "dtype": str(self.kv_move_buffer.dtype), + # PD buffer 固定为二维 uint8，页长度就是单页字节数。 + "page_bytes": self.kv_move_buffer.shape[1], } ) @@ -112,10 +112,7 @@ def connect_add_remote_agent(self, remote_agent: PDAgentMetadata): ), f"Peer name {metadata.agent_name} does not match remote name {remote_agent.agent_name}" page = pickle.loads(remote_agent.page_reg_desc) - if page.get("page_shape") != tuple(self.kv_move_buffer.shape[1:]) or page.get("dtype") != str( - self.kv_move_buffer.dtype - ): - raise ValueError("PD peers must use the same page shape and dtype") + assert page["page_bytes"] == self.kv_move_buffer.shape[1], "PD peers must use the same page byte capacity" self.remote_agents[remote_agent.agent_name] = remote_agent logger.info(f"Added NCCL remote agent {remote_agent.agent_name} at {metadata.host_ip}:{metadata.control_port}") return @@ -449,20 +446,23 @@ def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[Thr "allow_setattr": True, }, ) - # SO_REUSEADDR permits simultaneous binds before either socket - # listens. Claim the port here so EADDRINUSE retries in this - # thread instead of killing the background server after return. + # RPyC 构造时只 bind；启用 SO_REUSEADDR 后，并发进程可能都绑定同一端口成功。 + # 在当前线程先 listen，让端口冲突进入下面的重试；否则后台线程启动失败时， + # 外层已经返回该端口，无法捕获异常并切换到下一个端口。 server.listener.listen(server.backlog) threading.Thread(target=server.start, daemon=True).start() logger.info(f"NCCL RPyC control channel listen on {host_ip}:{cur_port}") return server, cur_port except OSError as e: + # 构造失败时 server 尚未赋值；已构造但 listen 失败时，关闭 socket 后再重试。 if server is not None: server.close() last_error = e if e.errno == errno.EADDRINUSE: + # 当前端口已被占用，或并发启动时被其他进程抢先监听，继续尝试下一个端口。 logger.info(f"NCCL RPyC control port {host_ip}:{cur_port} is in use, try next port") else: + # 其他 socket 错误保留告警；整个端口范围都失败后，携带最后一个异常向上抛出。 logger.warning(f"Create NCCL RPyC control channel on {host_ip}:{cur_port} failed: {e}") raise RuntimeError(f"can not allocate NCCL control port in [{port_min}, {port_max}]") from last_error diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py index 147ba3ee25..ae4b62ff2b 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py @@ -62,17 +62,17 @@ def get_new_notifs(self) -> Dict[str, list[bytes]]: def _register_kv_move_buffer(self, kv_move_buffer: Tensor): self.num_pages = kv_move_buffer.shape[0] - self.page_len = kv_move_buffer[0].numel() * kv_move_buffer.element_size() + self.page_bytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() self.page_reg_desc = self.nixl_agent.register_memory(kv_move_buffer) self.page_local_xfer_handles = self._create_paged_xfer_handles(self.page_reg_desc, self.num_pages) def _create_paged_xfer_handles(self, reg_desc: "nixlBind.nixlRegDList", page_num: int, agent_name: str = ""): base_addr, registered_bytes, device_id, _ = reg_desc[0] - if registered_bytes != page_num * self.page_len: - raise ValueError("PD peers must use the same page byte capacity") + # NIXL 原生描述记录整块注册内存的字节数，按对端页数检查单页容量是否一致。 + assert registered_bytes == page_num * self.page_bytes, "PD peers must use the same page byte capacity" pages_data = [] for page_id in range(page_num): - pages_data.append((base_addr + page_id * self.page_len, self.page_len, device_id)) + pages_data.append((base_addr + page_id * self.page_bytes, self.page_bytes, device_id)) descs = self.nixl_agent.get_xfer_descs(pages_data, "VRAM") return self.nixl_agent.prep_xfer_dlist(agent_name, descs, "VRAM") diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py index 6af1ebcc8a..8f99c6add3 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py @@ -193,29 +193,31 @@ def test_receiver_rejects_payload_with_decode_local_mark(): def test_nixl_rejects_different_peer_page_capacity_before_registering_handles(): transporter = NixlKVTransporter.__new__(NixlKVTransporter) - transporter.page_len = 128 + transporter.page_bytes = 128 transporter.nixl_agent = SimpleNamespace(get_xfer_descs=Mock(), prep_xfer_dlist=Mock()) - with pytest.raises(ValueError, match="same page byte capacity"): + with pytest.raises(AssertionError, match="same page byte capacity"): transporter._create_paged_xfer_handles([(4096, 512, 0, None)], 2, "peer") transporter.nixl_agent.get_xfer_descs.assert_not_called() transporter._create_paged_xfer_handles([(4096, 256, 0, None)], 2, "peer") transporter.nixl_agent.get_xfer_descs.assert_called_once_with([(4096, 128, 0), (4224, 128, 0)], "VRAM") -def test_nccl_checks_peer_page_shape_before_accepting_agent(): +def test_nccl_checks_peer_page_bytes_before_accepting_agent(): from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclAgentMetadata transporter = NcclKVTransporter.__new__(NcclKVTransporter) transporter.remote_agents = {} - transporter.kv_move_buffer = SimpleNamespace(shape=(2, 128), dtype="torch.uint8") + transporter.kv_move_buffer = SimpleNamespace(shape=(2, 128)) + transporter.num_pages = 2 + assert pickle.loads(transporter.local_page_mem_desc) == {"num_pages": 2, "page_bytes": 128} peer = SimpleNamespace( agent_name="peer", agent_metadata=pickle.dumps(NcclAgentMetadata("peer", "127.0.0.1", 1234, 0)), - page_reg_desc=pickle.dumps({"page_shape": (256,), "dtype": "torch.uint8"}), + page_reg_desc=pickle.dumps({"page_bytes": 256}), ) - with pytest.raises(ValueError, match="same page shape"): + with pytest.raises(AssertionError, match="same page byte capacity"): transporter.connect_add_remote_agent(peer) assert not transporter.remote_agents - peer.page_reg_desc = pickle.dumps({"page_shape": (128,), "dtype": "torch.uint8"}) + peer.page_reg_desc = transporter.local_page_mem_desc transporter.connect_add_remote_agent(peer) assert transporter.remote_agents["peer"] is peer From 46943c4ce7790f8f141a027d455727ac1dedd0aa Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 04:31:31 +0000 Subject: [PATCH 21/56] fix: resolve INT8, FP8 and NSA attention regressions Allow independent INT8 KV buffers, support missing Q calibration with graph-safe dynamic quantization, and use the correct NSA logical layer field. Update the resource design documentation. Validation: targeted checks passed; four complete GSM8K runs finished across the three affected configurations. --- docs/cache_manager_design.md | 2 +- .../full_att/static_fp8_state.py | 14 ++++---- .../common/basemodel/attention/fa3/fp8.py | 35 +++++++++++++++---- .../basemodel/attention/triton/int8kv.py | 7 +--- .../quantization/q_per_head_fp8_quant.py | 3 +- .../layer_infer/transformer_layer_infer.py | 2 +- 6 files changed, 41 insertions(+), 22 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 2c938ca743..843828c552 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1680,7 +1680,7 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `create_full_att_states()`、`create_bounded_att_states()`、`create_att_tail_states()`，再通过 `set_att_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_creation()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 -4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 仍必须有 K/V 和 Q 标定。随后 backend 调用 `initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 +4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 必须有 K/V 标定，Q 标定可选：有 Q 标定时使用静态量化，未提供时按本次请求的 Q 分组动态量化，prefill/decode 使用同一规则。随后 backend 调用 `initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index 9e0d38f2fb..26bdbe3721 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -76,13 +76,14 @@ def bind_calibration( 行起点由管理器按 main、draft0..draft35 的完整 State 集合确定，不是逻辑层号。 同组逻辑层共用一行；State 内的行顺序与 logical_layer_groups 的物理层顺序一致。 - 无文件时仅 per_tensor 可使用单位 scale；per_head 必须包含 K/V 和 Q 标定。 + 无文件时仅 per_tensor 可使用单位 scale；per_head 必须有 K/V 标定，Q 标定可选。 + 未提供 Q 标定时由 attention backend 动态量化 Q。 """ assert self.k_storage is None and self.scales is None, "calibration must be bound once before storage allocation" config = self.config storage_layer_count = len(config.logical_layer_groups) if calibration is None: - assert config.scale_mode == "per_tensor", "per-head FP8 requires a calibration file with Q scales" + assert config.scale_mode == "per_tensor", "per-head FP8 requires a KV calibration file" scales = torch.ones(size=(storage_layer_count, 2), dtype=torch.float32) q_scales = None else: @@ -121,10 +122,11 @@ def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: k, v = all_scales.split(split_size=calibrated_heads, dim=-1) scales = torch.cat(tensors=(shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) q_config = calibration.get("q_calibration") - assert q_config, "per-head FP8 requires Q calibration" - q_scales = shard( - scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), q_config["num_head"], - ) + q_scales = None + if q_config is not None: + q_scales = shard( + scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), q_config["num_head"], + ) else: scales = scale_tensor(calibration, (calibrated_layer_count, 2)) q_scales = None diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index 78f4e2887b..f2233c8128 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -3,7 +3,10 @@ import torch from ..base_att import AttControl from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import q_per_head_static_fp8_quant +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 .fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState @@ -100,9 +103,18 @@ def _fp8_prefill_att( f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", ) storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - static_q_scales = state.q_scales[storage_layer_index] - q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) + q = q.reshape(q.shape[0], k_head_num, -1) + if state.q_scales is None: + # 旧标定文件仅包含 K/V；按本次请求的 Q 长度动态计算每个 KV head 的 scale。 + q, q_scale = q_per_head_fp8_quant( + q=q, + seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], + b1_start_loc=self.cu_seqlens_q, + ) + else: + static_q_scales = state.q_scales[storage_layer_index] + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) o = flash_attn_with_kvcache( q=q.reshape(-1, q_head_num, q_head_dim), k_cache=cache_k, @@ -201,9 +213,18 @@ def _fp8_decode_att( q_head_num = q.shape[1] att_batch_size = self.b_att_seq_len.shape[0] - static_q_scales = state.q_scales - q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales[storage_layer_index]) - q_scale = static_q_scales[storage_layer_index].view(1, k_head_num).expand(att_batch_size, k_head_num) + q = q.reshape(q.shape[0], k_head_num, -1) + if state.q_scales is None: + # 使用 FA3 的实际 Q 分组，兼容普通 decode、MTP 多 token 及零长度填充组。 + q, q_scale = q_per_head_fp8_quant( + q=q, + seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], + b1_start_loc=self.cu_seqlens_q, + ) + else: + static_q_scales = state.q_scales[storage_layer_index] + 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, diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py index d863ac4014..9da88626c8 100644 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ b/lightllm/common/basemodel/attention/triton/int8kv.py @@ -69,12 +69,7 @@ def _groupsize_quant_prefill_att( v_scale: torch.Tensor, alloc_func=torch.empty, ) -> torch.Tensor: - # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) - # batch_size = self.infer_state.b_seq_len.shape[0] - - assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr() - assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr() - + # 反量化内核分别使用 K/V 及 scale 的指针和 stride，兼容共享或独立存储。 total_token_num = self.infer_state.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) diff --git a/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py b/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py index a31292b4fc..08e548d779 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py @@ -82,7 +82,8 @@ def q_per_head_fp8_quant(q, seq_lens, b1_start_loc, scale_out=None, token_batch_ if scale_out is None: scale_out = torch.empty((B, H), dtype=torch.float32, device=q.device) if token_batch_ids is None: - token_batch_ids = torch.repeat_interleave(torch.arange(B, device=q.device), seq_lens) + # 输出 token 数已知，避免读取 GPU 上的长度总和，允许在 CUDA Graph 中调用。 + token_batch_ids = torch.repeat_interleave(torch.arange(B, device=q.device), seq_lens, output_size=T) _per_head_max_reduce_kernel[(B, H)]( q, diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index 1122d6badc..6d5097a00f 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -212,7 +212,7 @@ def _get_indices( state = infer_state.att_resource_manager.get_att_state( model_name=infer_state.att_resource_role, att_type="mla", - logical_layer_id=self.layer_num_, + logical_layer_id=self.layer_idx_, state_class=FullAttentionState, ) destindex_copy_indexer_ks( From bb1500e6feb6b8674546e857d5d161fb5d7623b3 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 05:27:07 +0000 Subject: [PATCH 22/56] Optimize K-pool tail transfer with circular window kernel --- docs/cache_manager_design.md | 20 +- .../att_resources/att_tail/kpool_state.py | 84 +---- .../att_resources/att_tail/kpool_transfer.py | 211 +++++++++++ .../att_resources/att_tail/transfer_utils.py | 49 --- lightllm/common/att_resources/types.py | 23 +- .../att_resources/test_kpool_transfer.py | 332 ++++++++++++++++++ .../common/att_resources/test_tail_state.py | 14 +- unit_tests/models/glm5_next/test_pd_cache.py | 4 +- 8 files changed, 600 insertions(+), 137 deletions(-) create mode 100644 lightllm/common/att_resources/att_tail/kpool_transfer.py delete mode 100644 lightllm/common/att_resources/att_tail/transfer_utils.py create mode 100644 unit_tests/common/att_resources/test_kpool_transfer.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 843828c552..80516a62db 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -323,7 +323,7 @@ lightllm/common/ │ └── att_tail/ │ ├── __init__.py │ ├── base_state.py # AttentionTailState：尾部状态抽象操作契约 -│ ├── transfer_utils.py # att_tail_buffer_copies：本地 ring 与固定逻辑窗口之间搬运 +│ ├── kpool_transfer.py # transfer_kpool_pages：单次 kernel 双向搬运环形逻辑窗口 │ ├── create_kpool_state.py # K-pool 统一创建入口 │ └── kpool_state.py # KPoolAttentionTailState：原生 ring 与恢复规则 └── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 @@ -364,7 +364,7 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/transfer_utils.py` 接收本地 ring 和 prefill 前缀长度，按逻辑时间顺序搬运最近 `pool_size` 个位置；短前缀的不存在位置在写出时补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；tail 的复制视图由 State 在调用处遍历执行 `dst.copy_(src, non_blocking=True)`，不再包装仅转发循环的搬运函数。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时清零其余本地候选历史。三类路径都在提交前断言实际复制视图的 shape 一致。读回时交换本地和外部视图，不重新构造反向轴映射。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttResourceManager.get_cpu_transfer_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 @@ -1073,7 +1073,7 @@ CPU buffer 使用 `(page_num, page_bytes)` 的二维页池，通过 `cpu_page_bu | `regions` | 本 State 参与传输的物理字段或片段列表，同一字段可分为多个映射项 | | 每项的 `field`、`dtype`、`axes`、`shape` | 标识 State 中的字段及外部页区域；axes 命名各维度，shape 用元素数量；外部区域采用连续布局，stride 由 shape 和 dtype 推导 | | 每项的 `byte_start`、`byte_end` | 字段的左闭右开字节区间，组合后相对整页；长度等于 `prod(shape) * dtype.itemsize`，不包含字段间的对齐空隙，必须位于所属 State 区间内 | -| 每项的 `axis_slices` | `tuple[AxisSlice, ...]`，通过 `axis`、`src_start`、`dst_start`、`length` 描述轴名、源起点、目标起点和元素数量；一个 region 内同一轴只出现一次，同轴不连续范围拆成多个 region | +| 每项的 `axis_slices` | `tuple[AxisSlice, ...]`，通过 `axis`、`src_start`、`dst_start`、`length` 描述轴名、源起点、目标起点和元素数量；一个 region 内同一轴只出现一次，src_circular/dst_circular 描述两侧环形寻址；普通非连续范围拆成多个 region | | 每项的 `write_enabled` | 指定 TP 副本是否负责写入该区域，避免多个副本重复写入；恢复时各目标副本均可读取 | 以下冻结数据类只描述传输位置，不包含搬运方法： @@ -1089,6 +1089,8 @@ class AxisSlice: src_start: int dst_start: int length: int + src_circular: bool = False + dst_circular: bool = False @dataclass(frozen=True) @@ -1118,7 +1120,7 @@ class StateTransferLayout: 层维度统一完整、同序复制，搬运前检查源目标层数一致。逻辑层到物理存储位置的映射由 State 自身维护，例如逻辑层 `[3, 7]` 紧凑存为 `[0, 1]`，外部页沿用相同层顺序。CPU/PD 区域和 DP State 复制均不提供层子集选择或层重排参数。 -token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_slices`。其余搬运轴通过 `AxisSlice` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。使用多个连续片段表达同一轴的非连续范围。 +token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_slices`。其余搬运轴通过 `AxisSlice` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。普通非连续范围使用多个连续片段表达；环形轴使用 `src_circular` / `dst_circular`，按两侧 Tensor 的实际轴容量对 `start + i` 取模。标记始终按本地 State → 外部页解释，读回时反向执行；搬运长度不超过各环形轴容量。 `shape` 表达该外部字段区域的完整范围，而 `axis_slices` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： @@ -1647,15 +1649,17 @@ class AttentionTailState(ABC): | 接口 | 数据路径及约束 | | --- | --- | | `write_to_pd_buffer(...)` | 按 `state_token_len` 提取指定请求的逻辑窗口 → `pd_page_index` 指定的 PD 状态页区域，使用 layout 的资源偏移与 TP 写入规则 | -| `read_from_pd_buffer(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；清零该请求的其余候选历史，调用方建立接收完成依赖 | +| `read_from_pd_buffer(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；窗口外的 MTP 槽位保持原值，调用方建立接收完成依赖 | -这些方法传输 tail 自身的 prefill 逻辑窗口，不承担历史 KV、池化 key/scale、bounded 状态或 MTP proposal 的迁移。发送端取最近 `P` 个位置并按时间顺序排列；接收端先清零该请求的本地 ring，再把有效位置放回对应物理槽位，其他请求保持不变。后续 MTP 候选历史由 D 端执行过程中产生。 +这些方法传输 tail 自身的 prefill 逻辑窗口，不承担历史 KV、池化 key/scale、bounded 状态或 MTP proposal 的迁移。发送端取最近 `P` 个位置并按时间顺序排列；接收端同样按逻辑窗口遍历，单次 kernel 恢复窗口数据及短前缀补零，窗口之外的 MTP 槽位与其他请求保持不变。后续 MTP 候选历史由 D 端执行过程中产生。 -`StateTransferLayout` 描述外部页位置与切片。K-pool region 的轴为 `("layer", "window", "channel")`，形状为 `[layers, P, 2*head_dim]`，与本地 `mtp_step` 无关。`window` 按时间顺序对应 `[state_token_len-P, state_token_len)`；短前缀中小于 0 的位置在发送时补零，接收时忽略。独立状态页的 `tokens_per_page=None`，该窗口不是 full-attention 的页内 token 轴。 +`StateTransferLayout` 描述外部页位置与切片。K-pool region 的轴为 `("layer", "window", "channel")`，形状为 `[layers, P, 2*head_dim]`，与本地 `mtp_step` 无关。`window` 按时间顺序对应 `[state_token_len-P, state_token_len)`；短前缀中小于 0 的位置在发送和恢复时均补零，窗口外槽位保持原值。独立状态页的 `tokens_per_page=None`，该窗口不是 full-attention 的页内 token 轴。 + +K-pool 的 `window` 切片设置 `src_circular=True`。`kpool_transfer.py` 从 `state_token_len` 计算本次本地窗口起点；可通过 `axis_slice_overrides` 覆盖本次轴映射，不修改缓存的 layout。窗口起点及短前缀的无效位置数量作为运行时参数传入定制 Triton kernel，不随请求位置编译新变体。写出和读回各一次 launch，处理量均为 `layers * pool_size * channels`，不遍历窗口之外的 MTP 槽位；连续 channel 在对齐允许时使用最大 8 字节整数视图搬运，保留原始位模式。 P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_step`。位置 `pos` 在 P 端从 `pos % (P+S_p)` 读取，在 D 端写入 `pos % (P+S_d)`；不能直接截取 ring 的前 P 个物理槽位。模型算法、池化步长、层顺序、字段精度和逻辑维度仍须兼容。现有 main/draft prefill 保持相同的序列位置，因此可共用任务的 `end_kv_index`；位置变化不影响缓存的布局描述。 -例如 `P=4`、`S_p=0`、`S_d=2`、前缀长度 11，外部窗口保存位置 `[7,8,9,10]`；P 端读取槽位 `[3,0,1,2]`，D 端恢复到 `[1,2,3,4]`。D 端其余 ring 槽位清零。PD 块仍使用 KV 页确定的固定容量，减少 tail 有效区域和本地复制量不代表网络整页传输字节数同步减少。 +例如 `P=4`、`S_p=0`、`S_d=2`、前缀长度 11，外部窗口保存位置 `[7,8,9,10]`；P 端读取槽位 `[3,0,1,2]`，D 端恢复到 `[1,2,3,4]`。D 端其余 ring 槽位保持原值。PD 块仍使用 KV 页确定的固定容量，减少 tail 有效区域和本地复制量不代表网络整页传输字节数同步减少。 长度、前缀身份和格式来自保存/传输控制元数据，不能从原始 tail 字节反推。调用层负责核对元数据并维持 main/draft 的一致恢复点；State 负责布局和数据操作。所有 GPU 操作遵循当前 stream，返回不表示完成；源请求与目标页在读写完成前不能复用。 diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py index 2c1f7560b6..044a563583 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -8,7 +8,7 @@ from ..utils import dtype_name from ..types import AxisSlice, TransferRegion from .base_state import AttentionTailState -from .transfer_utils import att_tail_buffer_copies +from .kpool_transfer import transfer_kpool_pages if TYPE_CHECKING: from lightllm.common.att_resources.types import StateTransferLayout @@ -100,7 +100,6 @@ def reset_attention_tail_state( assert isinstance(req_indices, list), "request indices must be a list" for req_idx in req_indices: self.tail_buffer[:, req_idx].fill_(0) - # ======================================================================== # 三、PD 分离传输 # ======================================================================== @@ -123,6 +122,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: src_start=0, dst_start=0, length=window_size, + src_circular=True, ), AxisSlice( axis="channel", @@ -144,16 +144,15 @@ def write_to_pd_buffer( layout: StateTransferLayout, ) -> None: """将 prefill 末尾的逻辑窗口写入 PD 页；state_token_len 为已执行的前缀长度。""" - copies = self._pd_buffer_copies( - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - req_idx=req_idx, + assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + transfer_kpool_pages( + tail=self.tail_buffer[:, req_idx], + page_buffers=pd_page_buffers, + page_index=pd_page_index, state_token_len=state_token_len, layout=layout, write=True, ) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) def read_from_pd_buffer( self, @@ -163,74 +162,13 @@ def read_from_pd_buffer( state_token_len: int, layout: StateTransferLayout, ) -> None: - """按前缀长度将逻辑窗口放回本地 ring，清零其余候选历史，允许两端 MTP 步数不同。""" - # TODO：PD 恢复可按目标 ring 槽位一次写入有效历史或零，融合清零和逐窗口复制； - # 避免同一 kernel 内先清零再恢复造成竞争，同时清除窗口外的旧候选历史。 - copies = self._pd_buffer_copies( - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - req_idx=req_idx, - state_token_len=state_token_len, - layout=layout, - write=False, - ) - self.tail_buffer[:, req_idx].fill_(0) - for src_tensor, dst_tensor in copies: - dst_tensor.copy_(src_tensor, non_blocking=True) - - # ======================================================================== - # 四、私有辅助实现 - # ======================================================================== - - def _pd_buffer_copies( - self, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - req_idx: int, - state_token_len: int, - layout: StateTransferLayout, - write: bool, - ) -> list[tuple[torch.Tensor, torch.Tensor]]: - """为 PD prefill 交接生成复制视图，允许 P/D 使用不同的 MTP 步数。 - - 为什么不能直接复制物理 ring： - - 本地 ring 的容量为 P+S，其中 P=pool_size，S=本地 mtp_step；绝对位置 pos - 存在 pos % (P+S) 槽位。P/D 的 S 不同时，形状和同一位置对应的物理槽位都不同。 - - 即使 P 端没有执行 MTP decode，prefill 算子也按实际 ring 容量取模写入， - 最近 P 个位置不一定落在 ring 的前 P 个槽位，因此不能直接截取 [:P] 传输。 - - 外部页固定保存最近 P 个逻辑位置，按时间顺序排列，形状为 [layers, P, 2*head_dim]。 - 这样 PD 布局不依赖两端的 S；D 端的额外候选历史由后续 MTP 执行产生。 - - 位置和复制方向： - - state_token_len=L 是该 State 已执行完 prefill 的 token 数，不是最后一个 token - 的下标，也不包含刚采样但尚未执行的输出 token；来自 PD 任务的 end_kv_index。 - - 外部 window[i] 对应绝对位置 pos=L-P+i。write=True 从本地 ring 的 - pos % (P+S) 读取并写入 window[i]；write=False 按相同规则恢复到本地 ring。 - - 例如 P=4、L=11，窗口保存位置 [7,8,9,10]。S_p=0 时从槽位 [3,0,1,2] - 写出，S_d=2 时恢复到槽位 [1,2,3,4]；直接复制同名物理槽位会恢复错误历史。 - - L

= 0, "PD tail requires a nonnegative prefix length" - # helper 按 L 和本地 ring 实际容量逐位置取视图，不创建索引 Tensor 或读取 GPU 游标。 - return att_tail_buffer_copies( + transfer_kpool_pages( tail=self.tail_buffer[:, req_idx], page_buffers=pd_page_buffers, - layout=layout, page_index=pd_page_index, state_token_len=state_token_len, - write=write, + layout=layout, + write=False, ) diff --git a/lightllm/common/att_resources/att_tail/kpool_transfer.py b/lightllm/common/att_resources/att_tail/kpool_transfer.py new file mode 100644 index 0000000000..06aed5d7a7 --- /dev/null +++ b/lightllm/common/att_resources/att_tail/kpool_transfer.py @@ -0,0 +1,211 @@ +"""K-pool 的 PD 逻辑窗口传输：一次 kernel 完成环形映射及窗口数据搬运。""" + +import torch +import triton +import triton.language as tl + +from ..types import AxisSlice, StateTransferLayout + + +# ======================================================================== +# 一、K-pool 传输：固定布局 + 本次窗口位置，直接操作原始存储 +# ======================================================================== + + +def transfer_kpool_pages( + tail: torch.Tensor, + page_buffers: torch.Tensor, + page_index: int, + state_token_len: int, + layout: StateTransferLayout, + write: bool, + axis_slice_overrides: tuple[AxisSlice, ...] = (), + grid_num: int = 24, +) -> None: + """在当前 CUDA stream 交换一个请求的 tail 和选定 PD 页，不创建中间 Tensor。 + + 本地 tail 为 [layer, ring_slot, channel]，ring 容量为 pool_size + 本地 mtp_step； + 外部页保存最近 pool_size 个逻辑位置，按时间顺序排列。P/D 的 ring 容量可以不同。 + 窗口起点由 state_token_len 推导；axis_slice_overrides 可覆盖本次 window 映射， + 其余轴沿用缓存布局。K-pool 只支持 window 环形寻址，channel 保持完整、同序对应。 + 两个方向都只遍历逻辑窗口，短前缀对应的窗口内位置补零； + 本地窗口之外的 MTP 槽位保持原值，不在传输时清理。 + 调用方负责传输完成依赖及存储生命周期；返回不表示数据已经就绪。 + """ + # 阶段一：解析选定页的 tail 区域，统一为 [layer, window, channel] 视图。 + # K-pool 布局由 State 创建；此处直接使用约定的单个 tail region。 + region = layout.regions[0] + if write and not region.write_enabled: + return + + page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) + field_bytes = page_bytes[region.byte_start:region.byte_end] + external_tail = field_bytes.view(region.dtype).view(region.shape) + axis_order = tuple(region.axes.index(axis) for axis in ("layer", "window", "channel")) + external_tail = external_tail.permute(axis_order) + local_ring_size = tail.shape[1] + external_ring_size = external_tail.shape[1] + channel_count = tail.shape[2] + + # 阶段二：确定本次窗口。显式覆盖的起点直接使用；否则从 prefill 长度推导。 + # 只计算整数位置，不修改 layout，也不创建新的 AxisSlice 对象。 + layout_slices = {item.axis: item for item in region.axis_slices} + slice_overrides = {item.axis: item for item in axis_slice_overrides} + window_slice = slice_overrides.get("window", layout_slices["window"]) + channel_slice = slice_overrides.get("channel", layout_slices["channel"]) + window_size = window_slice.length + local_window_start = window_slice.src_start + external_window_start = window_slice.dst_start + + if window_slice.src_circular: + if "window" not in slice_overrides: + # 外部 window[i] 对应绝对位置 state_token_len - window_size + i。 + local_window_start += state_token_len - window_size + local_window_start %= local_ring_size + if window_slice.dst_circular: + external_window_start %= external_ring_size + invalid_prefix = max(window_size - state_token_len, 0) + + # 定制 kernel 完整处理 key/gate；窗口不能重复覆盖或越过线性存储边界。 + assert ( + channel_slice.src_start == channel_slice.dst_start == 0 + and channel_slice.length == channel_count + and not channel_slice.src_circular and not channel_slice.dst_circular + ), "K-pool transfers require the complete channel axis" + assert ( + 0 < window_size <= min(local_ring_size, external_ring_size) + and (window_slice.src_circular or 0 <= local_window_start <= local_ring_size - window_size) + and (window_slice.dst_circular or 0 <= external_window_start <= external_ring_size - window_size) + ), "window exceeds storage capacity" + + # 阶段三：只扩大连续 channel 的搬运单位，layer/window 维度和环形语义保持不变。 + # 例如 BF16 的 [layers, ring_size, 256] 按 INT64 搬运时变为 [layers, ring_size, 64]。 + local_tail = tail + if local_tail.stride(-1) == external_tail.stride(-1) == 1: + element_bytes = local_tail.element_size() + for unit_bytes, unit_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): + if unit_bytes <= element_bytes: + break + elements_per_unit = unit_bytes // element_bytes + # 每个 channel 行都必须能完整分组，不能跨 window 拼接搬运单元。 + if channel_count % elements_per_unit != 0: + continue + # Tensor 起点和 storage_offset 都须满足宽类型 view 的对齐要求。 + if ( + local_tail.data_ptr() % unit_bytes != 0 + or external_tail.data_ptr() % unit_bytes != 0 + or local_tail.storage_offset() % elements_per_unit != 0 + or external_tail.storage_offset() % elements_per_unit != 0 + ): + continue + # layer/window 的步长必须能换算为新类型的整数元素数。 + if ( + any(stride % elements_per_unit != 0 for stride in local_tail.stride()[:-1]) + or any(stride % elements_per_unit != 0 for stride in external_tail.stride()[:-1]) + ): + continue + local_tail = local_tail.view(unit_dtype) + external_tail = external_tail.view(unit_dtype) + break + + # 阶段四：正反方向都只处理逻辑窗口，窗口外的 MTP 槽位不参与读写。 + # grid_num 是发射上限，大窗口由 kernel 的跳步循环处理。 + element_count = local_tail.shape[0] * window_size * local_tail.shape[2] + block_elements = 1024 + program_count = min(grid_num, triton.cdiv(element_count, block_elements)) + _transfer_kpool_pages[(program_count,)]( + LOCAL=local_tail, + EXTERNAL=external_tail, + local_window_start=local_window_start, + external_window_start=external_window_start, + invalid_prefix=invalid_prefix, + LOCAL_STRIDES=local_tail.stride(), + EXTERNAL_STRIDES=external_tail.stride(), + LAYERS=local_tail.shape[0], + LOCAL_RING_SIZE=local_ring_size, + EXTERNAL_RING_SIZE=external_ring_size, + WINDOW_SIZE=window_size, + CHANNELS=local_tail.shape[2], + ELEMENT_BYTES=local_tail.element_size(), + SRC_CIRCULAR=window_slice.src_circular, + DST_CIRCULAR=window_slice.dst_circular, + WRITE=write, + BLOCK_ELEMENTS=block_elements, + num_warps=4, + ) + + +# ======================================================================== +# 二、定制 kernel：计算窗口位置 → 读取有效历史 → 写出数据或零 +# ======================================================================== + + +@triton.jit(do_not_specialize=["local_window_start", "external_window_start", "invalid_prefix"]) +def _transfer_kpool_pages( + LOCAL, # 当前请求的本地 ring：[layer, ring_slot, channel]。 + EXTERNAL, # 已选定页中的 tail 字段：[layer, window, channel]。 + local_window_start, # 本次逻辑窗口在本地 ring 的起点，保持为运行时参数。 + external_window_start, # 本次窗口在外部字段中的起点，通常为 0。 + invalid_prefix, # 短前缀中不存在的窗口前部位置数量。 + LOCAL_STRIDES: tl.constexpr, + EXTERNAL_STRIDES: tl.constexpr, + LAYERS: tl.constexpr, + LOCAL_RING_SIZE: tl.constexpr, + EXTERNAL_RING_SIZE: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + CHANNELS: tl.constexpr, + ELEMENT_BYTES: tl.constexpr, + SRC_CIRCULAR: tl.constexpr, + DST_CIRCULAR: tl.constexpr, + WRITE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, +): + if ELEMENT_BYTES == 8: + copy_dtype = tl.uint64 + elif ELEMENT_BYTES == 4: + copy_dtype = tl.uint32 + elif ELEMENT_BYTES == 2: + copy_dtype = tl.uint16 + else: + copy_dtype = tl.uint8 + LOCAL = LOCAL.to(tl.pointer_type(copy_dtype)) + EXTERNAL = EXTERNAL.to(tl.pointer_type(copy_dtype)) + local_window_start = tl.cast(local_window_start, tl.int64) + external_window_start = tl.cast(external_window_start, tl.int64) + element_count = LAYERS * WINDOW_SIZE * CHANNELS + program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS + lanes = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) + + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=1): + # 两个方向共用逻辑窗口坐标，只在 load/store 时交换数据流向。 + element = block_start + lanes + valid = element < element_count + channel = element % CHANNELS + window = element // CHANNELS % WINDOW_SIZE + layer = element // (CHANNELS * WINDOW_SIZE) + local_slot = local_window_start + window + if SRC_CIRCULAR: + local_slot %= LOCAL_RING_SIZE + external_slot = external_window_start + window + if DST_CIRCULAR: + external_slot %= EXTERNAL_RING_SIZE + has_history = (window >= invalid_prefix) & (window < WINDOW_SIZE) + + # 所有 stride 在参与地址乘加前转为 int64，支持请求切片产生的大 layer stride。 + local_offset = ( + layer * tl.cast(LOCAL_STRIDES[0], tl.int64) + + local_slot * tl.cast(LOCAL_STRIDES[1], tl.int64) + + channel * tl.cast(LOCAL_STRIDES[2], tl.int64) + ) + external_offset = ( + layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) + + external_slot * tl.cast(EXTERNAL_STRIDES[1], tl.int64) + + channel * tl.cast(EXTERNAL_STRIDES[2], tl.int64) + ) + if WRITE: + value = tl.load(LOCAL + local_offset, mask=valid & has_history, other=0) + tl.store(EXTERNAL + external_offset, value, mask=valid) + else: + value = tl.load(EXTERNAL + external_offset, mask=valid & has_history, other=0) + tl.store(LOCAL + local_offset, value, mask=valid) diff --git a/lightllm/common/att_resources/att_tail/transfer_utils.py b/lightllm/common/att_resources/att_tail/transfer_utils.py deleted file mode 100644 index c77174cb2e..0000000000 --- a/lightllm/common/att_resources/att_tail/transfer_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Attention tail 的 PD 页传输：在本地 ring 和固定逻辑窗口之间搬运。""" - -import torch - -from ..types import StateTransferLayout - - -def att_tail_buffer_copies( - tail: torch.Tensor, - page_buffers: torch.Tensor, - layout: StateTransferLayout, - page_index: int, - state_token_len: int, - write: bool, -) -> list[tuple[torch.Tensor, torch.Tensor]]: - """按绝对位置映射 ring 与逻辑窗口，短前缀写出时将不存在的位置补零。 - - tail 轴为 (layer, ring_slot, channel)，外部页轴为 (layer, window, channel)。 - window 按时间顺序对应 [state_token_len - pool_size, state_token_len)， - 本地物理槽位使用绝对位置对 tail.shape[1] 取模，因此不要求 P/D 的 MTP 步数相同。 - 返回有效位置的 (源视图, 目标视图)，由 State 提交复制；读回前由 State 清零本地 ring。 - """ - assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" - page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) - copies = [] - for region in layout.regions: - assert region.field == "tail", f"unknown transfer field {region.field}" - external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) - external = external.permute(tuple(region.axes.index(axis) for axis in ("layer", "window", "channel"))) - assert tail.dtype == external.dtype, "local and external field dtypes must match" - assert (tail.shape[0], tail.shape[2]) == (external.shape[0], external.shape[2]), ( - "local and external tail layer/channel shapes must match" - ) - if write and not region.write_enabled: - continue - # TODO：用融合传输 kernel 完成逻辑窗口与 ring 的取模映射、写出补零和数据搬运， - # 减少逐窗口槽位 copy_/fill_；必须保留 P/D 不同 MTP 步数及 TP 副本只由 rank 0 写出的语义。 - # 请求级低频搬运；逐槽位取视图，不创建索引 Tensor，也不读取 GPU 上的游标。 - window_start = state_token_len - external.shape[1] - for window_index in range(external.shape[1]): - position = window_start + window_index - external_view = external[:, window_index] - if position < 0: - if write: - external_view.fill_(0) - continue - local_view = tail[:, position % tail.shape[1]] - copies.append((local_view, external_view) if write else (external_view, local_view)) - return copies diff --git a/lightllm/common/att_resources/types.py b/lightllm/common/att_resources/types.py index c9cb2c49ac..8082fc70ca 100644 --- a/lightllm/common/att_resources/types.py +++ b/lightllm/common/att_resources/types.py @@ -66,13 +66,21 @@ class TensorField: @dataclass(frozen=True) class AxisSlice: - """某一轴上的连续范围映射，起点与长度均以元素为单位。 + """某一轴上的逻辑连续范围映射，起点与长度均以元素为单位。 TransferRegion 中的方向固定为本地 State 到外部页，读取时反向应用。 + 环形轴按对应 Tensor 的实际轴长度取模，不要求两侧容量相同。 示例：本地 head[0:2] 写到全局传输区域 head[2:4]，读回时反向复制。 >>> head_slice = AxisSlice(axis="head", src_start=0, dst_start=2, length=2) >>> assert head_slice.dst_start + head_slice.length == 4 + + 环形示例：本地 window 容量为 4，从槽位 3 起搬运 4 个位置到线性外部页。 + 本地位置依次为 [3, 0, 1, 2]，外部位置依次为 [0, 1, 2, 3]。 + >>> window_slice = AxisSlice( + ... axis="window", src_start=3, dst_start=0, length=4, src_circular=True, + ... ) + >>> assert window_slice.src_circular and not window_slice.dst_circular """ # 要切片的轴名，必须同时存在于源、目标字段中，例如 head、channel、dim。 @@ -88,10 +96,19 @@ class AxisSlice: # 用于将本地 TP 分片放到全局 head/channel 的对应位置。 dst_start: int - # 搬运的元素数量，必须为正；源和目标的半开区间均为 [start, start + length)。 + # 搬运的元素数量，必须为正；普通轴使用半开区间 [start, start + length)。 + # 环形轴对 start + i 取模，length 不得超过该侧 Tensor 的轴容量，避免重复覆盖。 # 它是该轴上的长度，不是字节数，也不是整个 Tensor 的元素总数。 length: int + # 本地 State 的该轴采用环形寻址；容量取自本地 Tensor，而不是外部 region.shape。 + # 例如 tail/SWA 的 window 轴。读回时仍作用于本地侧，不随实际搬运方向交换含义。 + src_circular: bool = False + + # 外部区域的该轴采用环形寻址；容量取自外部 Tensor 的对应轴。 + # 常规 CPU/PD 页按逻辑顺序线性保存，使用默认 False。 + dst_circular: bool = False + # ======================================================================== # 三、CPU/PD 外部页布局 @@ -153,6 +170,8 @@ class TransferRegion: # 描述本地字段到外部区域的轴范围，每个非 layer/token 轴必须恰好出现一次。 # 无分片的轴也显式填写完整范围，例如 dim: 0 -> 0、length=head_dim。 # 读回时交换源、目标起点；同一组切片应用于全部层和本次有效 token。 + # 环形寻址由支持它的具体 State 搬运实现处理；请求窗口起点通过本次调用的 + # axis_slice_overrides 覆盖，不回写缓存布局，起点也不作为 kernel 的编译期常量。 axis_slices: tuple[AxisSlice, ...] # 当前 TP rank 是否负责向外部区域写入；False 不影响从该区域读回。 diff --git a/unit_tests/common/att_resources/test_kpool_transfer.py b/unit_tests/common/att_resources/test_kpool_transfer.py new file mode 100644 index 0000000000..e58399c1e9 --- /dev/null +++ b/unit_tests/common/att_resources/test_kpool_transfer.py @@ -0,0 +1,332 @@ +"""K-pool 融合传输：环形寻址、完整位模式、窗口外槽位保留及编译缓存。""" + +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.att_resources.att_tail import kpool_transfer +from lightllm.common.att_resources.types import AxisSlice, StateTransferLayout + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("dtype,head_dim,byte_start", [ + (torch.bfloat16, 128, 16), + (torch.float16, 33, 2), + (torch.float32, 128, 4), +]) +@pytest.mark.parametrize("page_device", ["cuda", "cpu"]) +def test_kpool_single_launch_preserves_bits_and_other_slots( + dtype: torch.dtype, + head_dim: int, + byte_start: int, + page_device: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = KPoolStateConfig( + logical_layer_ids=(0, 3, 7), + pool_size=4, + head_dim=head_dim, + dtype=dtype, + ) + source = KPoolAttentionTailState( + config=config, + request_capacity=3, + mtp_step=1, + device='cuda', + ) + target = KPoolAttentionTailState( + config=config, + request_capacity=3, + mtp_step=5, + device='cuda', + ) + source.tail_buffer.view(torch.uint8).random_(0, 256) + original = source.tail_buffer.clone() + target.tail_buffer.view(torch.uint8).fill_(173) + region = source.get_pd_transfer_regions()[0] + # 同时验证非默认的外部轴顺序和字段起点；不能假定 window/channel 连续排列。 + region = replace( + region, + axes=("channel", "layer", "window"), + shape=(2 * head_dim, 3, 4), + byte_start=byte_start, + byte_end=byte_start + region.byte_end, + ) + layout = StateTransferLayout( + resource_id=source.resource_id, + tokens_per_page=None, + byte_start=byte_start, + byte_end=region.byte_end, + regions=(region,), + ) + pages = torch.full( + size=(3, layout.byte_end + 16), + fill_value=253, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + kernel = kpool_transfer._transfer_kpool_pages + launches = [] + + class CountLaunches: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs: object) -> None: + launches.append((grid, kwargs["WRITE"])) + kernel[grid](**kwargs) + return launch + + def reject_copy(*args: object, **kwargs: object) -> None: + pytest.fail("K-pool transfer must copy in one Triton launch") + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(kpool_transfer, "_transfer_kpool_pages", CountLaunches()) + for method in ("copy_", "fill_", "zero_", "index_select", "index_copy_", "roll"): + patch.setattr(torch.Tensor, method, reject_copy) + with torch.cuda.stream(stream): + source.write_to_pd_buffer( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=1, + state_token_len=11, + layout=layout, + ) + target.read_from_pd_buffer( + pd_page_buffers=pages, + req_idx=2, + pd_page_index=1, + state_token_len=11, + layout=layout, + ) + stream.synchronize() + assert len(launches) == 2 and [item[1] for item in launches] == [True, False] + expected = torch.empty_like(target.tail_buffer[:, 2]) + expected.view(torch.uint8).fill_(173) + for position in range(7, 11): + expected[:, position % 9] = source.tail_buffer[:, 1, position % 5] + assert torch.equal(target.tail_buffer[:, 2].view(torch.uint8), expected.view(torch.uint8)) + assert (target.tail_buffer[:, :2].view(torch.uint8) == 173).all() + assert torch.equal(source.tail_buffer.view(torch.uint8), original.view(torch.uint8)) + assert (pages[[0, 2]] == 253).all() + assert (pages[1, :byte_start] == 253).all() and (pages[1, layout.byte_end:] == 253).all() + + +@pytest.mark.parametrize("src_circular", [False, True]) +@pytest.mark.parametrize("dst_circular", [False, True]) +def test_kpool_runtime_axis_overrides_are_bidirectional( + src_circular: bool, + dst_circular: bool, +) -> None: + config = KPoolStateConfig( + logical_layer_ids=(0, 2), + pool_size=4, + head_dim=128, + dtype=torch.bfloat16, + ) + state = KPoolAttentionTailState( + config=config, + request_capacity=2, + mtp_step=3, + device='cuda', + ) + state.tail_buffer.view(torch.uint8).random_(0, 256) + original = state.tail_buffer[:, 1].clone() + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + window = AxisSlice( + axis="window", + src_start=5 if src_circular else 1, + dst_start=3 if dst_circular else 0, + length=4, + src_circular=src_circular, + dst_circular=dst_circular, + ) + kpool_transfer.transfer_kpool_pages( + tail=state.tail_buffer[:, 1], + page_buffers=pages, + page_index=0, + state_token_len=11, + layout=layout, + write=True, + axis_slice_overrides=(window,), + grid_num=1, + ) + external = pages[0].view(config.dtype).view(2, 4, 256) + for index in range(4): + assert torch.equal( + external[:, (window.dst_start + index) % 4].view(torch.uint8), + original[:, (window.src_start + index) % 7].view(torch.uint8), + ) + state.tail_buffer[:, 1].fill_(37) + kpool_transfer.transfer_kpool_pages( + tail=state.tail_buffer[:, 1], + page_buffers=pages, + page_index=0, + state_token_len=11, + layout=layout, + write=False, + axis_slice_overrides=(window,), + grid_num=1, + ) + expected = torch.full_like(original, 37) + for index in range(4): + slot = (window.src_start + index) % 7 + expected[:, slot] = original[:, slot] + assert torch.equal(state.tail_buffer[:, 1].view(torch.uint8), expected.view(torch.uint8)) + assert layout.regions[0].axis_slices[0].src_start == 0 + + +def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: + config = KPoolStateConfig( + logical_layer_ids=(0,), + pool_size=4, + head_dim=128, + dtype=torch.bfloat16, + tp_rank=1, + tp_world_size=2, + ) + state = KPoolAttentionTailState( + config=config, + request_capacity=2, + mtp_step=2, + device='cuda', + ) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + pages.view(torch.bfloat16).fill_(13) + state.write_to_pd_buffer( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=0, + state_token_len=11, + layout=layout, + ) + assert (pages.view(torch.bfloat16) == 13).all() + state.read_from_pd_buffer( + pd_page_buffers=pages, + req_idx=1, + pd_page_index=0, + state_token_len=11, + layout=layout, + ) + expected = torch.zeros_like(state.tail_buffer[:, 1]) + for position in range(7, 11): + expected[:, position % 6].fill_(13) + assert torch.equal(state.tail_buffer[:, 1], expected) + + +def test_kpool_window_positions_do_not_create_kernel_variants() -> None: + config = KPoolStateConfig( + logical_layer_ids=(0,), + pool_size=4, + head_dim=128, + dtype=torch.bfloat16, + ) + state = KPoolAttentionTailState( + config=config, + request_capacity=2, + mtp_step=5, + device='cuda', + ) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=state.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + kernel = kpool_transfer._transfer_kpool_pages + cache_size = None + for length in (*range(30), (1 << 33) + 7): + state.write_to_pd_buffer( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=0, + state_token_len=length, + layout=layout, + ) + state.read_from_pd_buffer( + pd_page_buffers=pages, + req_idx=0, + pd_page_index=0, + state_token_len=length, + layout=layout, + ) + current_size = len(kernel.device_caches[torch.cuda.current_device()][0]) + if cache_size is None: + cache_size = current_size + assert current_size == cache_size + torch.cuda.synchronize() + + +@pytest.mark.parametrize("length", [0, 3, 515]) +def test_kpool_large_window_with_one_program(length: int) -> None: + config = KPoolStateConfig( + logical_layer_ids=(0, 3, 9, 12), + pool_size=128, + head_dim=128, + dtype=torch.bfloat16, + ) + source = KPoolAttentionTailState( + config=config, + request_capacity=2, + mtp_step=2, + device="cuda", + ) + target = KPoolAttentionTailState( + config=config, + request_capacity=2, + mtp_step=7, + device="cuda", + ) + source.tail_buffer.view(torch.uint8).random_(0, 256) + target.tail_buffer.fill_(37) + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + resource_id=source.resource_id, + tokens_per_page=None, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + for state, write in ((source, True), (target, False)): + kpool_transfer.transfer_kpool_pages( + tail=state.tail_buffer[:, 1], + page_buffers=pages, + page_index=0, + state_token_len=length, + layout=layout, + write=write, + grid_num=1, + ) + expected = torch.full_like(target.tail_buffer[:, 1], 37) + for position in range(length - 128, length): + if position < 0: + expected[:, position % 135].zero_() + else: + expected[:, position % 135] = source.tail_buffer[:, 1, position % 130] + assert torch.equal(target.tail_buffer[:, 1].view(torch.uint8), expected.view(torch.uint8)) + assert (target.tail_buffer[:, 0] == 37).all() diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py index 490bd91aaf..426f79b781 100644 --- a/unit_tests/common/att_resources/test_tail_state.py +++ b/unit_tests/common/att_resources/test_tail_state.py @@ -45,6 +45,7 @@ def tail_layout(state: KPoolAttentionTailState): src_start=0, dst_start=0, length=shape[1], + src_circular=True, ), AxisSlice( axis="channel", @@ -66,7 +67,8 @@ def test_tail_reset_clears_selected_requests_only(): assert (state.tail_buffer[:, 2] == 17).all() -@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("device", ["cuda"]) @pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) @pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) def test_pd_logical_window_restores_different_mtp_rings( @@ -105,9 +107,13 @@ def test_pd_logical_window_restores_different_mtp_rings( state_token_len=length, layout=layout, ) - expected = torch.zeros_like(target.tail_buffer[:, 2]) - for position in range(max(0, length - 4), length): - expected[:, position % (4 + target_step)] = source.tail_buffer[:, 1, position % (4 + source_step)] + expected = torch.full_like(target.tail_buffer[:, 2], 37) + for position in range(length - 4, length): + slot = position % (4 + target_step) + if position < 0: + expected[:, slot].zero_() + else: + expected[:, slot] = source.tail_buffer[:, 1, position % (4 + source_step)] assert torch.equal(target.tail_buffer[:, 2], expected) assert (target.tail_buffer[:, :2] == 37).all() assert torch.equal(source.tail_buffer, original) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 326aa13fae..c36d9e610d 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -166,6 +166,7 @@ def test_pd_kv_and_runtime_state_roundtrip( offset = end untouched = [] + original_tails = [] for mem in dest: for state in mem.full_states: state.storage.zero_() @@ -173,6 +174,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.ssm_state.normal_() for state in mem.tails: state.tail_buffer.normal_() + original_tails.append(torch.cat([state.tail_buffer[:, dst_req] for state in mem.tails]).clone()) untouched.append( [ x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() @@ -231,7 +233,7 @@ def test_pd_kv_and_runtime_state_roundtrip( assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) restored_tail = torch.cat([state.tail_buffer[:, dst_req] for state in mem.tails]) - expected_tail = torch.zeros_like(restored_tail) + expected_tail = original_tails[rank].clone() for position in range(length - 4, length): expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] assert torch.equal(restored_tail, expected_tail) From d63930da17062edaf69a86356df997fd28e775b1 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 06:12:21 +0000 Subject: [PATCH 23/56] Refine bounded state transfers and logical layer queries --- docs/cache_manager_design.md | 99 ++++++++----- .../att_resources/att_resource_manager.py | 66 ++++++--- .../att_resources/att_tail/base_state.py | 4 +- .../att_resources/att_tail/kpool_state.py | 4 +- .../att_resources/bounded_att/base_state.py | 38 +++-- .../att_resources/bounded_att/gdn_state.py | 44 ++++-- .../att_resources/bounded_att/kda_state.py | 44 ++++-- .../bounded_att/transfer_utils.py | 62 ++++----- .../att_resources/full_att/base_state.py | 4 +- .../att_resources/full_att/gqa_state.py | 4 +- .../full_att/group_quantized_state.py | 4 +- .../att_resources/full_att/mla_state.py | 4 +- .../full_att/packed_mla_state.py | 4 +- .../full_att/static_fp8_state.py | 4 +- lightllm/common/att_resources/types.py | 4 +- .../server/router/model_infer/infer_batch.py | 13 +- .../mode_backend/multi_level_kv_cache.py | 1 + .../att_resources/test_backend_startup.py | 14 +- .../att_resources/test_bounded_state.py | 51 ++++++- .../common/att_resources/test_capacity.py | 8 +- .../common/att_resources/test_layer_groups.py | 4 +- .../test_layer_state_selection.py | 8 +- .../att_resources/test_manager_helpers.py | 131 +++++++++++++++++- .../att_resources/test_model_resources.py | 8 +- .../att_resources/test_role_transfers.py | 8 +- .../att_resources/test_runtime_transport.py | 9 +- unit_tests/models/glm5_next/test_cache.py | 14 +- .../router/dynamic_prompt/test_radix_cache.py | 23 ++- .../model_infer/test_cpu_resource_layout.py | 9 +- .../model_infer/test_multimodal_chunking.py | 18 ++- 30 files changed, 537 insertions(+), 171 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 80516a62db..a13fac9c6a 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -364,7 +364,7 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 特化 `_copy_pair`，仅处理状态轴切片。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttResourceManager.get_cpu_transfer_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 @@ -482,9 +482,9 @@ State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get #### State 层归属与模型装配 -三类 State 均提供 `has_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 +三类 State 均提供 `contains_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 -`AttResourceManager.get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_name` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`att_type` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`has_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `att_type=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 +`AttResourceManager.get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_name` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`att_type` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `att_type=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_states_created=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。三类模型工厂只返回 State 列表，`prepare_att_resources()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_att_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 @@ -536,7 +536,7 @@ class FullAttentionState: # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... def get_storage_layer_index(self, logical_layer_id: int) -> int: ... @@ -605,7 +605,7 @@ class FullAttentionState: | 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | | `initialize_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | | `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | -| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | +| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | | `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | | `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | @@ -758,8 +758,8 @@ IPC 发布按用途调用 `create_pd_ipc_copy()` 或 `create_dp_ipc_copy()`， | --- | --- | --- | | full-attention token 复制 | `copy_full_att_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | | bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页 | -| bounded 状态页保存 | `save_bounded_att_pages(...)` | 将指定运行态保存到全部参与 bounded State 的同类 pool、同一组槽位 | -| bounded 状态页恢复 | `restore_bounded_att_pages(...)` | 按调用方选定的槽位，将各 bounded State 保存的页恢复到对应运行态位置 | +| bounded 状态页保存 | `save_bounded_att_pages(...)` | 按 state_token_lens 保存全部参与 bounded State 的同类 pool、同一组槽位 | +| bounded 状态页恢复 | `restore_bounded_att_pages(...)` | 按 state_token_lens 将选定快照恢复到对应请求运行态及 ring 位置 | | bounded 验证后提交 | `commit_bounded_att_state(req_indices, accepted_steps, *, role)` | 将指定角色的接受步数分发给参与 bounded State；由 State 选择已接受运行态，不暴露游标 Tensor | | attention tail 初始化与无快照恢复 | `reset_attention_tail_state(req_indices)` | 清零指定请求在全部参与 tail 中的状态；调用方保证空前缀或可重置的缓存边界，保留其他状态 | | CPU 写出 / 读回 | `write_req_to_cpu(...)` / `read_req_from_cpu(...)` | 组合 KV 和状态页，按调用参数提供的长度读写各 State 在外部页中的区域 | @@ -808,11 +808,19 @@ class AttResourceManager: # Bounded Attention 状态页：供 radix/CPU cache 复用，调用方直接分配和归还槽位。 def save_bounded_att_pages( - self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], ) -> None: ... def restore_bounded_att_pages( - self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], ) -> None: ... # MTP 推测解码：只提交指定角色；接受步数包含本轮已经执行的基准输入步。 @@ -844,6 +852,7 @@ class AttResourceManager: *, layouts: Mapping[str, StateTransferLayout], big_page_slot_ids: List[int] | None = None, req_idx: int | None = None, + state_token_len: int | None = None, ) -> None: ... # DP 缓存迁移：索引已由调用方分配，空 Tensor/列表表示不搬运该类数据。 @@ -879,7 +888,7 @@ class AttResourceManager: 调用方根据 `page_kind` 直接选择相应 allocator；该成员为 None 时不能申请该域。分配与归还使用 allocator 自身的接口，管理器不再增加包装方法。allocator 只管理空闲索引，不修改 State 的页数据或有效长度。 -`req_indices` 只定位运行态存储，不传递请求对象。保存/恢复时，请求状态槽位与页槽位逐项对应；目标槽位唯一且已分配。调用方维护快照与实际前缀的关联，资源方法不再接收重复的长度参数。 +`req_indices` 只定位运行态存储，不传递请求对象。保存/恢复时，请求状态槽位与页槽位逐项对应；目标槽位唯一且已分配。调用方维护快照与实际前缀的关联，并通过 `state_token_lens` 提供各快照的绝对长度，供 ring 等状态定位和恢复；页池不重复保存长度元数据。 #### bounded 空前缀运行态初始化 @@ -936,7 +945,7 @@ bounded_att_big_page_allocator 分配一次 slot_id = 17 └── GDNState C.big_pages_pool[17] ``` -调用 `save_bounded_att_pages(..., page_kind="big", slot_ids=[17])` 保存三个 State；`restore_bounded_att_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 +调用 `save_bounded_att_pages(..., state_token_lens=[256], page_kind="big", slot_ids=[17])` 保存三个 State；`restore_bounded_att_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 @@ -992,7 +1001,7 @@ CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能 | `"kv"` | `token_indices`；`req_idx=None` | 按单页容量读写全部参与 full-attention 的 KV、scale 和 packed 字段 | | `"att_state"` | `req_idx`；有 tail 时还需 `state_token_len`；`token_indices=None` | 读写 prefill 基准 bounded 状态与固定 tail 逻辑窗口，不经过大小页 pool | -PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再仅转交 tail State 完成 ring 重排。full/bounded State 的接口不增加长度参数；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 +PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再转交 bounded 和 tail State，用于各自的窗口定位及恢复。full-attention 的 token 传输接口保持不变；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 @@ -1417,7 +1426,7 @@ class BoundedAttentionState(ABC): # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== @abstractmethod - def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... def get_layer_view(self, logical_layer_id: int) -> Any: ... @@ -1426,12 +1435,20 @@ class BoundedAttentionState(ABC): @abstractmethod def save_bounded_att_pages( - self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], ) -> None: ... @abstractmethod def restore_bounded_att_pages( - self, req_indices: List[int], *, page_kind: PageKind, slot_ids: List[int], + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], ) -> None: ... @abstractmethod @@ -1478,18 +1495,30 @@ class BoundedAttentionState(ABC): @abstractmethod def write_to_pd_buffer( - self, req_idx: int, pd_page_buffers: torch.Tensor, *, pd_page_index: int, layout: StateTransferLayout, + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: ... @abstractmethod def read_from_pd_buffer( - self, pd_page_buffers: torch.Tensor, req_idx: int, *, pd_page_index: int, layout: StateTransferLayout, + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, ) -> None: ... ``` -bounded State 使用 `write_page_state_to_cpu` / `read_page_state_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 +bounded State 使用 `write_page_state_to_cpu` / `read_page_state_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 -CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。PD 状态页按单请求处理。`req_indices` 是 ReqManager 已分配的请求槽位，`slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_other_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 +CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`state_token_lens`、`slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。`state_token_lens` 表示各快照已经执行到的绝对 token 长度，不包含新采样但尚未执行的 token，也不是本批搬运的 token 数。调用方保证长度与选定快照一致；页池不新增长度记录。KDA/GDN 的 Conv/SSM 布局不依赖长度定位，保留此参数供 ring 等实现使用。 + +PD 状态页按单请求处理，bounded 和 tail 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_other_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 #### 运行态与推理访问 @@ -1497,7 +1526,7 @@ CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_pag | --- | --- | | 初始化 | 构造函数按 `request_capacity` 和 `mtp_step` 创建运行态及层映射，大小页池先为 `None`；不创建 allocator，不分配请求号或页槽位 | | `initialize_page_pools(...)` | 容量规划完成后创建大小页池，覆盖共享 allocator 的槽位域；容量为零的页池保持 `None`，不重建运行态或重置 MTP 游标 | -| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | +| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | | `reset_bounded_att_state(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | | `commit_bounded_att_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | @@ -1514,10 +1543,10 @@ prefill 的最终状态、普通 decode 的结果和已提交验证状态都可 | 接口 | 必需语义 | | --- | --- | -| `save_bounded_att_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req 与 slot 逐项对齐，前缀关联由调用方维护 | -| `restore_bounded_att_pages(...)` | 从调用方选定的槽位恢复运行态和算法控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | +| `save_bounded_att_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req、state_token_lens 与 slot 逐项对齐，ring 实现按长度确定窗口；前缀关联由调用方维护 | +| `restore_bounded_att_pages(...)` | 按 state_token_lens 从选定槽位恢复运行态及 ring 位置等控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | -调用方负责将快照关联到实际执行后的位置；KDA/GDN 的页保存契约只覆盖传入请求的 prefill 完成状态，不接收用于声称历史位置的长度参数。chunk 跨过中间边界但没有保留对应运行态时，预留槽位不能作为已保存的页面发布。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 +调用方负责将快照关联到实际执行后的位置；KDA/GDN 的页保存契约只覆盖传入请求的 prefill 完成状态；`state_token_lens` 必须是该实际状态的绝对长度，传入其他长度不能生成历史快照。chunk 跨过中间边界但没有保留对应运行态时，预留槽位不能作为已保存的页面发布。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 保存或 CPU 导入时直接写入指定槽位；异步完成前由上层保持 pending，完成后才发布或读取。发生失败直接透传异常，调用方不发布该批页面，并等已提交操作结束后才可回收。恢复前由调用方保证前缀关联和来源就绪，State 按兼容布局搬运数据。 @@ -1525,6 +1554,10 @@ prefill 的最终状态、普通 decode 的结果和已提交验证状态都可 #### CPU、DP、PD 的状态单位 +CPU 页与本地大小页池之间使用 `get_cpu_transfer_regions()` 返回的映射：`region.field` 选择页池中的 Conv/SSM 字段，`byte_start/byte_end` 和 `shape` 定位 CPU 全局字段，`axis_slices` 将本地 TP 分片对应到全局 head/channel 范围。`src_start` 固定指本地字段，`dst_start` 固定指外部页，实际读回时只反转搬运方向，不交换两者含义。`read_page_state_from_cpu` 显式准备目标 `pool_page_fields`，共用搬运流程直接展开这些映射，不通过 `_copy_pair` 隐藏切片过程。 + +窗口顺序的约定为：请求 ring → 按长度保存为逻辑时间顺序的大小页快照 → CPU 页；恢复时为 CPU 页 → 大页快照（完整页用普通槽位，碎页用 LOAD 槽位）→ 按绝对长度恢复请求 ring。CPU 页与页池之间只处理 TP 映射，不执行环形展开或恢复。 + | 路径 | 传输单位与接口职责 | | --- | --- | | CPU 写出 | `write_page_state_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，长度由缓存任务携带，不写入页内；每个待发布页的参与状态必须完整；CPU pool 到 CPU 页同样是独立的格式搬运 | @@ -1532,7 +1565,7 @@ prefill 的最终状态、普通 decode 的结果和已提交验证状态都可 | DP 大页复制 | `copy_big_page_state_from_other_dp(src_page_ids, dst_page_ids, remote_dp_bounded_state)` 在同一 TP 分片的 DP 副本之间复制大页 Conv/SSM 快照；不分配页、不修改请求运行态或小页，前缀关联与源页生命周期由调用方维护 | | DP 运行态复制 | `copy_run_time_state_from_other_dp(src_req_indices, dst_req_indices, remote_dp_bounded_state)` 复制同一 TP 分片的 prefill 基准 Conv/SSM；KDA/GDN 源游标必须为 0，目标游标归零，不修改大小页或其他候选区 | | PD 写出 | `write_to_pd_buffer` 将一个请求的 prefill 基准状态写入 `pd_page_index` 指定的状态页区域，按协议关联其前缀长度 | -| PD 读回 | `read_from_pd_buffer` 将完整快照恢复到指定请求槽位，不接收前缀长度或读取页内长度、版本字段；不自动生成历史 checkpoint | +| PD 读回 | `read_from_pd_buffer` 按 state_token_len 将完整快照恢复到指定请求槽位及对应 ring 位置，不读取页内长度、版本字段；不自动生成历史 checkpoint | CPU 尾页若需要保存 prefill 完成后的运行态，管理器先使用保留槽位调用 `save_bounded_att_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 @@ -1544,7 +1577,7 @@ CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字 #### 依赖、同步与使用示例 -资源方法接收请求槽位、checkpoint 槽位和明确的资源引用，只有 tail 无快照恢复额外接收前缀长度，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttResourceManager`单独协调 tail 在恢复边界上的重置或迁移。 +资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttResourceManager`单独协调 tail 在恢复边界上的重置或迁移。 GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 @@ -1553,14 +1586,18 @@ GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制 state.reset_bounded_att_state(req_indices) # backend 经层视图完成 prefill；调用方维护实际 chunk 终点。 state.save_bounded_att_pages( - req_indices, - page_kind="big", slot_ids=slot_ids, + req_indices=req_indices, + state_token_lens=state_token_lens, + page_kind="big", + slot_ids=slot_ids, ) # 确认所有参与状态已按指定长度保存并完成传输后，缓存调用方发布联合页。 state.restore_bounded_att_pages( - dst_req_indices, - page_kind="big", slot_ids=slot_ids, + req_indices=dst_req_indices, + state_token_lens=state_token_lens, + page_kind="big", + slot_ids=slot_ids, ) # target 请求在相关 KV、linear、tail 均就绪后继续执行。 @@ -1605,7 +1642,7 @@ class AttentionTailState(ABC): # 二、Attention tail：计算视图与无快照重置 # ======================================================================== @abstractmethod - def has_logical_layer(self, logical_layer_id: int) -> bool: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... def get_layer_view(self, logical_layer_id: int) -> Any: ... @@ -1636,7 +1673,7 @@ class AttentionTailState(ABC): | 接口 | 语义 | | --- | --- | | 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | -| `has_logical_layer(logical_layer_id)` | 判断当前 State 是否包含该逻辑层，不依赖实际存储分配，共享逻辑层同样返回 True | +| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | | `reset_attention_tail_state(req_indices)` | 清零指定请求的 tail，适用于新请求或可重置的缓存前缀；K-pool 清零完整 ring，其他实现清零对应数据与控制信息。模型配置和缓存流程负责保证边界，接口不接收长度 | diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/att_resources/att_resource_manager.py index a27a940e11..e449fc2a65 100644 --- a/lightllm/common/att_resources/att_resource_manager.py +++ b/lightllm/common/att_resources/att_resource_manager.py @@ -103,7 +103,7 @@ def get_att_state( and isinstance(state, (FullAttentionState, BoundedAttentionState, AttentionTailState)) and att_type in state.get_state_type() and (state_class is None or isinstance(state, state_class)) - and state.has_logical_layer(logical_layer_id) + and state.contains_logical_layer(logical_layer_id) ] assert len(matches) == 1, ( f"expected exactly one attention State: {model_name=}, {att_type=}, {logical_layer_id=}, " @@ -387,28 +387,50 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: # Bounded attention 状态页：供 radix/CPU cache 复用；未启用对应页池时操作报错。 # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 - def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """将 prefill 运行态保存到全部参与 State 的同类页池、同一组槽位。""" + def save_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" bounded_states = self.get_bounded_attention_states() assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(slot_ids) == len(req_indices), "one slot is required per request" + assert len(slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" for state in bounded_states: - state.save_bounded_att_pages(req_indices, page_kind=page_kind, slot_ids=slot_ids) + state.save_bounded_att_pages( + req_indices=req_indices, + state_token_lens=state_token_lens, + page_kind=page_kind, + slot_ids=slot_ids, + ) - def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """将调用方选定的页快照恢复到对应请求槽位的运行态。""" + def restore_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" bounded_states = self.get_bounded_attention_states() assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(slot_ids) == len(req_indices), "one slot is required per request" + assert len(slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" assert not slot_ids or bounded_states, "state restore requires participating bounded states" for state in bounded_states: - state.restore_bounded_att_pages(req_indices, page_kind=page_kind, slot_ids=slot_ids) + state.restore_bounded_att_pages( + req_indices=req_indices, + state_token_lens=state_token_lens, + page_kind=page_kind, + slot_ids=slot_ids, + ) # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 def commit_bounded_att_state( @@ -548,10 +570,12 @@ def read_req_from_cpu( layouts: Mapping[str, StateTransferLayout], big_page_slot_ids: List[int] | None = None, req_idx: int | None = None, + state_token_len: int | None = None, ) -> None: """读回单个请求本次指定的有序前缀页面，并从末页恢复 req_idx 的 bounded 运行态。 big_page_slot_ids 只包含完整 CPU 页对应的普通槽位；末尾碎页使用内部预留的 LOAD 大页。 + 存在 bounded 时须提供 state_token_len，即恢复后的绝对前缀长度，不是本批 token 数。 存在 bounded 或 tail 时必须提供 req_idx；调用方保证最终前缀位于 tail 可重置边界。 """ full_states = self.get_full_attention_states() @@ -560,6 +584,7 @@ def read_req_from_cpu( if bounded_states or tail_states: assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" if bounded_states: + assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" # token 索引只用于判断本批末页是否完整，不读取 GPU 上的索引值。 page_tokens = layouts[bounded_states[0].resource_id].tokens_per_page has_partial_page = token_indices.numel() % page_tokens != 0 @@ -590,7 +615,12 @@ def read_req_from_cpu( # 阶段三：从末页恢复请求运行态；纯碎页也走 LOAD 槽位。 if bounded_states: - self.restore_bounded_att_pages([req_idx], page_kind="big", slot_ids=[load_big_page_slot_ids[-1]]) + self.restore_bounded_att_pages( + req_indices=[req_idx], + state_token_lens=[state_token_len], + page_kind="big", + slot_ids=[load_big_page_slot_ids[-1]], + ) # 阶段四：清零目标请求的 tail；调用方保证前缀终点满足重置边界。 for state in tail_states: @@ -707,7 +737,7 @@ def write_to_pd_buffer( req_idx: int | None = None, state_token_len: int | None = None, ) -> None: - """按任务布局写出 PD 页；含 tail 时传入 prefill 前缀长度 state_token_len。""" + """按任务布局写出 PD 页；att_state 必须传入已执行的绝对前缀长度 state_token_len。""" match kind: case "kv": # KV 页只写出 full-attention 数据；各 State 负责自身布局与容量约束。 @@ -721,11 +751,10 @@ def write_to_pd_buffer( ) case "att_state": - # 提交搬运前确认请求位置；tail 额外使用前缀长度定位逻辑窗口。 + # 前缀长度由 PD 任务提供，bounded/tail 据此定位各自的运行态窗口。 assert req_idx is not None, "state PD pages require a request index" tail_states = self.get_attention_tail_states() - if tail_states: - assert state_token_len is not None and state_token_len >= 0, "PD tail requires its prefill prefix length" + assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" # 阶段一：写出 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): @@ -733,6 +762,7 @@ def write_to_pd_buffer( req_idx=req_idx, pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, + state_token_len=state_token_len, layout=layouts[state.resource_id], ) @@ -759,7 +789,7 @@ def read_from_pd_buffer( req_idx: int | None = None, state_token_len: int | None = None, ) -> None: - """按布局恢复 PD 页；含 tail 时使用 state_token_len 适配本地 MTP ring。""" + """按布局恢复 PD 页；att_state 使用绝对前缀长度 state_token_len 恢复运行态位置。""" match kind: case "kv": # KV 页只恢复 full-attention 数据；各 State 负责自身布局与容量约束。 @@ -773,11 +803,10 @@ def read_from_pd_buffer( ) case "att_state": - # 提交搬运前确认请求位置；tail 额外使用前缀长度定位逻辑窗口。 + # 前缀长度由 PD 任务提供，bounded/tail 据此定位各自的运行态窗口。 assert req_idx is not None, "state PD pages require a request index" tail_states = self.get_attention_tail_states() - if tail_states: - assert state_token_len is not None and state_token_len >= 0, "PD tail requires its prefill prefix length" + assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" # 阶段一：恢复 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): @@ -785,6 +814,7 @@ def read_from_pd_buffer( pd_page_buffers=pd_page_buffers, req_idx=req_idx, pd_page_index=pd_page_index, + state_token_len=state_token_len, layout=layouts[state.resource_id], ) diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/att_resources/att_tail/base_state.py index abf90c9c42..2aa784e123 100644 --- a/lightllm/common/att_resources/att_tail/base_state.py +++ b/lightllm/common/att_resources/att_tail/base_state.py @@ -41,8 +41,8 @@ def get_state_type(self) -> str: # ======================================================================== @abstractmethod - def has_logical_layer(self, logical_layer_id: int) -> bool: - """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" ... @abstractmethod diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/att_resources/att_tail/kpool_state.py index 044a563583..fdb4330f0c 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/att_resources/att_tail/kpool_state.py @@ -84,8 +84,8 @@ def get_state_type(self) -> str: # 二、Attention tail：计算视图与无快照重置 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/att_resources/bounded_att/base_state.py index b529184bec..635b68edb8 100644 --- a/lightllm/common/att_resources/bounded_att/base_state.py +++ b/lightllm/common/att_resources/bounded_att/base_state.py @@ -56,8 +56,8 @@ def initialize_page_pools( # ======================================================================== @abstractmethod - def has_logical_layer(self, logical_layer_id: int) -> bool: - """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" ... @abstractmethod @@ -71,13 +71,33 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... @abstractmethod - def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """将 prefill 完成后的基准运行态保存到指定页槽位，前缀关联由调用方维护。""" + def save_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 + + 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 + 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 + """ ... @abstractmethod - def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """从调用方选定的页槽位恢复运行态和控制信息，保留源快照数据。""" + def restore_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 + + req_indices、state_token_lens、slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， + 不表示本批搬运的 token 数，调用方保证与所选快照的前缀一致。 + """ ... @abstractmethod @@ -168,9 +188,10 @@ def write_to_pd_buffer( req_idx: int, pd_page_buffers: torch.Tensor, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将指定槽位的 prefill 基准运行态写入 PD 页。 + """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 """ @@ -182,7 +203,8 @@ def read_from_pd_buffer( pd_page_buffers: torch.Tensor, req_idx: int, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + """按 state_token_len 恢复请求运行态及 ring 位置等控制信息，不写入大小页池。""" ... diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/att_resources/bounded_att/gdn_state.py index dd57007b6b..03f0a9360c 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/att_resources/bounded_att/gdn_state.py @@ -106,8 +106,8 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: @@ -125,10 +125,16 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" + def save_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(slot_ids), "request and page counts differ" + assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 @@ -139,10 +145,16 @@ def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, sl copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" + def restore_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(slot_ids), "request and page counts differ" + assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): @@ -171,7 +183,8 @@ def commit_bounded_att_state( def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - 起止位置相对本 State，不包含 MTP 候选区、游标或大小页池的其他槽位。 + region.shape 描述 CPU 页中的全局字段；axis_slices 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 """ return self._get_transfer_regions() @@ -225,9 +238,12 @@ def read_page_state_from_cpu( assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_pages 按长度恢复。 for copies in bounded_att_buffer_copies( - page_fields=page_fields, + page_fields=pool_page_fields, page_buffers=cpu_page_buffers, layout=layout, page_indices=cpu_page_indices, @@ -307,9 +323,10 @@ def write_to_pd_buffer( req_idx: int, pd_page_buffers: torch.Tensor, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" + """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) @@ -327,9 +344,10 @@ def read_from_pd_buffer( pd_page_buffers: torch.Tensor, req_idx: int, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/att_resources/bounded_att/kda_state.py index f5628878a9..f61d97fadc 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/att_resources/bounded_att/kda_state.py @@ -106,8 +106,8 @@ def initialize_page_pools( # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: @@ -125,10 +125,16 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """将 prefill 完成后的 Conv/SSM 基准状态直接保存到指定页池槽位。""" + def save_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(slot_ids), "request and page counts differ" + assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 @@ -139,10 +145,16 @@ def save_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, sl copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def restore_bounded_att_pages(self, req_indices: List[int], page_kind: PageKind, slot_ids: List[int]) -> None: - """从指定页池槽位恢复 Conv/SSM 到目标基准位置，并将游标归零。""" + def restore_bounded_att_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(slot_ids), "request and page counts differ" + assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, slot_ids): @@ -171,7 +183,8 @@ def commit_bounded_att_state( def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - 起止位置相对本 State，不包含 MTP 候选区、游标或大小页池的其他槽位。 + region.shape 描述 CPU 页中的全局字段；axis_slices 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 """ return self._get_transfer_regions() @@ -225,9 +238,12 @@ def read_page_state_from_cpu( assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_pages 按长度恢复。 for copies in bounded_att_buffer_copies( - page_fields=page_fields, + page_fields=pool_page_fields, page_buffers=cpu_page_buffers, layout=layout, page_indices=cpu_page_indices, @@ -307,9 +323,10 @@ def write_to_pd_buffer( req_idx: int, pd_page_buffers: torch.Tensor, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将 prefill 基准运行态写入 PD 页；调用方保证源 MTP 游标为 0。""" + """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) @@ -327,9 +344,10 @@ def read_from_pd_buffer( pd_page_buffers: torch.Tensor, req_idx: int, pd_page_index: int, + state_token_len: int, layout: StateTransferLayout, ) -> None: - """将 PD 页恢复到指定请求槽位并规范化控制信息，不写入大小页池。""" + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) diff --git a/lightllm/common/att_resources/bounded_att/transfer_utils.py b/lightllm/common/att_resources/bounded_att/transfer_utils.py index 7f9cb8bb9c..54a0117c73 100644 --- a/lightllm/common/att_resources/bounded_att/transfer_utils.py +++ b/lightllm/common/att_resources/bounded_att/transfer_utils.py @@ -4,7 +4,7 @@ import torch -from ..types import StateTransferLayout, TensorField, TransferRegion +from ..types import StateTransferLayout, TensorField # ======================================================================== @@ -21,6 +21,10 @@ def bounded_att_buffer_copies( """按传输区域收集所有页面的复制视图，每个返回批次只包含一种 shape。 page_fields 与 page_indices 一一对应，调用方已排除 ready 页。 + CPU 路径的 page_fields 来自本地大小页池，PD 路径来自请求运行态；均为本地 TP 分片。 + layout.regions 描述外部页中的全局字段，以及本地 TP 分片在该字段中的位置。 + axis_slices 的 src_start 固定指本地，dst_start 固定指外部页，读回时也不交换含义。 + CPU 页和大小页快照的窗口均按逻辑时间排列，此处只处理普通轴切片，不展开或恢复 ring。 Conv 的 Q/K/V 切片与 SSM 分别成批，不将不同字段或不同 TP 切片混入同一次拷贝。 page_buffers 的第一维为物理页；这里只创建视图，由调用方逐批提交 Triton 算子。 """ @@ -30,37 +34,33 @@ def bounded_att_buffer_copies( if write and not region.write_enabled: continue copies = [] - for fields, page_index in zip(page_fields, page_indices): - # 阶段一：按字节区间定位当前页的字段，构造外部存储视图。 - field = fields[region.field] + for local_fields, page_index in zip(page_fields, page_indices): + # 阶段一：region.field 选择本地 Conv/SSM；CPU 读回时这里就是目标大页槽位。 + local_field = local_fields[region.field] + local_tensor = local_field.tensor + + # 阶段二：region 的字节区间和 shape 定位外部页中的全局字段，再对齐轴顺序。 page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) - external = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) - # 阶段二：同一区域使用相同的 TP 切片，所有页面生成的视图 shape 一致。 - local_view, external_view = _copy_pair(field, external, region) + external_tensor = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + assert local_tensor.dtype == external_tensor.dtype, "local and external field dtypes must match" + external_tensor = external_tensor.permute(tuple(region.axes.index(axis) for axis in local_field.axes)) + + # 阶段三：region 的两侧切片定位同一个 TP 分片，不代表本次复制的源和目标。 + # 例如 rank 1 的本地 head[0:4] 对应外部页 head[4:8]，写出和读回使用相同映射。 + local_slices = [slice(None)] * local_tensor.ndim + external_slices = [slice(None)] * external_tensor.ndim + for axis_slice in region.axis_slices: + axis = local_field.axes.index(axis_slice.axis) + local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) + external_slices[axis] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + local_view = local_tensor[tuple(local_slices)] + external_view = external_tensor[tuple(external_slices)] assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" - copies.append((local_view, external_view) if write else (external_view, local_view)) + + # 阶段四：只在此确定搬运方向；CPU 读回为外部页 → 本地大页池，不直接恢复请求。 + if write: + copies.append((local_view, external_view)) + else: + copies.append((external_view, local_view)) batches.append(copies) return batches - - -# ======================================================================== -# 二、私有辅助：本类状态的复制视图 -# ======================================================================== - -def _copy_pair( - field: TensorField, - external: torch.Tensor, - region: TransferRegion, -) -> tuple[torch.Tensor, torch.Tensor]: - """按 bounded 状态的 channel/head 等映射取视图，不处理 token 裁剪。""" - assert field.tensor.dtype == external.dtype, "local and external field dtypes must match" - # 外部布局可采用不同轴顺序，先转换为本地字段的顺序；只创建视图。 - external = external.permute(tuple(region.axes.index(axis) for axis in field.axes)) - local_slices = [slice(None)] * field.tensor.ndim - external_slices = [slice(None)] * external.ndim - # 布局已明确每个分片的两侧起点和长度，其余轴完整保留。 - for axis_slice in region.axis_slices: - axis = field.axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) - return field.tensor[tuple(local_slices)], external[tuple(external_slices)] diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/att_resources/full_att/base_state.py index 941e7d65af..372d226de4 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/att_resources/full_att/base_state.py @@ -129,8 +129,8 @@ def get_token_bytes(self) -> int: # ======================================================================== @abstractmethod - def has_logical_layer(self, logical_layer_id: int) -> bool: - """判断当前 State 是否服务该逻辑层；仅查询层映射，不依赖存储分配。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" ... @abstractmethod diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/att_resources/full_att/gqa_state.py index 172133380f..6221800ac9 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/att_resources/full_att/gqa_state.py @@ -84,8 +84,8 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_storage_layer_index(self, logical_layer_id: int) -> int: diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/att_resources/full_att/group_quantized_state.py index 948440b310..58d432ee11 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/att_resources/full_att/group_quantized_state.py @@ -101,8 +101,8 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_storage_layer_index(self, logical_layer_id: int) -> int: diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/att_resources/full_att/mla_state.py index 90844abd9f..5395c9ca06 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/att_resources/full_att/mla_state.py @@ -72,8 +72,8 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_storage_layer_index(self, logical_layer_id: int) -> int: diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/att_resources/full_att/packed_mla_state.py index ee2f5d5502..20bfc19e20 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/att_resources/full_att/packed_mla_state.py @@ -126,8 +126,8 @@ def indexer_offset(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_storage_layer_index(self, logical_layer_id: int) -> int: diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/att_resources/full_att/static_fp8_state.py index 26bdbe3721..61398c4e7e 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/att_resources/full_att/static_fp8_state.py @@ -166,8 +166,8 @@ def get_token_bytes(self) -> int: # 二、Full attention：计算视图、KV 写入与完整 token 复制 # ======================================================================== - def has_logical_layer(self, logical_layer_id: int) -> bool: - """查询逻辑层归属；共享同一物理层的逻辑层均返回 True。""" + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer def get_storage_layer_index(self, logical_layer_id: int) -> int: diff --git a/lightllm/common/att_resources/types.py b/lightllm/common/att_resources/types.py index 8082fc70ca..b2eaf0d226 100644 --- a/lightllm/common/att_resources/types.py +++ b/lightllm/common/att_resources/types.py @@ -28,8 +28,8 @@ # "small" 对应 small_pages_pool 与 bounded_att_small_page_allocator。 # 两个域中的同一个 slot_id 指向不同状态页。调用方先在对应 allocator 分配槽位， # 再将类别传给 save_bounded_att_pages / restore_bounded_att_pages 等操作，例如： -# manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[3]) -# 上例假定请求槽位 0 的已接受长度为 128，且大页槽位 3 已由调用方分配。 +# manager.save_bounded_att_pages(req_indices=[0], state_token_lens=[256], page_kind="big", slot_ids=[3]) +# 上例假定请求槽位 0 的已接受长度为 256，且大页槽位 3 已由调用方分配。 PageKind = Literal["big", "small"] diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 1899f69f49..cf04e3be4f 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -448,7 +448,8 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): selected = [(req, slot) for req, slot in zip(reqs, big_page_buffer_ids) if slot != -1] self.att_resource_manager.save_bounded_att_pages( - [req.req_idx for req, _ in selected], + req_indices=[req.req_idx for req, _ in selected], + state_token_lens=[req.get_chuncked_input_token_len() for req, _ in selected], page_kind="big", slot_ids=[slot for _, slot in selected], ) @@ -468,7 +469,8 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR if req.tail_small_page_buffer_id is not None: dst_buffer_idx = req.tail_small_page_buffer_id self.att_resource_manager.save_bounded_att_pages( - [req.req_idx], + req_indices=[req.req_idx], + state_token_lens=[req.hybrid_cache_len], page_kind="small", slot_ids=[dst_buffer_idx], ) @@ -675,7 +677,12 @@ def _match_radix_cache(self): def _restore_attention_page(self, node, page_kind): resources = g_infer_context.att_resource_manager slot = node.big_page_buffer_idx if page_kind == "big" else node.small_page_buffer_idx - resources.restore_bounded_att_pages([self.req_idx], page_kind=page_kind, slot_ids=[slot]) + resources.restore_bounded_att_pages( + req_indices=[self.req_idx], + state_token_lens=[node.node_prefix_total_len], + page_kind=page_kind, + slot_ids=[slot], + ) resources.reset_attention_tail_state([self.req_idx]) def _hybrid_match_radix_cache(self): diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 728c002dd0..ae4edf83c7 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -184,6 +184,7 @@ def read_req_from_cpu( layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=big_page_slot_ids, req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, + state_token_len=page_lengths[-1], ) if big_page_slot_ids is not None: req.hybrid_len_to_big_page_id.update(zip(complete_page_lengths, big_page_slot_ids)) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index ea01afd357..0f88538c65 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -341,10 +341,20 @@ def start_runtime(): for i, state in enumerate(manager.get_bounded_attention_states()): state.conv_state.fill_(i + 3) state.ssm_state.fill_(i + 7) - manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + manager.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + slot_ids=[slot], + ) manager.reset_bounded_att_state([0]) assert all(not state.conv_state[:, 0].any() for state in manager.get_bounded_attention_states()) - manager.restore_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + manager.restore_bounded_att_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + slot_ids=[slot], + ) for i, state in enumerate(manager.get_bounded_attention_states()): assert (state.conv_state[:, 0, :, : state.config.conv_width] == i + 3).all() assert (state.ssm_state[:, 0] == i + 7).all() diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 4c4290449f..7a9ffe5aad 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -164,7 +164,12 @@ def test_prefill_state_is_saved_and_restored( assert type(state.small_pages_pool) is pool_cls assert state.big_pages_pool is not state.small_pages_pool assert state.mtp_state_index.tolist() == [0, 0, 1] - state.save_bounded_att_pages([0], page_kind=page_kind, slot_ids=[1]) + state.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + slot_ids=[1], + ) if device == "cuda": torch.cuda.synchronize() pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool @@ -172,7 +177,12 @@ def test_prefill_state_is_saved_and_restored( assert torch.equal(pool.conv_state[1, ..., 1], torch.full_like(pool.conv_state[1, ..., 1], 11)) assert (pool.ssm_state[1] == 20).all() state.mtp_state_index[1] = 2 - state.restore_bounded_att_pages([1], page_kind=page_kind, slot_ids=[1]) + state.restore_bounded_att_pages( + req_indices=[1], + state_token_lens=[8], + page_kind=page_kind, + slot_ids=[1], + ) assert state.mtp_state_index[1].item() == 0 assert torch.equal(state.conv_state[:, 1, :, :2].cpu(), pool.conv_state[1]) assert torch.equal(state.ssm_state[:, 3].cpu(), pool.ssm_state[1]) @@ -200,10 +210,20 @@ def test_gdn_dimensions_and_page_restore(): state.conv_state[:, 0, :, :2].copy_(conv) state.ssm_state[:, 0].copy_(ssm) for page_kind in ("big", "small"): - state.save_bounded_att_pages([0], page_kind=page_kind, slot_ids=[0]) + state.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + slot_ids=[0], + ) state.reset_bounded_att_state([0]) for req, page_kind in enumerate(("big", "small"), start=1): - state.restore_bounded_att_pages([req], page_kind=page_kind, slot_ids=[0]) + state.restore_bounded_att_pages( + req_indices=[req], + state_token_lens=[8], + page_kind=page_kind, + slot_ids=[0], + ) assert torch.equal(state.conv_state[:, req, :, :2].cpu(), conv) assert torch.equal(state.ssm_state[:, req * 3].cpu(), ssm) with pytest.raises(AssertionError, match="equal Q/K/V"): @@ -215,7 +235,12 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): state = make_state(cls, device="cuda") state.conv_state[:, 0].fill_(7) state.ssm_state[:, 0:3].fill_(11) - state.save_bounded_att_pages([0], page_kind="small", slot_ids=[0]) + state.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[6], + page_kind="small", + slot_ids=[0], + ) # 推进候选区和游标后，重置应清零全部运行态，同时保留已保存的页快照。 state.conv_state[:, 0, :, 2:4].fill_(19) state.ssm_state[:, 2].fill_(23) @@ -234,7 +259,12 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): source = make_state(cls, device="cuda") source.conv_state[:, 0].fill_(13) source.ssm_state[:, 0].fill_(29) - source.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) + source.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=[0], + ) layout = page_layout(source) buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) source.write_page_state_to_cpu( @@ -253,7 +283,12 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): big_page_slot_ids=[1], layout=layout, ) - target.restore_bounded_att_pages([2], page_kind="big", slot_ids=[1]) + target.restore_bounded_att_pages( + req_indices=[2], + state_token_lens=[8], + page_kind="big", + slot_ids=[1], + ) assert (target.conv_state[:, 2, :, :2] == 13).all() assert (target.ssm_state[:, 6] == 29).all() # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 @@ -300,6 +335,7 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( req_idx=0, pd_page_buffers=allocation, pd_page_index=1, + state_token_len=8, layout=layout, ) target = make_state(cls, device=device, mtp_step=1) @@ -307,6 +343,7 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( pd_page_buffers=allocation, req_idx=2, pd_page_index=1, + state_token_len=8, layout=layout, ) assert (target.conv_state[:, 2, :, :2] == 23).all() diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index a8b2d31039..a8d42164fd 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -821,7 +821,12 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas state.conv_state.fill_(i + 3) state.ssm_state.fill_(i + 7) source_slots = manager.bounded_att_big_page_allocator.alloc(1) - manager.save_bounded_att_pages([0], page_kind="big", slot_ids=source_slots) + manager.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=source_slots, + ) pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) manager.write_req_to_cpu( token_indices=torch.tensor([], dtype=torch.long), @@ -839,6 +844,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=torch.tensor([], dtype=torch.long), + state_token_len=8, layouts=layouts, big_page_slot_ids=target_slots, req_idx=0, diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 543ca01046..116fbd05d8 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -34,8 +34,8 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - token_capacity=8, device="cpu", ) - assert all(state.has_logical_layer(logical_layer_id) for logical_layer_id in (0, 4, 8, 11)) - assert not state.has_logical_layer(1) + assert all(state.contains_logical_layer(logical_layer_id) for logical_layer_id in (0, 4, 8, 11)) + assert not state.contains_logical_layer(1) assert all(buffer.shape[0] == 2 for buffer in full_state_test_buffers(state).values()) assert state.get_token_bytes() == template.get_token_bytes() assert allocated_bytes(state, token_only=True) == state.get_token_bytes() * (8 + state.page_size) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 35c266c7a1..94c089bf32 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -139,8 +139,8 @@ def test_bounded_state_layer_ownership(state_class: type[KDAState] | type[GDNSta manager = AttResourceManager(request_capacity=3) manager.set_att_state("main", state) manager.finish_state_creation() - assert state.has_logical_layer(1) and state.has_logical_layer(3) - assert not state.has_logical_layer(0) + assert state.contains_logical_layer(1) and state.contains_logical_layer(3) + assert not state.contains_logical_layer(0) assert manager.get_att_state("main", state.get_state_type(), 3) is state @@ -156,8 +156,8 @@ def test_mla_selection_excludes_same_layer_tail() -> None: manager.set_att_state("main", state) manager.set_att_state("main", tail) manager.finish_state_creation() - assert tail.has_logical_layer(3) and tail.has_logical_layer(7) - assert not tail.has_logical_layer(0) + assert tail.contains_logical_layer(3) and tail.contains_logical_layer(7) + assert not tail.contains_logical_layer(0) assert manager.get_att_state("main", "mla", 7) is state assert manager.get_att_state("main", "kpool", 7) is tail assert manager.get_att_state( diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 72a2e794c9..04aa3124bc 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -12,6 +12,7 @@ from lightllm.common.att_resources.allocators import FullAttKVAllocator, StatePageAllocator from lightllm.common.att_resources import AttResourceManager +from lightllm.common.att_resources.types import PageKind, StateTransferLayout from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager from lightllm.common.att_resources.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state, layout_for @@ -66,7 +67,12 @@ def init_source(manager: AttResourceManager): state.tail_buffer[:, 0].fill_(3) state.tail_buffer[:, 1].fill_(9) slots = manager.bounded_att_big_page_allocator.alloc(2) - manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) + manager.save_bounded_att_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + slot_ids=slots, + ) if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: torch.cuda.synchronize() return slots @@ -76,7 +82,12 @@ def init_source(manager: AttResourceManager): def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.MonkeyPatch): manager = composition(device="cuda") slots = init_source(manager) - manager.restore_bounded_att_pages([2], page_kind="big", slot_ids=[slots[0]]) + manager.restore_bounded_att_pages( + req_indices=[2], + state_token_lens=[4], + page_kind="big", + slot_ids=[slots[0]], + ) for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): assert (state.conv_state[:, 2, :, :2] == 11 + i).all() assert (state.ssm_state[:, 6] == 31 + i).all() @@ -86,7 +97,12 @@ def fail(*args: object, **kwargs: object): monkeypatch.setattr(manager.draft0_gdn_conv_fp32_ssm_fp32_state, "save_bounded_att_pages", fail) with pytest.raises(RuntimeError): - manager.save_bounded_att_pages([0, 1], page_kind="big", slot_ids=slots) + manager.save_bounded_att_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + slot_ids=slots, + ) assert manager.bounded_att_big_page_allocator.free_count == 1 manager.bounded_att_big_page_allocator.free(slots) assert manager.bounded_att_big_page_allocator.free_count == 3 @@ -100,13 +116,103 @@ def test_disabled_domains(): manager.reset_bounded_att_state([0]) manager.reset_attention_tail_state([0]) with pytest.raises(AssertionError, match="disabled"): - manager.save_bounded_att_pages([0], page_kind="big", slot_ids=[0]) + manager.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + slot_ids=[0], + ) empty_manager = AttResourceManager(request_capacity=3) empty_manager.finish_state_creation() with pytest.raises(AssertionError, match="no full-attention"): empty_manager.copy_full_att_tokens(torch.tensor([0], dtype=torch.long), torch.tensor([1], dtype=torch.long)) +@pytest.mark.parametrize("operation", ["save_bounded_att_pages", "restore_bounded_att_pages"]) +def test_page_operations_forward_each_request_length(operation: str, monkeypatch: pytest.MonkeyPatch): + manager = composition(with_tail=False) + calls = [] + + def record_lengths( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + slot_ids: list[int], + ) -> None: + calls.append((req_indices, state_token_lens, page_kind, slot_ids)) + + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, operation, record_lengths) + getattr(manager, operation)( + req_indices=[0, 2], + state_token_lens=[17, 29], + page_kind="big", + slot_ids=[0, 1], + ) + assert calls == [([0, 2], [17, 29], "big", [0, 1])] * len(manager.get_bounded_attention_states()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for token indices") +def test_cpu_restore_forwards_absolute_length_for_partial_prefix(monkeypatch: pytest.MonkeyPatch): + manager = composition(with_tail=False) + layouts = layouts_for(manager) + manager.bounded_att_cpu_load_big_page_slot = 2 + for state in manager.get_full_attention_states(): + monkeypatch.setattr(state, "read_from_cpu", lambda **kwargs: None) + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, "read_page_state_from_cpu", lambda **kwargs: None) + restored = [] + + def record_restore( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + slot_ids: list[int], + ) -> None: + restored.append((req_indices, state_token_lens, page_kind, slot_ids)) + + monkeypatch.setattr(manager, "restore_bounded_att_pages", record_restore) + # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 + manager.read_req_from_cpu( + cpu_page_buffers=torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), + cpu_page_indices=[0, 1], + token_indices=torch.arange(6, device="cuda"), + layouts=layouts, + big_page_slot_ids=[0], + req_idx=1, + state_token_len=70, + ) + assert restored == [([1], [70], "big", [2])] + + +@pytest.mark.parametrize("operation", ["write_to_pd_buffer", "read_from_pd_buffer"]) +def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): + manager = composition(with_tail=False) + layouts = layouts_for(manager, kind="att_state") + lengths = [] + + def record_length( + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + lengths.append((req_idx, state_token_len)) + + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, operation, record_length) + getattr(manager, operation)( + pd_page_buffers=torch.empty((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), + pd_page_index=0, + kind="att_state", + layouts=layouts, + req_idx=2, + state_token_len=29, + ) + assert lengths == [(2, 29)] * len(manager.get_bounded_attention_states()) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_reset_tail_clears_all_tail_states_and_preserves_linear(): first = make_tail() @@ -180,6 +286,7 @@ def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | tor dtype=torch.long, device=device, ), + state_token_len=8, layouts=layouts, big_page_slot_ids=target_slots, req_idx=2, @@ -214,7 +321,12 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i small_slots = [] if token_count % 4: small_slots = source.bounded_att_small_page_allocator.alloc(1) - source.save_bounded_att_pages([page_count - 1], page_kind="small", slot_ids=small_slots) + source.save_bounded_att_pages( + req_indices=[page_count - 1], + state_token_lens=[token_count], + page_kind="small", + slot_ids=small_slots, + ) torch.cuda.synchronize() source.write_req_to_cpu( token_indices=torch.arange(token_count, device="cuda"), @@ -237,6 +349,7 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i cpu_page_buffers=pages, cpu_page_indices=page_ids, token_indices=torch.arange(4, 4 + token_count, device="cuda"), + state_token_len=token_count, layouts=layouts, big_page_slot_ids=slots, req_idx=2, @@ -284,6 +397,7 @@ def test_cpu_restore_resets_selected_tail_after_all_pages(tail_pool_size: int): cpu_page_buffers=pages, cpu_page_indices=[0, 1], token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + state_token_len=8, layouts=layouts, big_page_slot_ids=target_slots, req_idx=0, @@ -448,7 +562,12 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo source = composition("cuda", with_tail=False) big_slots = init_source(source) small_slots = source.bounded_att_small_page_allocator.alloc(1) - source.save_bounded_att_pages([1], page_kind="small", slot_ids=small_slots) + source.save_bounded_att_pages( + req_indices=[1], + state_token_lens=[6], + page_kind="small", + slot_ids=small_slots, + ) torch.cuda.synchronize() layouts = layouts_for(source) page_ids = [2, 0, 3] diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index e37db5d1c8..68de6067a5 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -137,7 +137,12 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( for h in range(cfg.num_v_heads): state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_v_heads + h + 7) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + resources.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=[slot], + ) torch.cuda.synchronize() expected.append((state.big_pages_pool.conv_state[slot].clone(), state.big_pages_pool.ssm_state[slot].clone())) size = sources[0][0].get_cpu_page_bytes() @@ -164,6 +169,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( dtype=torch.long, device="cuda", ), + state_token_len=8, layouts=resources.get_cpu_transfer_layouts(), big_page_slot_ids=[slot], req_idx=1, diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 8fcb74d304..92bc2eeaed 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -95,7 +95,12 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: assert not tail[:, 2].any() # DP fetch resets the aligned prefix. slot = source.bounded_att_big_page_allocator.alloc(1)[0] - source.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + source.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=[slot], + ) torch.cuda.synchronize() cpu = torch.empty((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) source.write_req_to_cpu( @@ -116,6 +121,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: dtype=torch.long, device="cuda", ), + state_token_len=8, layouts=target.get_cpu_transfer_layouts(), big_page_slot_ids=[dst_slot], req_idx=1, diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 0279153752..485822701c 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -325,7 +325,12 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( pool.ssm_state.zero_() slot = resources.bounded_att_big_page_allocator.alloc(1) with torch.cuda.device(source_device): - resources.save_bounded_att_pages([0], page_kind="big", slot_ids=slot) + resources.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=slot, + ) mapping = torch.zeros((3, 16), dtype=torch.int32, device=f"cuda:{source_device}") mapping[1, 0] = 7 torch.cuda.synchronize(source_device) @@ -414,6 +419,7 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe token_indices=torch.tensor([0, 3], dtype=torch.long, device="cuda"), task_kind=kind, req_idx=1, + state_token_len=8, write_to_pd_buffer=True, ) copy_page( @@ -425,6 +431,7 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe token_indices=torch.tensor([4, 7], dtype=torch.long, device="cuda"), task_kind=kind, req_idx=2, + state_token_len=8, write_to_pd_buffer=False, ) for rank, (manager, full) in enumerate(target): diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 720c18909a..0875d3f9b5 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -143,13 +143,23 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() - resources.save_bounded_att_pages([0], page_kind=kind, slot_ids=[slot]) + resources.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=kind, + slot_ids=[slot], + ) torch.cuda.synchronize() linear.conv_state[:, 0].zero_() linear.ssm_state[:, 0].zero_() tail.tail_buffer.normal_() unchanged = tail.tail_buffer[:, [0, 1, 3]].clone() - resources.restore_bounded_att_pages([2], page_kind=kind, slot_ids=[slot]) + resources.restore_bounded_att_pages( + req_indices=[2], + state_token_lens=[8], + page_kind=kind, + slot_ids=[slot], + ) resources.reset_attention_tail_state([2]) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index 5bd3451fb9..6d42fdd052 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -268,15 +268,30 @@ def test_flush_cache_requires_released_references_and_returns_allocator_pages(): assert tree.root_node.ref_counter == 1 -def test_hybrid_flush_releases_shared_kv_and_state_slots(): +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_hybrid_flush_releases_shared_kv_and_state_slots(monkeypatch): from unit_tests.common.att_resources.test_manager_helpers import composition - resources = composition(with_tail=False) + monkeypatch.setattr( + "lightllm.common.att_resources.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda", with_tail=False) allocator = resources.full_att_allocator big = resources.bounded_att_big_page_allocator.alloc(1)[0] small = resources.bounded_att_small_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[big]) - resources.save_bounded_att_pages([1], page_kind="small", slot_ids=[small]) + resources.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=[big], + ) + resources.save_bounded_att_pages( + req_indices=[1], + state_token_lens=[12], + page_kind="small", + slot_ids=[small], + ) values = allocator.alloc(12) tree = HybridAttPagedRadixCache(12, 106, 4, 2, resources) tree.insert( diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index fa5686578d..ed5fe0585b 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -269,7 +269,12 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( linear.conv_state.fill_(3) linear.ssm_state.fill_(5) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages([0], page_kind="big", slot_ids=[slot]) + resources.save_bounded_att_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + slot_ids=[slot], + ) kwargs = dict(big_page_slot_ids=[slot]) torch.cuda.synchronize() expected.append( @@ -307,6 +312,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( dtype=torch.long, device="cuda", ), + state_token_len=8, layouts=resources.get_cpu_transfer_layouts(), **kwargs, ) @@ -434,6 +440,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=[pages[1]], token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + state_token_len=8, layouts=resources.get_cpu_transfer_layouts(), **kwargs, ) diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 0284054b3f..904cad0f17 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -33,16 +33,26 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): @pytest.mark.usefixtures("kv_allocator_environment") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from unit_tests.common.att_resources.test_manager_helpers import composition, init_source - resources = composition() + monkeypatch.setattr( + "lightllm.common.att_resources.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda") slots = init_source(resources) tail = resources.main_kpool_c4_fp32_state resources.reset_bounded_att_state([1]) assert (tail.tail_buffer[:, 1] == 9).all() - resources.restore_bounded_att_pages([1], page_kind="big", slot_ids=[slots[-1]]) + resources.restore_bounded_att_pages( + req_indices=[1], + state_token_lens=[8], + page_kind="big", + slot_ids=[slots[-1]], + ) for state in resources.get_bounded_attention_states(): state.mtp_state_index.fill_(2) @@ -129,9 +139,9 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( context.att_resource_manager = SimpleNamespace( bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: [next(slots)]), bounded_att_small_page_allocator=SimpleNamespace(alloc=lambda count: [7]), - save_bounded_att_pages=lambda reqs, page_kind, slot_ids: ( + save_bounded_att_pages=lambda req_indices, state_token_lens, page_kind, slot_ids: ( big_snapshots if page_kind == "big" else small_snapshots - ).append(req.get_chuncked_input_token_len()), + ).extend(state_token_lens), ) while req.cur_kv_len < req.get_cur_total_len(): From 1a362139bc3b9d446de9704bb3ceb452212d773f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 06:52:49 +0000 Subject: [PATCH 24/56] refactor: unify attention state management and transfer naming --- docs/CN/source/models/add_new_model.md | 4 +- docs/EN/source/models/add_new_model.md | 4 +- docs/cache_manager_design.md | 460 +++++++++--------- lightllm/common/att_resources/__init__.py | 5 - lightllm/common/attention_state/__init__.py | 5 + .../allocators/__init__.py | 0 .../allocators/full_att_kv_allocator.py | 0 .../allocators/state_page_allocator.py | 14 +- .../att_tail/__init__.py | 0 .../att_tail/base_state.py | 8 +- .../att_tail/create_kpool_state.py | 0 .../att_tail/kpool_state.py | 30 +- .../att_tail/kpool_transfer.py | 52 +- .../bounded_att/__init__.py | 0 .../bounded_att/base_state.py | 34 +- .../bounded_att/conv_ssm_types.py | 0 .../bounded_att/copy_kernel.py | 0 .../bounded_att/create_gdn_state.py | 0 .../bounded_att/create_kda_state.py | 0 .../bounded_att/gdn_pages_pool.py | 0 .../bounded_att/gdn_state.py | 110 ++--- .../bounded_att/kda_pages_pool.py | 0 .../bounded_att/kda_state.py | 110 ++--- .../bounded_att/transfer_utils.py | 10 +- .../bounded_att/utils.py | 0 .../full_att/__init__.py | 0 .../full_att/base_state.py | 16 +- .../full_att/create_gqa_state.py | 0 .../full_att/create_mla_state.py | 0 .../full_att/dp_transfer.py | 0 .../full_att/gqa_state.py | 34 +- .../full_att/gqa_transfer.py | 10 +- .../full_att/group_quantized_state.py | 34 +- .../full_att/group_quantized_transfer.py | 10 +- .../full_att/mla_state.py | 34 +- .../full_att/mla_transfer.py | 10 +- .../full_att/packed_mla_state.py | 34 +- .../full_att/packed_mla_transfer.py | 10 +- .../full_att/static_fp8_state.py | 36 +- .../full_att/static_fp8_transfer.py | 10 +- .../manager.py} | 234 ++++----- .../types.py | 68 +-- .../utils.py | 0 .../common/basemodel/attention/fa3/fp8.py | 16 +- .../basemodel/attention/flashinfer/fp8.py | 8 +- .../common/basemodel/attention/linear/gdn.py | 14 +- .../common/basemodel/attention/linear/kda.py | 14 +- lightllm/common/basemodel/basemodel.py | 54 +- lightllm/common/basemodel/cuda_graph.py | 4 +- lightllm/common/basemodel/infer_struct.py | 6 +- .../transformer_layer_infer_template.py | 8 +- .../common/basemodel/prefill_cuda_graph.py | 4 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/models/bloom/model.py | 4 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/models/deepseek2/model.py | 4 +- .../layer_infer/transformer_layer_infer.py | 20 +- lightllm/models/deepseek3_2/model.py | 4 +- lightllm/models/gemma3/model.py | 4 +- .../layer_infer/transformer_layer_infer.py | 8 +- lightllm/models/gemma4/model.py | 4 +- lightllm/models/gemma_2b/model.py | 4 +- lightllm/models/glm5_next/indexer.py | 16 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/models/glm5_next/model.py | 12 +- lightllm/models/glm5_next_mtp/model.py | 8 +- .../layer_infer/transformer_layer_infer.py | 14 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/models/llama/model.py | 4 +- lightllm/models/mistral/model.py | 4 +- lightllm/models/mixtral/model.py | 4 +- lightllm/models/qwen3_dflash/model.py | 2 +- .../layer_infer/transformer_layer_infer.py | 14 +- lightllm/models/qwen3next/model.py | 12 +- lightllm/models/starcoder/model.py | 4 +- lightllm/models/starcoder2/model.py | 4 +- lightllm/server/pd_io_struct.py | 2 +- .../dynamic_prompt/hybrid_att_radix_cache.py | 22 +- .../router/dynamic_prompt/radix_cache.py | 8 +- lightllm/server/router/manager.py | 2 +- .../server/router/model_infer/infer_batch.py | 48 +- .../model_infer/mode_backend/base_backend.py | 46 +- .../mode_backend/diverse_backend/impl.py | 2 +- .../dp_backend/dp_shared_kv_trans.py | 32 +- .../mode_backend/multi_level_kv_cache.py | 12 +- .../pd/decode_node_impl/decode_impl.py | 2 +- .../decode_node_impl/decode_trans_process.py | 10 +- .../prefill_trans_process.py | 10 +- .../mode_backend/pd/resource_io.py | 26 +- .../model_infer/mtp_speculative/utils.py | 2 +- .../model_infer/shared_att_resources_utils.py | 6 +- lightllm/utils/cpu_cache_layout.py | 14 +- test/att_resources/validate_model_cache.py | 42 +- .../static_inference/static_benchmark.py | 22 +- test/cpu_cache_kernel/test_speed.py | 26 +- .../common/att_resources/manager_fixtures.py | 32 +- .../common/att_resources/model_helpers.py | 38 +- .../common/att_resources/test_allocators.py | 2 +- .../att_resources/test_backend_startup.py | 56 +-- .../att_resources/test_bounded_copy_kernel.py | 10 +- .../att_resources/test_bounded_state.py | 100 ++-- .../common/att_resources/test_capacity.py | 232 ++++----- .../att_resources/test_compute_binding.py | 14 +- .../att_resources/test_fp8_compute_binding.py | 26 +- .../att_resources/test_full_att_cpu_kernel.py | 18 +- .../test_full_att_dp_transfer.py | 8 +- .../att_resources/test_full_att_transfer.py | 20 +- .../common/att_resources/test_full_state.py | 46 +- .../test_full_storage_initialization.py | 12 +- .../att_resources/test_gqa_pd_kernel.py | 30 +- .../common/att_resources/test_gqa_transfer.py | 20 +- .../att_resources/test_kpool_transfer.py | 58 +-- .../common/att_resources/test_layer_groups.py | 22 +- .../test_layer_state_selection.py | 136 +++--- .../att_resources/test_manager_helpers.py | 108 ++-- .../att_resources/test_model_calibration.py | 38 +- .../test_model_initialization.py | 26 +- .../att_resources/test_model_resources.py | 20 +- .../test_model_storage_definition.py | 38 +- .../att_resources/test_mtp_state_commit.py | 22 +- .../att_resources/test_network_transport.py | 10 +- .../att_resources/test_no_attention_draft.py | 24 +- .../att_resources/test_page_pool_storage.py | 4 +- .../att_resources/test_paged_storage.py | 2 +- .../test_quantized_full_state.py | 38 +- .../test_quantized_paged_storage.py | 4 +- .../att_resources/test_role_transfers.py | 22 +- .../att_resources/test_runtime_transport.py | 54 +- .../att_resources/test_state_creation.py | 16 +- .../att_resources/test_state_identity.py | 108 ++-- .../att_resources/test_state_queries.py | 50 +- .../common/att_resources/test_tail_state.py | 28 +- .../common/att_resources/test_tp_config.py | 38 +- .../basemodel/attention/linear/test_gdn.py | 14 +- .../attention/test_fp8_logical_layer.py | 12 +- unit_tests/common/test_req_manager_page.py | 30 +- unit_tests/conftest.py | 6 +- unit_tests/models/glm5_next/test_cache.py | 14 +- unit_tests/models/glm5_next/test_kernels.py | 8 +- unit_tests/models/glm5_next/test_layers.py | 26 +- unit_tests/models/glm5_next/test_mtp.py | 14 +- unit_tests/models/glm5_next/test_pd_cache.py | 26 +- .../models/test_qwen3_dspark_model_output.py | 4 +- .../test_multimodal_cache_matching.py | 8 +- .../router/dynamic_prompt/test_radix_cache.py | 24 +- .../mode_backend/test_dp_big_pages.py | 14 +- .../mode_backend/test_multi_level_kv_cache.py | 6 +- .../test_paged_side_allocations.py | 10 +- .../mode_backend/test_pd_token_indices.py | 2 +- .../model_infer/test_cpu_resource_layout.py | 32 +- .../test_hybrid_cpu_state_cache.py | 16 +- .../model_infer/test_multimodal_chunking.py | 16 +- unit_tests/server/router/test_kv_capacity.py | 4 +- unit_tests/test_static_benchmark_paged_kv.py | 2 +- 154 files changed, 1947 insertions(+), 1925 deletions(-) delete mode 100644 lightllm/common/att_resources/__init__.py create mode 100644 lightllm/common/attention_state/__init__.py rename lightllm/common/{att_resources => attention_state}/allocators/__init__.py (100%) rename lightllm/common/{att_resources => attention_state}/allocators/full_att_kv_allocator.py (100%) rename lightllm/common/{att_resources => attention_state}/allocators/state_page_allocator.py (67%) rename lightllm/common/{att_resources => attention_state}/att_tail/__init__.py (100%) rename lightllm/common/{att_resources => attention_state}/att_tail/base_state.py (94%) rename lightllm/common/{att_resources => attention_state}/att_tail/create_kpool_state.py (100%) rename lightllm/common/{att_resources => attention_state}/att_tail/kpool_state.py (90%) rename lightllm/common/{att_resources => attention_state}/att_tail/kpool_transfer.py (85%) rename lightllm/common/{att_resources => attention_state}/bounded_att/__init__.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/base_state.py (89%) rename lightllm/common/{att_resources => attention_state}/bounded_att/conv_ssm_types.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/copy_kernel.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/create_gdn_state.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/create_kda_state.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/gdn_pages_pool.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/gdn_state.py (87%) rename lightllm/common/{att_resources => attention_state}/bounded_att/kda_pages_pool.py (100%) rename lightllm/common/{att_resources => attention_state}/bounded_att/kda_state.py (87%) rename lightllm/common/{att_resources => attention_state}/bounded_att/transfer_utils.py (87%) rename lightllm/common/{att_resources => attention_state}/bounded_att/utils.py (100%) rename lightllm/common/{att_resources => attention_state}/full_att/__init__.py (100%) rename lightllm/common/{att_resources => attention_state}/full_att/base_state.py (96%) rename lightllm/common/{att_resources => attention_state}/full_att/create_gqa_state.py (100%) rename lightllm/common/{att_resources => attention_state}/full_att/create_mla_state.py (100%) rename lightllm/common/{att_resources => attention_state}/full_att/dp_transfer.py (100%) rename lightllm/common/{att_resources => attention_state}/full_att/gqa_state.py (95%) rename lightllm/common/{att_resources => attention_state}/full_att/gqa_transfer.py (97%) rename lightllm/common/{att_resources => attention_state}/full_att/group_quantized_state.py (96%) rename lightllm/common/{att_resources => attention_state}/full_att/group_quantized_transfer.py (97%) rename lightllm/common/{att_resources => attention_state}/full_att/mla_state.py (94%) rename lightllm/common/{att_resources => attention_state}/full_att/mla_transfer.py (97%) rename lightllm/common/{att_resources => attention_state}/full_att/packed_mla_state.py (96%) rename lightllm/common/{att_resources => attention_state}/full_att/packed_mla_transfer.py (97%) rename lightllm/common/{att_resources => attention_state}/full_att/static_fp8_state.py (96%) rename lightllm/common/{att_resources => attention_state}/full_att/static_fp8_transfer.py (97%) rename lightllm/common/{att_resources/att_resource_manager.py => attention_state/manager.py} (84%) rename lightllm/common/{att_resources => attention_state}/types.py (82%) rename lightllm/common/{att_resources => attention_state}/utils.py (100%) diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 1f0412cddc..64cdddd139 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -128,9 +128,9 @@ class TpPartBaseModel: | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | -全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_att_resources()` 准备自身 State；随后管理器依次完成 `finish_state_creation()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 -公共 `prepare_att_resources()` 为 main 创建空管理器，draft 引用 main 的实例。三个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded/tail 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 +公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。三个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded/tail 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 ## 2. 添加 bloom 模型的示例说明 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index dc18d4da91..34838b80e1 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -128,9 +128,9 @@ Commonly used interfaces that need to be inherited and implemented | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | -After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_att_resources()`, then calls `finish_state_creation()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. +After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. -The public `prepare_att_resources()` stage creates an empty manager for main and shares it with drafts. The three State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded/tail also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. +The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The three State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded/tail also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. ## 2. Example of Adding Bloom Model diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index a13fac9c6a..0f90a09e0d 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1,15 +1,15 @@ # Attention Resources 架构设计草案 -本文定义 attention 资源的架构设计，包括 `AttResourceManager` 的职责、共享 allocator、各类 State 的接口、目录组织及 CPU/DP/PD 传输约定。文中的结构与接口为设计契约，待细化事项集中列于第 9 节。 +本文定义 attention 资源的架构设计，包括 `AttentionStateManager` 的职责、共享 allocator、各类 State 的接口、目录组织及 CPU/DP/PD 传输约定。文中的结构与接口为设计契约，待细化事项集中列于第 9 节。 ## 1. 设计约定 -统一在 `lightllm/common/att_resources/` 下组织可复用的资源实现，由 `AttResourceManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的三类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 +统一在 `lightllm/common/attention_state/` 下组织可复用的资源实现，由 `AttentionStateManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的三类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： ```python -class AttResourceManager: +class AttentionStateManager: full_att_allocator: FullAttKVAllocator | None bounded_att_big_page_allocator: StatePageAllocator | None bounded_att_small_page_allocator: StatePageAllocator | None @@ -20,15 +20,34 @@ class AttResourceManager: - 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att`、`main_att_tail`、`draft_att_tail` 成员。具体模型按需组织资源组合与实现类型，拥有的 State 成员统一遵守第 3.5 节的角色、类型和存储方案命名。 - `full_att`、`bounded_att`、`att_tail` 作为可复用实现的目录分类。full attention 保存历史 KV；bounded attention 由 `BoundedAttentionState` 抽象基类约定操作，当前包含 `KDAState`、`GDNState`，以后可加入 sliding-window 等有界状态，各自定义运行态与可选页池；attention tail 保存索引、池化、压缩等模块在增量计算中需要跨步骤保留的尾部数据及控制信息；具体实现决定恢复边界和快照要求。 - main/draft 可以采用不同方案；同一共享分配域的参与方使用同一个 allocator 实例，遵守统一的索引和回收约定。资源对象持有 buffer 和布局信息；管理器创建或绑定共享 allocator，调用方直接使用对应 allocator 分配与归还索引。 -- `AttResourceManager` 提供运行态清零、清零指定请求的 tail、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 +- `AttentionStateManager` 提供运行态清零、清零指定请求的 tail、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 - `req_to_token_indexs` 和请求槽位仍由 `ReqManager` 持有。资源管理器接收显式传入的物理索引、请求槽位及恢复长度，提供资源操作接口。 - 内部调用契约使用 `assert`，例如页类型、物理索引范围、Tensor 形状及共享分配域的一致性；用户配置、外部页面内容及格式兼容性继续显式校验，资源不足等运行失败保留相应异常。GPU 上的 MTP 状态选择保留异步断言，避免引入主机同步。 - 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 -本文区分已确定的职责边界和接口草案。main/draft 的资源由共享的 `AttResourceManager` 统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 +本文区分已确定的职责边界和接口草案。main/draft 的资源由共享的 `AttentionStateManager` 统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 ### 1.1 接口命名、排列与注释规范 +模块统一命名为 `attention_state`，公共管理器为 `AttentionStateManager`，模型和计算上下文通过 +`att_state_manager` 持有同一实例。管理器实现集中在 `attention_state/manager.py`，由包根目录导出。 + +| 命名 | 明确含义 | +| --- | --- | +| `register_state(model_role, state)` / `get_state(...)` | 按 `main`、`draft0..draft35` 注册或选择 State；`state_type_filter` 为算法/方案标记的包含匹配条件 | +| `finish_state_registration()` | 全部模型结束注册，允许缓存查询；不代表实际存储已分配 | +| `initialize_storage(...)` / `allocate_storage(...)` | manager 的前者规划容量并初始化存储，后者按已确定容量分配；FullAttentionState 直接提供 `allocate_storage(token_capacity)` | +| `save_bounded_att_runtime_to_pages` / `restore_bounded_att_runtime_from_pages` | 请求运行态与本地大小页快照之间的数据操作，不使用外部页 Region | +| `write_state_pages_to_cpu` / `read_state_pages_from_cpu` | 本地快照页与 CPU 联合页之间的数据操作，按布局处理字段位置和 TP 分片 | +| `write_to_pd_page` / `read_from_pd_page` | 操作 PD 页集合中 `pd_page_index` 指定的单个页面 | +| `get_cpu_page_layouts`、`get_pd_token_page_layouts`、`get_pd_runtime_page_layouts` | 分别返回 CPU 联合页、PD token 页和 PD 运行态页的布局 | +| `state_name`、`page_slot_ids`、`cpu_page_ready_flags` | 分别表示规范 State 成员名、本地状态页槽位列表和 CPU 页就绪标记 | + +管理器的 `write_req_to_cpu` / `read_req_from_cpu` 保留原名；绝对状态长度仍统一使用 +`state_token_len` / `state_token_lens`，不改为 prefix 命名。`AxisMapping` 的 +`local_start` / `external_start` 和 `local_circular` / `external_circular` 始终表示本地与外部两侧， +读回时仅反转数据方向；映射集合统一命名为 `axis_mappings`，单次覆盖为 `axis_mapping_overrides`。 + 新增 Python 文件、类和接口按 **公共函数在前、私有辅助函数在后** 排列。已有基础流程文件以保留功能、减少 review 差异为优先，保留原有函数相对顺序，不为满足分组规范整体重排。 新增和改动较多的文件不使用函数声明中的独立 `*` 分隔符。函数参数不少于 4 个时（不计 `self`、`cls`），声明每行一个参数；调用处不少于 4 个实参时，使用 `参数名=值` 并逐行排列。外部 API 的位置参数约束及 `*args`、`**kwargs` 保持其原有语义；改动较少的存量文件不作整文件格式调整。 @@ -45,7 +64,7 @@ class AttResourceManager: State 的基类与具体实现同样使用中文大块分隔注释，依次组织基础定义、该类资源的数据操作、CPU Cache、DP、PD、私有辅助实现。KDA/GDN 直接继承 `BoundedAttentionState`；所有具体 full-attention State 直接继承 `FullAttentionState`。具体 State 之间不继承实现，分别在自己的文件中按上述分块组织完整操作。 -- 分类数据操作统一带 `full_att`、`bounded_att` 或 `att_tail` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_pages`、`reset_attention_tail_state`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 +- 分类数据操作统一带 `full_att`、`bounded_att` 或 `att_tail` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_runtime_to_pages`、`reset_attention_tail_state`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 - `get_layer_view` 用于计算访问，CPU/DP/PD 保持按传输方向命名。State 不增加仅用于统一名字的转发方法。 - 所有 attention 配置字段必须提供中文注释，说明业务含义、单位或对应 Tensor 轴、全局/TP 本地范围，以及适用方案和默认值的特殊语义；派生属性还需说明计算关系。资源身份只保存在 State，不放入 config。 - 已知的调用参数使用显式 `key=value`，如 `size=...`、`dtype=...`、`device=...`、`pin_memory=...`；不先组装 `options` 字典再通过 `**options` 传入。 @@ -57,8 +76,7 @@ State 的基类与具体实现同样使用中文大块分隔注释，依次组 ```mermaid flowchart TD - Model[模型初始化] --> Concrete[模型专属 AttResourceManager] - Concrete -.->|实现公共接口| API[AttResourceManager] + Model[模型初始化] --> Concrete[共享 AttentionStateManager] Context[InferenceContext] --> Req[ReqManager] Req --> Slots[请求槽位与 KV 映射] Context -->|前缀复用| Radix[Radix Cache 功能] @@ -74,17 +92,17 @@ flowchart TD Transport[DP/PD 调度] -->|调用具体资源操作| Concrete ``` -图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享共享的 `AttResourceManager`，分别绑定各自的计算 State。 +图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享同一个 `AttentionStateManager`，分别绑定各自的计算 State。 | 层次 | 组件 | 职责 | | --- | --- | --- | | 请求流程 | `InferenceContext`、`ReqManager` | 创建、暂停、恢复和完成请求；由 ReqManager 管理请求槽位、KV 索引映射及采样参数 | -| 资源数据操作 | `AttResourceManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | -| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用`AttResourceManager`操作资源数据 | +| 资源数据操作 | `AttentionStateManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | +| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用`AttentionStateManager`操作资源数据 | | 资源实现 | 模型所需的 State、存储布局、checkpoint pool | 持有 buffer 和布局，实现读写、复制、重置和快照；IPC 导出范围由管理器统一处理 | | 传输调度 | DP/PD 与多级缓存模块 | 管理任务、队列、通信、事件同步及完成通知 | -`AttResourceManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，`AttResourceManager`执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 +`AttentionStateManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，`AttentionStateManager`执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 代码依赖从调用流程向资源层单向传递。radix cache 依赖管理器的通用状态页操作；管理器与 State 不导入或访问 radix 节点、树和策略。资源实现也不反向导入请求调度模块；传输调度使用数据读写与复制接口。 @@ -95,19 +113,19 @@ ReqManager ├── request_slots └── req_to_token_indexs # 共享请求到 KV 的索引映射 -AttResourceManager # 公共接口与 allocator 入口 +AttentionStateManager # 公共接口与 allocator 入口 ├── full_att_allocator: FullAttKVAllocator | None ├── bounded_att_big_page_allocator: StatePageAllocator | None └── bounded_att_small_page_allocator: StatePageAllocator | None -AttResourceManager # 模型实现，内部结构由其自行定义 +AttentionStateManager # 模型实现，内部结构由其自行定义 ├── 按需创建或绑定所需 allocator ├── 初始化模型所需的资源组合 ├── 实现内部状态页聚合操作与 CPU / DP / PD 数据传输 └── 绑定推理视图、组织传输所需状态 # 搬运直接调用 State 方法 ``` -纯 full-attention 与混合模型都使用 `AttResourceManager`，区别在于模型工厂返回的 State 组合。管理器按角色持有各 State，共享 allocator 按分配域创建；公共调用方使用同一组聚合接口。 +纯 full-attention 与混合模型都使用 `AttentionStateManager`，区别在于模型工厂返回的 State 组合。管理器按角色持有各 State，共享 allocator 按分配域创建；公共调用方使用同一组聚合接口。 以下列出可复用资源组件的结构示例。linear attention 直接使用 `KDAState`、`GDNState` 等具体状态对象，由 `BoundedAttentionState` 约定必需操作，不增加资源包装实例。管理器内部的成员遵守统一命名及唯一性约束，checkpoint pool 仅在相应功能启用时存在： @@ -137,25 +155,25 @@ AttentionTailState # 抽象操作契约，不固定内部存储 └── tail_buffer # 原始 key/gate ring，按层、请求组织 ``` -同一共享分配域中的 allocator 只初始化一次，由 main/draft 共用的唯一 `AttResourceManager` 管理生命周期。调用方解除相关引用后，向对应 allocator 单次归还索引，不因模型数增加而重复分配或释放。 +同一共享分配域中的 allocator 只初始化一次，由 main/draft 共用的唯一 `AttentionStateManager` 管理生命周期。调用方解除相关引用后，向对应 allocator 单次归还索引，不因模型数增加而重复分配或释放。 | 资源 | 持有者 | 分配单位 / 生命周期 | | --- | --- | --- | -| full-attention 空闲页与容量状态 | `AttResourceManager.full_att_allocator` | 每次为请求统一分配/回收物理页索引，main/draft 共用 | -| 大页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_big_page_allocator` | 为 main/draft 参与的大页 pool 统一分配槽位 ID | -| 小页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | +| full-attention 空闲页与容量状态 | `AttentionStateManager.full_att_allocator` | 每次为请求统一分配/回收物理页索引，main/draft 共用 | +| 大页 checkpoint 空闲槽位 | `AttentionStateManager.bounded_att_big_page_allocator` | 为 main/draft 参与的大页 pool 统一分配槽位 ID | +| 小页 checkpoint 空闲槽位 | `AttentionStateManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | | 请求 KV 索引表 | `ReqManager` | 记录逻辑 token 到共享物理索引的映射；负责绑定、更新和清理 | -| HOLD 保留物理页及索引描述 | `AttResourceManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | -| KV、量化 scale、packed indexer buffer 与层映射 | `AttResourceManager` 组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | -| KDA/GDN 的 Conv/SSM、MTP 游标 | `AttResourceManager`持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | +| HOLD 保留物理页及索引描述 | `AttentionStateManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | +| KV、量化 scale、packed indexer buffer 与层映射 | `AttentionStateManager` 组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | +| KDA/GDN 的 Conv/SSM、MTP 游标 | `AttentionStateManager`持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | | 大小页 checkpoint buffer | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID 保存数据；不持有独立空闲列表或长度记录 | -| indexer 原始尾部窗口 | `AttResourceManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | +| indexer 原始尾部窗口 | `AttentionStateManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | | 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | | 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复关联资源数据，解除引用后直接归还 allocator | ### 3.1 main/draft 共享 full-attention allocator -调用方直接通过 `AttResourceManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，`AttResourceManager`恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： +调用方直接通过 `AttentionStateManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，`AttentionStateManager`恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： ```text 逻辑 token 位置 t @@ -169,13 +187,13 @@ AttentionTailState # 抽象操作契约，不固定内部存储 共享索引表示存储位置的对应关系，不表示 main/draft 在每个位置都已经写入有效 KV。MTP 预留区和回退后的有效长度仍由请求执行流程维护。 -释放时，请求流程与`AttResourceManager`确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 +释放时，请求流程与`AttentionStateManager`确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 ### 3.2 main/draft 共享 linear checkpoint allocator main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAState`，draft 使用 `GDNState` 或无需 linear 状态。各状态对象持有自身运行态及启用的大页或小页 pool。对同一前缀的一个大页或小页 checkpoint，调用方直接从管理器持有的对应 allocator 分配一次槽位 ID，再将该 ID 传给管理器的聚合数据操作，供参与的两侧 pool 使用。两侧的状态布局、dtype 和单槽位字节数可以不同，各自的存储必须覆盖共享 allocator 可返回的槽位范围。 -例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。`AttResourceManager`负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 +例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。`AttentionStateManager`负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 @@ -183,22 +201,22 @@ linear 运行态与两侧 indexer tail 均由 `ReqManager` 提供的 `req_idx` ### 3.3 分配状态与数据状态分离 -allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；`AttResourceManager` 组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 +allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；`AttentionStateManager` 组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 -对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由`AttResourceManager`维护。 +对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由`AttentionStateManager`维护。 ### 3.4 按具体 attention 实现定义状态 -定义只约束操作的 `BoundedAttentionState` 抽象基类，由 `KDAState`、`GDNState` 等子类负责自己的成员、布局、重置、快照与恢复语义。`AttResourceManager` 初始化并组合所需对象，公共缓存流程通过管理器访问它们。不同模型也可以复用同一种状态实现。 +定义只约束操作的 `BoundedAttentionState` 抽象基类，由 `KDAState`、`GDNState` 等子类负责自己的成员、布局、重置、快照与恢复语义。`AttentionStateManager` 初始化并组合所需对象，公共缓存流程通过管理器访问它们。不同模型也可以复用同一种状态实现。 `KDAState` 的成员定义如下： ```python class KDAState(BoundedAttentionState): config: ConvSSMStateConfig # 仅包含算法布局和 TP 信息 - resource_id: str | None # 构造时为 None，由管理器绑定模型角色时设置 + state_name: str | None # 构造时为 None，由管理器绑定模型角色时设置 request_capacity: int # 包含 HOLD 请求槽位 conv_state: torch.Tensor ssm_state: torch.Tensor @@ -219,21 +237,21 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 `GDNState` 按 GDN 算法独立定义 Conv/SSM 和 MTP 状态，其大小页使用 `GDNPagesPool`，支持不同的 K/V head 数及维度。`KDAPagesPool` 按 KDA 的相同 Q/K/V head 数和方阵递推状态定义布局。两类 pool 各自分配快照存储，不额外记录长度或页面有效性；前缀与槽位的关联由缓存调用流程维护。两种 State 可以复用兼容的 Conv/SSM 数据操作，共同实现 `BoundedAttentionState` 接口；基类不规定 Conv/SSM 等成员，也不意味着两种算法的状态内容可以互换。 -运行态按 `ReqManager` 的槽位容量预分配，状态对象接收 `req_idx` 执行操作。main/draft 可以持有同一种状态类的两个独立实例，也可以使用不同状态类型；各实例的层数、维度、dtype 和 MTP 配置由对应模型决定。allocator 仍由 `AttResourceManager` 创建或绑定并统一管理。 +运行态按 `ReqManager` 的槽位容量预分配，状态对象接收 `req_idx` 执行操作。main/draft 可以持有同一种状态类的两个独立实例，也可以使用不同状态类型；各实例的层数、维度、dtype 和 MTP 配置由对应模型决定。allocator 仍由 `AttentionStateManager` 创建或绑定并统一管理。 后续 sliding-window attention 可以定义实现 `BoundedAttentionState` 的 `SlidingWindowState`，按窗口 KV、有效范围和游标等实际需求组织成员及自己的页池。它可以使用适用的 bounded-attention 分配域，其页布局、恢复边界和已接受状态选择须按窗口语义实现；不要求包含 Conv/SSM 字段，也不复用 KDA/GDN 的页池。统一操作接口见 6.7 节。 ### 3.5 State 成员命名、资源身份与唯一性 -`AttResourceManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `set_att_state(model_name, state)` 绑定。`model_name` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.resource_id`。`FullAttentionConfig`（含量化与 MLA 子类）、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `resource_id` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_creation()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 +`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。`FullAttentionConfig`（含量化与 MLA 子类）、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 -管理器成员名与 `state.resource_id` 完全一致，统一使用 `{model_name}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState`、`AttentionTailState` 类型表达，名称中不再重复 `full_att`、`bounded_att`、`att_tail`。 +管理器成员名与 `state.state_name` 完全一致，统一使用 `{model_role}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState`、`AttentionTailState` 类型表达，名称中不再重复 `full_att`、`bounded_att`、`att_tail`。 三类 State 都必须实现 `get_state_type()`，只返回中间的算法及方案标记，使用小写 snake_case。方案标记区分量化位宽、group size、scale 方式/精度，KDA/GDN 分别标注 Conv 和 SSM 精度，K-pool 标注压缩比例与存储精度。层号、容量、设备、TP rank 和请求信息不进入名称。 同一角色、同一标记只能注册一个 State，包含该方案的全部逻辑层；不同存储方案可以共存，通过算法字符串、层归属及可选的 State 类别筛选。名字用于资源定位，不能替代完整布局兼容性检查。标记注册后保持稳定，meta 定义、实际存储和 CPU/DP/PD 进程使用相同名称，不提供旧名称别名。 -| 方案 | State 成员名及 `resource_id` | +| 方案 | State 成员名及 `state_name` | | --- | --- | | 普通 BF16 KV | `main_gqa_bf16_state` | | INT8，group size 8，BF16 scale | `main_gqa_int8_group8_bf16_scale_state` | @@ -248,7 +266,7 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： ```text -AttResourceManager +AttentionStateManager ├── full_att_allocator ├── bounded_att_big_page_allocator ├── bounded_att_small_page_allocator @@ -267,25 +285,25 @@ AttResourceManager main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 State 同时命名为 main 和 draft。当前启动流程为各模型分别创建 State 和存储；GLM 的每个模型也分别创建 indexer ring。三种 allocator 的名称不带模型角色，同一分配域只初始化或绑定一个实例。 -`AttResourceManager`只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_resource_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `set_att_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 +`AttentionStateManager`只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_state_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `register_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 -`InferStateInfo` 保存 `att_resource_manager`、`is_draft_model` 和资源角色 `att_resource_role`，不额外绑定 `full_att_state`。各层直接访问管理器上对应方案的 State 成员，再通过局部引用取得层视图或写入 KV；FP8 backend 从具体量化方案的成员读取固定参数。每个 draft 的层由自身编号对应的 State 管理；`att_resource_role` 为 `main` 或 `draft0` 到 `draft35`，不能仅根据 `is_draft_model` 选择资源。是否初始化 full-attention backend 根据管理器中该角色实际启用的资源判断；FFN-only 和 linear-only 模型不创建 full-attention backend，linear backend 由相应模型独立初始化。 +`InferStateInfo` 保存 `att_state_manager`、`is_draft_model` 和资源角色 `att_state_role`，不额外绑定 `full_att_state`。各层直接访问管理器上对应方案的 State 成员，再通过局部引用取得层视图或写入 KV；FP8 backend 从具体量化方案的成员读取固定参数。每个 draft 的层由自身编号对应的 State 管理；`att_state_role` 为 `main` 或 `draft0` 到 `draft35`，不能仅根据 `is_draft_model` 选择资源。是否初始化 full-attention backend 根据管理器中该角色实际启用的资源判断；FFN-only 和 linear-only 模型不创建 full-attention backend，linear backend 由相应模型独立初始化。 资源身份用于 CPU/PD 资源布局与 DP 映射，所有传输按完整角色和方案匹配。部署配置确定各资源的页布局；PD 参数兼容性由 PD master 在连接阶段校验，CPU 服务绑定时核对页大小，磁盘缓存按 main/draft 模型版本隔离。 ## 4. 目录设计 -目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器、公共 State、allocator 和 pool 文件按对应类名命名；具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 +目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器使用 `manager.py`，三类公共 State 契约使用各自目录下的 `base_state.py`；allocator 和 pool 文件按对应类名命名，具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 ```text lightllm/common/ ├── basemodel/ │ ├── basemodel.py # 模型基础流程与模型专属配置 hook │ └── basemodel_utils.py # 配置归一化、KV 选项与 full/linear State 定义构建 -├── att_resources/ +├── attention_state/ │ ├── __init__.py │ ├── utils.py # State 角色与精度命名 -│ ├── att_resource_manager.py # AttResourceManager：State 装配、容量与存储初始化、缓存和传输接口 +│ ├── manager.py # AttentionStateManager：State 装配、容量与存储初始化、缓存和传输接口 │ ├── types.py # TensorField / StateTransferLayout 等 │ ├── allocators/ # 管理器直接持有，同一分配域共享实例 │ │ ├── __init__.py @@ -340,9 +358,9 @@ lightllm/utils/ `full_att/base_state.py` 同时定义 `FullAttentionConfig` 和 `FullAttentionState`；`bounded_att/base_state.py`、`att_tail/base_state.py` 分别定义另外两类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 -`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_pages`、`restore_bounded_att_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 +`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 -`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `resource_id`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 +`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `state_name`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个私有流程函数顺序执行：`_view_min_dimensions` 合并连续尾维，`_view_max_dtype` 整理最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧不连续就停止，保留前面的维度及单元素轴，再按行宽、stride、storage_offset 及整批地址对齐条件创建最大 8 字节的整数 view，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 @@ -364,15 +382,15 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttResourceManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 -CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttResourceManager.get_cpu_transfer_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 +CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 -无需模型特有初始化的场景直接使用 `AttResourceManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `set_att_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 +无需模型特有初始化的场景直接使用 `AttentionStateManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `register_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 -模型特有的 State 组合由 `model.py` 中三类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttResourceManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 +模型特有的 State 组合由 `model.py` 中三类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttentionStateManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 -token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttResourceManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttResourceManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 resource_id 匹配时可构建临时映射，它不决定模型的资源组织方式。 +token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttentionStateManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttentionStateManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 state_name 匹配时可构建临时映射，它不决定模型的资源组织方式。 #### 按逻辑算法创建 State @@ -386,7 +404,7 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 | GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | | K-pool | `att_tail.create_kpool_state(config, request_capacity, mtp_step=0, *, device)` | 带池化尾部窗口及 MTP 历史槽位的 K-pool State | -每个算法的创建入口独立放在 `create_gqa_state.py`、`create_mla_state.py`、`create_kda_state.py`、`create_gdn_state.py` 和 `create_kpool_state.py`，通过各分类包导出。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `set_att_state`。 +每个算法的创建入口独立放在 `create_gqa_state.py`、`create_mla_state.py`、`create_kda_state.py`、`create_gdn_state.py` 和 `create_kpool_state.py`，通过各分类包导出。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 @@ -397,15 +415,15 @@ state = create_gqa_state( config=full_config, kv_type="int4kv", group_size=8, page_size=16, device="meta", ) -manager.set_att_state("main", state) +manager.register_state("main", state) # 全部 State 准备完成后，由 manager 统一规划容量并分配存储。 ``` -根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttResourceManager.initialize_resources()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 +根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 -管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttResourceManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 +管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttentionStateManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 -`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisSlice`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 +`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisMapping`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 ## 5. allocator 与状态页 pool @@ -417,9 +435,9 @@ manager.set_att_state("main", state) | `bounded_att_big_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 大页 | | `bounded_att_small_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 小页 | -`AttResourceManager`初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 +`AttentionStateManager`初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 -allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由`AttResourceManager`明确描述。 +allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由`AttentionStateManager`明确描述。 大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 @@ -433,16 +451,16 @@ HOLD 页位于普通 KV 分配域之后，由 full-attention 存储预留。初 CPU 大小页池通过 `bounded_att/utils.py` 的 `create_page_pool_tensor()` 统一选择分配方式。未启用 `enable_dp_prompt_cache_fetch` 时沿用 `torch.empty(..., pin_memory=True)`；启用后直接创建共享 Storage 并原地调用 `cudaHostRegister(Portable | Mapped)`，避免先申请普通 pinned 页池再转共享存储造成整块复制。显式关闭 pin_memory 时只创建共享 CPU 存储，GPU 存储保持原分配方式。注册随本地 Storage 生命周期管理，最后一个视图释放时等待 GPU 使用结束再注销；接收进程需要自行注册其映射地址。KDA/GDN 页池均提供 `register_pin_memory()`，复用 `register_page_pool_tensor()` 注册 Conv/SSM 的整个 Storage；已 pinned 或空存储重复调用为空操作。`create_dp_ipc_copy()` 保留共享 CPU 页池，`create_pd_ipc_copy()` 始终排除全部大小页池。DP 接收端设置本地设备并读出资源后，调用 manager 的 `register_page_pools_pin_memory()`，统一注册所有 bounded State 的大小页池。`Portable` 使本进程其他 CUDA context 也识别其 pinned 属性，`Mapped` 建立设备可访问映射；它们不会替另一个进程注册地址。 -`StatePageAllocator` 只维护可用槽位。`KDAPagesPool`、`GDNPagesPool` 等具体页池只负责分配和持有各自的状态 buffer；State 使用外部传入的槽位完成保存和恢复。`AttResourceManager` 持有 allocator 并提供跨 State 的聚合数据操作，调用方直接使用 allocator。操作顺序如下： +`StatePageAllocator` 只维护可用槽位。`KDAPagesPool`、`GDNPagesPool` 等具体页池只负责分配和持有各自的状态 buffer；State 使用外部传入的槽位完成保存和恢复。`AttentionStateManager` 持有 allocator 并提供跨 State 的聚合数据操作，调用方直接使用 allocator。操作顺序如下： 1. 调用方直接从 `bounded_att_big_page_allocator` 或 `bounded_att_small_page_allocator` 分配槽位。新分配的槽位不代表已有可复用内容，必须先完成保存或载入才能发布。 -2. `save_bounded_att_pages` 聚合调用参与 State 保存；异步操作通过 stream/event 保证后续读取顺序。保存失败直接向上传递异常，由调用方停止发布并处理槽位生命周期。 -3. `restore_bounded_att_pages` 按调用方选定的槽位恢复全部参与 State，调用方保证前缀关联正确并已满足同步条件。 +2. `save_bounded_att_runtime_to_pages` 聚合调用参与 State 保存；异步操作通过 stream/event 保证后续读取顺序。保存失败直接向上传递异常，由调用方停止发布并处理槽位生命周期。 +3. `restore_bounded_att_runtime_from_pages` 按调用方选定的槽位恢复全部参与 State，调用方保证前缀关联正确并已满足同步条件。 4. 调用方确认引用已解除且异步使用完成，直接向同一个 allocator 归还共享槽位一次；pool 中的数据可保留至下一次写入。 CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排除在大页 allocator 的普通分配范围外；参与操作的 main/draft 大页 pool 分别为这些槽位提供存储。成员名为 `bounded_att_cpu_load_big_page_slot`（从 CPU 加载末尾不足一页的状态时使用）和 `bounded_att_cpu_offload_big_page_slot`（为卸载到 CPU 预留，当前尚无调用点），均属于 bounded-attention 大页池。 -`AttResourceManager`声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 +`AttentionStateManager`声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 ### 5.3 前缀长度与状态页的关联 @@ -454,11 +472,11 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 ## 6. 对外接口与 attention 资源暴露方式 -`AttResourceManager` 负责共享资源初始化及缓存和传输操作，模型工厂决定 State 组合。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，`AttResourceManager` 组织内部 State 完成操作。 +`AttentionStateManager` 负责共享资源初始化及缓存和传输操作，模型工厂决定 State 组合。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，`AttentionStateManager` 组织内部 State 完成操作。 ### 6.1 FullAttentionState：公共接口与不同 storage 的完整处理 -本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttResourceManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 +本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttentionStateManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 #### 逻辑层与物理存储层 @@ -484,9 +502,9 @@ State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get 三类 State 均提供 `contains_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 -`AttResourceManager.get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_name` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`att_type` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `att_type=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 +`AttentionStateManager.get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_role` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`state_type_filter` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `state_type_filter=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 -该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_states_created=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。三类模型工厂只返回 State 列表，`prepare_att_resources()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_att_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 +该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_state_registration_finished=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。三类模型工厂只返回 State 列表，`prepare_attention_states()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层 tail 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 @@ -527,7 +545,7 @@ class FullAttentionState: ) -> None: ... - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: ... def get_token_bytes(self) -> int: @@ -558,7 +576,7 @@ class FullAttentionState: # layout 描述该 State 在外部页中的位置及本地轴到页内轴的映射。 def write_to_cpu( self, token_indices, cpu_page_buffers, cpu_page_indices, - *, layout: StateTransferLayout, cpu_page_readies: List[bool] | None = None, + *, layout: StateTransferLayout, cpu_page_ready_flags: List[bool] | None = None, ) -> None: raise NotImplementedError @@ -572,7 +590,7 @@ class FullAttentionState: # 四、DP 迁移 # ======================================================================== # 其他 DP 的同一 TP 分片到当前 State；调用方准备各自设备上的 CUDA 索引，目标为 int64。 - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, remote_dp_full_att_state: FullAttentionState, ) -> None: @@ -585,13 +603,13 @@ class FullAttentionState: raise NotImplementedError # pd_page_buffers 为页集合，pd_page_index 选择物理页；layout 区域偏移相对该页起点。 - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices, pd_page_buffers, *, pd_page_index: int, layout: StateTransferLayout, ) -> None: raise NotImplementedError - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers, token_indices, *, pd_page_index: int, layout: StateTransferLayout, ) -> None: @@ -603,7 +621,7 @@ class FullAttentionState: | 接口 | 完整语义 | | --- | --- | | 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | -| `initialize_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | +| `allocate_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | | `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | | `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | @@ -634,7 +652,7 @@ GQA、MLA 及各量化方案只导出 KV、动态 scale 和随 token 保存的 i manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字段的实际最大结束位置安排后续资源，并通过布局终点计算整页大小；manager 不解析 Conv/SSM/K-pool 的内部结构。CPU 服务复用相同 State 接口；当前 CPU 联合页仍只包含 full/bounded，tail 按恢复边界重置。DP 映射只包含 full/bounded attention，调用它们无参数的 PD 区域接口，只取字段名和本地轴切片，不使用区域的页形状与字节偏移；本次复制的 token/request 数量由调用参数决定。 -资源身份使用 `state.resource_id`；层映射、容量和 TP 分片事实继续由 State 的 config、capacity、page_size 及计算接口维护。CPU/DP/PD 由抽象接口约定并由具体 State 实现，管理器直接调用对应方法，不再维护额外的传输能力集合。信息保留在其所属对象，传输布局仅描述当前路径所需字段。 +资源身份使用 `state.state_name`；层映射、容量和 TP 分片事实继续由 State 的 config、capacity、page_size 及计算接口维护。CPU/DP/PD 由抽象接口约定并由具体 State 实现，管理器直接调用对应方法，不再维护额外的传输能力集合。信息保留在其所属对象，传输布局仅描述当前路径所需字段。 字段描述只读，引用的 Tensor 可以由合法计算操作更新。字段列表不表示数据已就绪，stream/event 同步和前缀关联仍由调用流程负责。每次增加辅助 buffer 或 packed 字段，都需要检查固定传输字段、容量计算、写入和复制/传输是否一起覆盖它。公共层不根据固定属性名推测资源组成。 @@ -668,9 +686,9 @@ manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字 #### 索引、异步执行与生命周期 -构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `initialize_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `initialize_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `initialize_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式由配置检查；静态标定参数在 `bind_calibration()` 中检查。 +构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `allocate_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `allocate_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `allocate_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式由配置检查；静态标定参数在 `bind_calibration()` 中检查。 -State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 构造时不接收 `scales` 或 `q_scales`；只能通过 `bind_calibration()` 绑定真实的 CPU 标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 +State 在后续 `allocate_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 构造时不接收 `scales` 或 `q_scales`；只能通过 `bind_calibration()` 绑定真实的 CPU 标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 `logical_layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 @@ -688,13 +706,13 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 | --- | --- | | `write_to_cpu(..., layout)` | 当前 State → CPU 缓存；layout 描述外部页区域、层/轴映射和写入副本选择 | | `read_from_cpu(..., layout)` | CPU 缓存 → 当前 State；layout 描述该目标 State 从页中读取的对应范围 | -| `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` | 其他 DP 同一 TP 分片 → 当前 State，复制完整 token 的 KV、动态量化参数及 indexer 字段，不接收 mapping | -| `write_to_pd_buffer(..., layout)` | 当前 State → `pd_page_index` 指定的 PD 页；按 layout 写入指定区域 | -| `read_from_pd_buffer(..., layout)` | `pd_page_index` 指定的 PD 页 → 当前 State；按 layout 恢复对应区域 | +| `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` | 其他 DP 同一 TP 分片 → 当前 State，复制完整 token 的 KV、动态量化参数及 indexer 字段，不接收 mapping | +| `write_to_pd_page(..., layout)` | 当前 State → `pd_page_index` 指定的 PD 页；按 layout 写入指定区域 | +| `read_from_pd_page(..., layout)` | `pd_page_index` 指定的 PD 页 → 当前 State；按 layout 恢复对应区域 | -`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，`AttResourceManager`按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 +`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，`AttentionStateManager`按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 -基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttResourceManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 +基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttentionStateManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 每次调用完整处理其映射涉及的数据及关联的动态 scale、packed 字段；一次组合传输可以由若干次调用完成，上层检查所有必需数据区域均已完成。固定参数按约定传递或验证兼容性。位置映射的组成、单位及生成责任见 6.5 节。 @@ -702,41 +720,41 @@ State 在后续 `initialize_storage()` 中自行分配数据 buffer：GQA 与静 形状相同不能证明不同模型的 KV 语义可互换。模型/资源身份和 token 前缀匹配由管理器与缓存流程校验；State 检查布局及固定参数兼容性。跨布局重编码需要明确实现，本地 `copy_full_att_tokens` 不承担这一职责。 -### 6.2 AttResourceManager：内部资源的聚合操作与数据传输 +### 6.2 AttentionStateManager：内部资源的聚合操作与数据传输 -`AttResourceManager` 对内部资源组合提供统一的数据操作。一个模型可以包含不同角色或不同具体方案的 FullAttentionState、KDAState、GDNState 等状态实现，同一角色和具体方案最多一个；token 复制覆盖 full-attention 参与集合，清零操作覆盖需要初始运行态的参与集合，状态页操作覆盖指定大小页域的参与集合。外部调用方无需知道具体数量、类型或成员名称。 +`AttentionStateManager` 对内部资源组合提供统一的数据操作。一个模型可以包含不同角色或不同具体方案的 FullAttentionState、KDAState、GDNState 等状态实现，同一角色和具体方案最多一个；token 复制覆盖 full-attention 参与集合，清零操作覆盖需要初始运行态的参与集合，状态页操作覆盖指定大小页域的参与集合。外部调用方无需知道具体数量、类型或成员名称。 -数据操作的依赖方向是 `radix cache / CPU cache / 其他调用方 → AttResourceManager → 各 State`；槽位分配与归还由调用方直接访问管理器持有的共享 allocator。状态页接口不接收 radix 节点、树或前缀匹配对象，也不读取 radix 内部状态；radix 只保存并使用管理器提供的页类别、槽位 ID 和状态长度等数据。状态页在未启用 radix 时同样可以保存、恢复或用于 CPU 缓存。 +数据操作的依赖方向是 `radix cache / CPU cache / 其他调用方 → AttentionStateManager → 各 State`；槽位分配与归还由调用方直接访问管理器持有的共享 allocator。状态页接口不接收 radix 节点、树或前缀匹配对象，也不读取 radix 内部状态；radix 只保存并使用管理器提供的页类别、槽位 ID 和状态长度等数据。状态页在未启用 radix 时同样可以保存、恢复或用于 CPU 缓存。 #### 资源初始化、State 查询与传输布局查询 | 接口 | 语义 | | --- | --- | -| `AttResourceManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | -| `set_att_state(model_name, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | -| `finish_state_creation()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_states_created = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | -| `load_full_att_calibration(args)` | 在 `finish_state_creation()` 后、存储分配前加载静态 FP8 标定；文件只读一次，逐 State 绑定物理层区间及 TP 参数；没有静态 FP8 State 时直接返回 | -| `initialize_resources(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：在完成创建标记和标定绑定后，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | -| `get_att_state(model_name, att_type, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_states_created` 并使用 `lru_cache` 缓存 | +| `AttentionStateManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | +| `register_state(model_role, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | +| `finish_state_registration()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_state_registration_finished = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | +| `load_full_att_calibration(args)` | 在 `finish_state_registration()` 后、存储分配前加载静态 FP8 标定；文件只读一次，逐 State 绑定物理层区间及 TP 参数；没有静态 FP8 State 时直接返回 | +| `initialize_storage(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：在完成创建标记和标定绑定后，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | +| `get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_state_registration_finished` 并使用 `lru_cache` 缓存 | | `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | | `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | | `get_attention_tail_states(role=None)` | 返回指定角色或全部 `AttentionTailState` 的 tuple | | `get_full_att_token_bytes()` | 返回当前 TP rank 全部 full-attention State 的单 token 存储字节数，包含动态 scale 和行 padding | -| `get_cpu_page_bytes()` | 依赖 `get_cpu_transfer_layouts()`，取全部 State 的最大 `byte_end` 并按 16 字节对齐；包含字段间及 State 间的必要 padding，页内不保存格式指纹、长度或版本 | -| `get_cpu_transfer_layouts()` | 按全局 `cpu_cache_token_page_size` 生成 CPU 联合页中各 State 的字段偏移和 TP 映射，使用 `lru_cache` 缓存 | -| `get_pd_full_att_transfer_layouts()` | 从全局启动参数读取 `pd_kv_page_size`，按 State 字段的实际偏移和结束位置生成 `kv` 布局及最终页跨度；不依赖 `get_pd_page_bytes()`，按 `resource_id` 索引并缓存 | -| `get_pd_bounded_att_or_tail_transfer_layouts()` | 按实际字段结束位置独立生成 `att_state` 布局及其自身所需大小，不查询 KV 布局或执行组合容量校验；每个 State 保存一份运行态，`tokens_per_page=None` | +| `get_cpu_page_bytes()` | 依赖 `get_cpu_page_layouts()`，取全部 State 的最大 `byte_end` 并按 16 字节对齐；包含字段间及 State 间的必要 padding，页内不保存格式指纹、长度或版本 | +| `get_cpu_page_layouts()` | 按全局 `cpu_cache_token_page_size` 生成 CPU 联合页中各 State 的字段偏移和 TP 映射，使用 `lru_cache` 缓存 | +| `get_pd_token_page_layouts()` | 从全局启动参数读取 `pd_kv_page_size`，按 State 字段的实际偏移和结束位置生成 `kv` 布局及最终页跨度；不依赖 `get_pd_page_bytes()`，按 `state_name` 索引并缓存 | +| `get_pd_runtime_page_layouts()` | 按实际字段结束位置独立生成 `att_state` 布局及其自身所需大小，不查询 KV 布局或执行组合容量校验；每个 State 保存一份运行态，`tokens_per_page=None` | | `get_pd_page_bytes()` | 查询两类布局大小并集中决定实际 PD 块容量：返回 KV 页与完整运行态页占用的最大值；有 KV 且 KV 页较小时记录 warning 并扩大物理页；无资源时返回 0。包含 field 与 State 间的 padding，可基于 meta State 查询，不分配传输 buffer | -| `initialize_storage(token_capacity, page_size, *, big_page_tokens, small_page_capacity, device)` | 为已创建 State 分配实际 full-attention 存储、共享 allocator 和 bounded 大小页 | +| `allocate_storage(token_capacity, page_size, *, big_page_tokens, small_page_capacity, device)` | 为已创建 State 分配实际 full-attention 存储、共享 allocator 和 bounded 大小页 | | `create_pd_ipc_copy()` | PD 导出：保留运行态 Tensor 和容量信息，排除 allocator 及所有大小页池 | | `create_dp_ipc_copy()` | DP shared 导出：保留运行态和共享 CPU 大小页池，仅排除 allocator | | `register_page_pools_pin_memory()` | 在当前 GPU 上统一注册全部 bounded State 的大小页池，已注册存储可重复调用 | -聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `set_att_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `resource_id` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数。`get_att_state` 提供创建完成后的单对象缓存查询，推理层通过该接口直接取得当前层使用的资源，不另存成员名。 +聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `register_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `state_name` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数。`get_state` 提供创建完成后的单对象缓存查询，推理层通过该接口直接取得当前层使用的资源，不另存成员名。 -State 装配期间由管理器内部遍历 `__dict__`，公共聚合查询及 `get_att_state` 均会断言失败，避免缓存不完整的 main/draft 集合；模型准备阶段只创建和注册 State，不查询推理资源。装配完成后先调用 `finish_state_creation()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 +State 装配期间由管理器内部遍历 `__dict__`，公共聚合查询及 `get_state` 均会断言失败，避免缓存不完整的 main/draft 集合；模型准备阶段只创建和注册 State，不查询推理资源。装配完成后先调用 `finish_state_registration()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 -容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `initialize_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded/indexer 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 +容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `allocate_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded/indexer 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 实际存储分配完成后执行 warmup。CPU 布局由第一次容量查询或读写时按需生成；PD 容量规划会通过 `get_pd_page_bytes()` 在 meta State 上提前生成 PD 布局，不分配实际传输 buffer。PD 页的 token 容量由全局 `get_env_start_args().pd_kv_page_size` 唯一确定，两个 PD 布局 getter 均不接收页 token 数，按 manager 实例缓存并直接返回 `Mapping[str, StateTransferLayout]`，没有对应资源时返回空字典。 @@ -758,13 +776,13 @@ IPC 发布按用途调用 `create_pd_ipc_copy()` 或 `create_dp_ipc_copy()`， | --- | --- | --- | | full-attention token 复制 | `copy_full_att_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | | bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页 | -| bounded 状态页保存 | `save_bounded_att_pages(...)` | 按 state_token_lens 保存全部参与 bounded State 的同类 pool、同一组槽位 | -| bounded 状态页恢复 | `restore_bounded_att_pages(...)` | 按 state_token_lens 将选定快照恢复到对应请求运行态及 ring 位置 | +| bounded 状态页保存 | `save_bounded_att_runtime_to_pages(...)` | 按 state_token_lens 保存全部参与 bounded State 的同类 pool、同一组槽位 | +| bounded 状态页恢复 | `restore_bounded_att_runtime_from_pages(...)` | 按 state_token_lens 将选定快照恢复到对应请求运行态及 ring 位置 | | bounded 验证后提交 | `commit_bounded_att_state(req_indices, accepted_steps, *, role)` | 将指定角色的接受步数分发给参与 bounded State；由 State 选择已接受运行态，不暴露游标 Tensor | | attention tail 初始化与无快照恢复 | `reset_attention_tail_state(req_indices)` | 清零指定请求在全部参与 tail 中的状态；调用方保证空前缀或可重置的缓存边界，保留其他状态 | | CPU 写出 / 读回 | `write_req_to_cpu(...)` / `read_req_from_cpu(...)` | 组合 KV 和状态页，按调用参数提供的长度读写各 State 在外部页中的区域 | -| DP 复制 | `copy_from_other_dp(remote_att_manager, src_token_indexes, dst_token_indexes, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，清零目标 tail；不分配资源 | -| PD 写出 / 读回 | `write_to_pd_buffer(...)` / `read_from_pd_buffer(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | +| DP 复制 | `copy_from_dp(remote_att_manager, src_token_indices, dst_token_indices, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，清零目标 tail；不分配资源 | +| PD 写出 / 读回 | `write_to_pd_page(...)` / `read_from_pd_page(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | #### 必须实现的聚合接口 @@ -781,15 +799,15 @@ from typing import Any, List, Literal PageKind = Literal["big", "small"] -class AttResourceManager: +class AttentionStateManager: full_att_allocator: FullAttKVAllocator | None bounded_att_big_page_allocator: StatePageAllocator | None bounded_att_small_page_allocator: StatePageAllocator | None def __init__(self, request_capacity: int): ... - def set_att_state( - self, model_name: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState, + def register_state( + self, model_role: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState, ) -> None: ... # ======================================================================== @@ -807,20 +825,20 @@ class AttResourceManager: def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... # Bounded Attention 状态页：供 radix/CPU cache 复用，调用方直接分配和归还槽位。 - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: ... - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: ... # MTP 推测解码：只提交指定角色；接受步数包含本轮已经执行的基准输入步。 @@ -834,14 +852,14 @@ class AttResourceManager: def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... # 基类使用 lru_cache 缓存布局，存储初始化后清除缓存。 - def get_cpu_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... # CPU 联合页：KV + 全部参与 linear State 的有效页。 # 纯 full-attention 模型没有 linear 页；混合模型不能省略它们。 def write_req_to_cpu( self, token_indices, cpu_page_buffers, cpu_page_indices, *, layouts: Mapping[str, StateTransferLayout], - cpu_page_readies: List[bool] | None = None, + cpu_page_ready_flags: List[bool] | None = None, big_page_slot_ids: List[int] | None = None, small_page_slot_ids: List[int] | None = None, ) -> None: ... @@ -856,26 +874,26 @@ class AttResourceManager: ) -> None: ... # DP 缓存迁移：索引已由调用方分配，空 Tensor/列表表示不搬运该类数据。 - def copy_from_other_dp( - self, remote_att_manager: AttResourceManager, - src_token_indexes: torch.Tensor, dst_token_indexes: torch.Tensor, + def copy_from_dp( + self, remote_att_manager: AttentionStateManager, + src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, src_big_page_ids: List[int], dst_big_page_ids: List[int], src_req_idx: int, dst_req_idx: int, ) -> None: ... # 基类直接构建两类 PD 布局，各自使用 lru_cache 缓存；无对应资源时返回空字典。 - def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + def get_pd_token_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... - def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ... + def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... # PD 分离：pd_page_buffers 为页集合，pd_page_index 必填；kind 沿用 kv / att_state 两类任务。 - def write_to_pd_buffer( + def write_to_pd_page( self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], layouts: Mapping[str, StateTransferLayout], token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, ) -> None: ... - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], layouts: Mapping[str, StateTransferLayout], token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, @@ -899,7 +917,7 @@ class AttResourceManager: - full-attention 没有需要清零的空前缀递推状态；该操作不遍历或清空历史 KV，不更改请求映射，也不失效大小页快照或归还索引。纯 full-attention 管理器合法地不做数据操作。 - 目标请求槽位须唯一且在有效范围内，main/draft 所需运行态全部初始化后才可使用。重复执行在没有并发访问时具有相同结果；方法返回仍遵守 stream 完成约定。 -这是长度 0 的数据初始化接口，不接收非零前缀长度。命中非零前缀时使用 `restore_bounded_att_pages` 或 CPU/DP/PD 读回；没有 tail 快照时在保证可重置边界后通过 `reset_attention_tail_state` 清零 tail；不能在 bounded 运行态恢复后调用此方法，否则会清除已经恢复的 bounded 数据。 +这是长度 0 的数据初始化接口，不接收非零前缀长度。命中非零前缀时使用 `restore_bounded_att_runtime_from_pages` 或 CPU/DP/PD 读回；没有 tail 快照时在保证可重置边界后通过 `reset_attention_tail_state` 清零 tail；不能在 bounded 运行态恢复后调用此方法，否则会清除已经恢复的 bounded 数据。 #### 验证后运行态提交 @@ -934,7 +952,7 @@ main/draft 共享管理器，因此必须显式指定 `role="main"` 或具体 dr #### 多个 State 的参与规则 -`AttResourceManager`在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 +`AttentionStateManager`在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 例如，一个管理器同时管理两个 KDAState 和一个 GDNState，大页槽位 17 对应： @@ -945,17 +963,17 @@ bounded_att_big_page_allocator 分配一次 slot_id = 17 └── GDNState C.big_pages_pool[17] ``` -调用 `save_bounded_att_pages(..., state_token_lens=[256], page_kind="big", slot_ids=[17])` 保存三个 State；`restore_bounded_att_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 +调用 `save_bounded_att_runtime_to_pages(..., state_token_lens=[256], page_kind="big", page_slot_ids=[17])` 保存三个 State；`restore_bounded_att_runtime_from_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 -调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 +调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_runtime_from_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 full-attention 历史 KV 和没有独立 pool 的 AttentionTailState 不属于这组状态页操作。需要同时恢复这些数据的调用方，使用相应资源操作，并保证历史 KV、bounded 快照和 tail 对应同一前缀且均已就绪。 #### 保存与直接归还槽位的完整性 - 调用方只从共享 allocator 分配一次槽位，全部参与 pool 覆盖同一索引范围。分配成功仅表示槽位可用，不表示其中已有有效状态。后续数据操作失败时，本次数据不得发布；调用方在已提交操作结束后决定重写或归还槽位。 -- `save_bounded_att_pages` 将同一组槽位传给全部参与 State 保存。只有所有参与者都完成对应位置的数据保存，调用方才能发布页面；一项失败时直接透传异常，不发布本次结果。 -- `restore_bounded_att_pages` 使用调用方选定的快照恢复各参与 State。复制阶段失败时直接透传异常，源页保持不变，调用方不得消费半恢复的数据。 +- `save_bounded_att_runtime_to_pages` 将同一组槽位传给全部参与 State 保存。只有所有参与者都完成对应位置的数据保存，调用方才能发布页面；一项失败时直接透传异常，不发布本次结果。 +- `restore_bounded_att_runtime_from_pages` 使用调用方选定的快照恢复各参与 State。复制阶段失败时直接透传异常，源页保持不变，调用方不得消费半恢复的数据。 - 归还槽位前，调用方必须解除相关使用引用并满足完成条件，再直接向对应 allocator 归还一次。只归还指定大小页域的槽位，不操作相同数字的 KV 索引或请求槽位；槽位中的旧内容不能作为重新分配后的有效快照使用。 保存时全部参与 State 必须实际具备所要求位置的状态；不能从较新的状态补造较早快照。main/draft 可以分别计算，但聚合操作的完整性条件不变；若具体实现分阶段捕获参与者，未全部完成前必须保持聚合页不可用。 @@ -966,7 +984,7 @@ GPU 数据操作遵守当前 stream 的依赖约定。调用方通过 stream/eve #### CPU / DP / PD 聚合接口的参数与数据路径 -`layouts` 是按完整 `resource_id`（包含 main/draft 角色）索引的 `StateTransferLayout` 映射，用于 CPU/PD 外部页。DP 不构造字段映射：manager 直接遍历本地 State，通过同名成员取得远端 State，将对应 token、页或请求索引交给具体实现。DP 不再维护映射缓存，也不通过 helper 转发。 +`layouts` 是按完整 `state_name`（包含 main/draft 角色）索引的 `StateTransferLayout` 映射，用于 CPU/PD 外部页。DP 不构造字段映射：manager 直接遍历本地 State，通过同名成员取得远端 State，将对应 token、页或请求索引交给具体实现。DP 不再维护映射缓存，也不通过 helper 转发。 CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。CPU 读写由管理器按 full-attention、bounded 状态页、末页运行态恢复及 tail 重置分阶段组织，高频读写路径不记录流程日志；buffer 与字段搬运约束交给 State，管理器检查恢复所需的目标请求、非空前缀及完整页槽位数量。异步搬运的完成仍由调用方的 stream/event 保证。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 @@ -976,7 +994,7 @@ CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能 - `token_indices: torch.Tensor` 是由调用方准备在 CUDA 上的 full-attention 物理 token 索引，`cpu_page_indices: List[int]` 是 CPU 页池的页 ID，仅接受 Python 列表，不接受 Tensor 或其他序列，也不在接口内部转换。索引按列表顺序排列，从各页槽位 0 开始；长度与页数遵守 6.5 节的整页/尾页约定。页列表按逻辑前缀顺序排列，不要求物理页 ID 连续或递增。只传状态的 CPU 格式使用空整数 Tensor 作为 token 索引。 - 格式包含 linear 页时，管理器写出接口分别接收 `big_page_slot_ids` 与 `small_page_slot_ids`；`cpu_page_indices` 前段对应大页、后段对应小页，每段与槽位列表一一对应。未参与的页池传 None 或空列表。full-attention KV 按完整页面列表写出一次，各 bounded State 再分别从大小页池写出快照；读回接口固定写入大页池，不接收 `page_kind`。管理器的 `big_page_slot_ids: List[int]` 只包含完整 CPU 页对应的普通大页槽位，不包含末尾碎页的临时槽位；没有 linear 页时为 None。槽位参数仅接受 Python 列表，不接受 Tensor。快照搬运不需要逐页长度数组。 -- `cpu_page_readies: List[bool] | None` 是本批 CPU 页的跳过写入标记，顺序与 `cpu_page_indices` 一致；True 保留整张联合页，None 表示全部写出。缓存层保留完整 token/page 序列并一次调用资源管理器，不逐页发起传输，也不通过过滤页 ID 改变 token 映射。资源层不查询或发布 ready 状态，全部 ready 时直接返回。 +- `cpu_page_ready_flags: List[bool] | None` 是本批 CPU 页的跳过写入标记，顺序与 `cpu_page_indices` 一致；True 保留整张联合页，None 表示全部写出。缓存层保留完整 token/page 序列并一次调用资源管理器，不逐页发起传输，也不通过过滤页 ID 改变 token 映射。资源层不查询或发布 ready 状态，全部 ready 时直接返回。 - FullAttentionState 将 ready 标记传给各自 CPU/PD 共用的传输入口；CPU 写出按 State 构造一次 GPU bool Tensor，全部字段复用。在 kernel 内按逻辑页屏蔽源 KV/量化字段读取和目标写入；若当前 block 没有待写元素，直接跳过后续索引读取、地址计算和数据搬运。跨页 block 混有 ready 与未就绪元素时保留逐元素 mask。PD 与读回不传标记、不分配此 Tensor，也不执行 ready 归约判断。bounded State 接收对应大小页段的 ready 列表，在取快照和构造复制视图前跳过 ready 页。 - CPU 页不保存 tail 数据，写出接口不接收请求索引；布局只包含 full-attention 与 bounded-attention。 - `read_req_from_cpu` 一次处理同一前缀的有序页面，将全部 KV 写到给定物理位置，将各页 linear 快照写入对应 pool 槽位。管理器根据 `token_indices.numel()` 和布局的页 token 数判断是否存在末尾碎页，将其完整 linear 快照写入预留 LOAD 大页；完整页写入 `big_page_slot_ids` 指定的大页槽位。读回后直接使用最后一页恢复 `req_idx` 的 linear 运行态，不修改调用方的槽位列表。仅含碎页时传 `big_page_slot_ids=[]`；只传状态的格式使用空 token 索引，按传入的普通槽位恢复。 @@ -992,7 +1010,7 @@ CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能 `src_manager` 明确来源 DP/TP 存储实例，也可以是传输进程中已绑定共享存储的管理器。源和目标 token 索引等长、逐项对应；`src_req_indices`、`dst_req_indices` 仅接受 `List[int]`，成对提供且数量相等，用来迁移 prefill 基准 bounded 运行态；KDA/GDN 源游标必须为 0。有 tail 时直接清零目标请求的 tail，不传递额外长度。token 索引与请求槽位是不同的寻址维度，无需等长；前缀对应关系及可重置边界由调用方保证。 -只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，`AttResourceManager` 将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 +只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，`AttentionStateManager` 将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 **PD 页读写** @@ -1009,15 +1027,15 @@ CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引 radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存和恢复状态页时调用上述通用接口，前缀与槽位的对应关系由 radix 调用流程保证；槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 -CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，`AttResourceManager`负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 +CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，`AttentionStateManager`负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 ### 6.3 模型推理使用具体资源视图 -Model 和 InferStateInfo 通过 `att_resource_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_gqa_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_resource_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 +Model 和 InferStateInfo 通过 `att_state_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_gqa_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_state_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与 tail 的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 -普通 KV 层可支持多种存储精度。推理层使用 `get_att_state(model_name=infer_state.att_resource_role, att_type="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState)` 取得资源。MLA 推理层和 GLM indexer 使用 `att_type="mla"`，indexer 使用自己的 `layer_idx`。通用写入模板使用空字符串和 `FullAttentionState` 类型，不推断算法。匹配结果由管理器缓存，推理层和 indexer 不保存重复成员名，也不需要模型启动时逐层绑定。不同精度的 State 由逻辑层归属区分。 +普通 KV 层可支持多种存储精度。推理层使用 `get_state(model_role=infer_state.att_state_role, state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState)` 取得资源。MLA 推理层和 GLM indexer 使用 `state_type_filter="mla"`，indexer 使用自己的 `layer_idx`。通用写入模板使用空字符串和 `FullAttentionState` 类型，不推断算法。匹配结果由管理器缓存，推理层和 indexer 不保存重复成员名，也不需要模型启动时逐层绑定。不同精度的 State 由逻辑层归属区分。 Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 State 将共享消费者映射到源物理层。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按自身压缩比例访问对应角色的 tail 成员。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 @@ -1025,11 +1043,11 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S ### 6.4 调用边界与独立传输进程 -单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttResourceManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 +单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttentionStateManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 -具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttResourceManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该`AttResourceManager`的具体实现提供。 +具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttentionStateManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该`AttentionStateManager`的具体实现提供。 -跨进程共享通过管理器的 `create_pd_ipc_copy()` / `create_dp_ipc_copy()` 保留 State Tensor 引用并由 IPC 传递句柄，不调用 `initialize_storage()` 重新分配 KV；两者均排除 allocator，PD 副本排除全部大小页池，DP 副本保留共享 CPU 大小页池，不导出整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 +跨进程共享通过管理器的 `create_pd_ipc_copy()` / `create_dp_ipc_copy()` 保留 State Tensor 引用并由 IPC 传递句柄，不调用 `allocate_storage()` 重新分配 KV；两者均排除 allocator，PD 副本排除全部大小页池，DP 副本保留共享 CPU 大小页池，不导出整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 ### 6.5 复杂 attention 组合与 TP 的位置映射 @@ -1043,7 +1061,7 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S | --- | --- | --- | | 本地存储布局与分片事实 | 具体 State 的配置、存储参数及内部 Tensor | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | | 传输字段及本地 TP 切片 | 具体 State 的 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` | 字段选择、全局形状、State 内字段的起止区间和各 rank 的读写切片 | -| 组合页布局及资源间分区 | 模型专属 `AttResourceManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | +| 组合页布局及资源间分区 | 模型专属 `AttentionStateManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | | 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `StateTransferLayout`；按资源与 CPU/PD 路径分别保存 | | 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | @@ -1053,17 +1071,17 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S #### StateTransferLayout：单个 State 的传输布局 -`StateTransferLayout` 通过 `resource_id` 绑定一个具体 State，描述它在 CPU/PD 外部页中的传输布局。它通过 `regions` 组织字段级映射，每个 `TransferRegion` 描述 State 内一个字段与外部页区域之间的对应关系，`AxisSlice` 进一步描述轴上的拷贝范围： +`StateTransferLayout` 通过 `state_name` 绑定一个具体 State，描述它在 CPU/PD 外部页中的传输布局。它通过 `regions` 组织字段级映射，每个 `TransferRegion` 描述 State 内一个字段与外部页区域之间的对应关系，`AxisMapping` 进一步描述轴上的拷贝范围： ```text StateTransferLayout # 一个 State 的传输布局 -├── resource_id # 对应的 State +├── state_name # 对应的 State ├── tokens_per_page # 外部页的 token 容量 ├── byte_start / byte_end # 本 State 的页内预留区间 [start, end) └── regions: tuple[TransferRegion, ...] └── TransferRegion # 一个字段与外部区域的映射 ├── field / byte_start / byte_end / dtype / axes / shape - ├── axis_slices: tuple[AxisSlice, ...] + ├── axis_mappings: tuple[AxisMapping, ...] └── write_enabled ``` @@ -1075,14 +1093,14 @@ CPU buffer 使用 `(page_num, page_bytes)` 的二维页池，通过 `cpu_page_bu | 字段 | 语义与单位 | | --- | --- | -| `resource_id` | 与管理器成员名和 State 的资源名称一致，用于选择当前资源的页区域 | +| `state_name` | 与管理器成员名和 State 的资源名称一致，用于选择当前资源的页区域 | | `tokens_per_page` | 含 KV 的页格式为正整数 token 容量；纯递推状态页为 `None`，按请求状态或 checkpoint 槽位处理 | | `byte_start` | 本 State 相对整页的起始字节，包含该位置；region 已使用页内坐标，搬运时不再重复加上此起点 | | `byte_end` | 本 State 的结束字节，不包含该位置，涵盖自身字段、对齐及预留空间；与起点之差是本 State 占用，不是物理页间距 | | `regions` | 本 State 参与传输的物理字段或片段列表，同一字段可分为多个映射项 | | 每项的 `field`、`dtype`、`axes`、`shape` | 标识 State 中的字段及外部页区域；axes 命名各维度，shape 用元素数量；外部区域采用连续布局，stride 由 shape 和 dtype 推导 | | 每项的 `byte_start`、`byte_end` | 字段的左闭右开字节区间，组合后相对整页；长度等于 `prod(shape) * dtype.itemsize`，不包含字段间的对齐空隙，必须位于所属 State 区间内 | -| 每项的 `axis_slices` | `tuple[AxisSlice, ...]`，通过 `axis`、`src_start`、`dst_start`、`length` 描述轴名、源起点、目标起点和元素数量；一个 region 内同一轴只出现一次，src_circular/dst_circular 描述两侧环形寻址；普通非连续范围拆成多个 region | +| 每项的 `axis_mappings` | `tuple[AxisMapping, ...]`，通过 `axis`、`local_start`、`external_start`、`length` 描述轴名、本地起点、外部起点和元素数量；一个 region 内同一轴只出现一次，local_circular/external_circular 描述两侧环形寻址；普通非连续范围拆成多个 region | | 每项的 `write_enabled` | 指定 TP 副本是否负责写入该区域，避免多个副本重复写入；恢复时各目标副本均可读取 | 以下冻结数据类只描述传输位置，不包含搬运方法： @@ -1093,13 +1111,13 @@ import torch @dataclass(frozen=True) -class AxisSlice: +class AxisMapping: axis: str - src_start: int - dst_start: int + local_start: int + external_start: int length: int - src_circular: bool = False - dst_circular: bool = False + local_circular: bool = False + external_circular: bool = False @dataclass(frozen=True) @@ -1110,13 +1128,13 @@ class TransferRegion: dtype: torch.dtype # 外部区域的数据类型 axes: tuple[str, ...] # 如 (token, layer, head, dim) shape: tuple[int, ...] # 外部区域在各轴上的完整尺寸，不含页池的 page 轴 - axis_slices: tuple[AxisSlice, ...] + axis_mappings: tuple[AxisMapping, ...] write_enabled: bool = True @dataclass(frozen=True) class StateTransferLayout: - resource_id: str # 含模型角色，如 main_gqa_fp16_state + state_name: str # 含模型角色，如 main_gqa_fp16_state tokens_per_page: int | None # KV 页为 token 容量，纯状态页为 None byte_start: int # 本 State 的页内起点，包含 byte_end: int # 本 State 的页内终点，不包含，含自身对齐与预留空间 @@ -1125,13 +1143,13 @@ class StateTransferLayout: `field` 在初始化时与 State 字段匹配；本地 Tensor 继续保留自身的 dtype、stride 和字段起点，可以是 packed buffer 的非连续视图。外部页中每个 region 的完整区域采用连续布局，`axes` 与 `shape` 长度相同且轴名唯一，逐维 stride 由后续维度大小乘 dtype 字节数推导，不再作为输入字段。多个 region 可以通过切片引用同一外部区域，例如 Conv 的 Q/K/V 分片。部署配置确定布局，CPU cache 读写直接使用字段区域与 TP 切片。 -`AxisSlice` 的起点和长度均以元素为单位。在 `TransferRegion` 中，方向固定为本地 State 到外部页，读回时交换源、目标起点。构造时使用字段名，搬运时通过属性访问，不再依赖四元组的位置。 +`AxisMapping` 的起点和长度均以元素为单位。在 `TransferRegion` 中，映射固定描述本地 State 与外部页的对应关系；写出和读回共用两侧起点，仅反转数据流向。构造时使用字段名，搬运时通过属性访问，不再依赖四元组的位置。 层维度统一完整、同序复制，搬运前检查源目标层数一致。逻辑层到物理存储位置的映射由 State 自身维护，例如逻辑层 `[3, 7]` 紧凑存为 `[0, 1]`，外部页沿用相同层顺序。CPU/PD 区域和 DP State 复制均不提供层子集选择或层重排参数。 -token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_slices`。其余搬运轴通过 `AxisSlice` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。普通非连续范围使用多个连续片段表达；环形轴使用 `src_circular` / `dst_circular`，按两侧 Tensor 的实际轴容量对 `start + i` 取模。标记始终按本地 State → 外部页解释，读回时反向执行；搬运长度不超过各环形轴容量。 +token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_mappings`。其余搬运轴通过 `AxisMapping` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。普通非连续范围使用多个连续片段表达；环形轴使用 `local_circular` / `external_circular`，按两侧 Tensor 的实际轴容量对 `start + i` 取模。标记始终按本地 State → 外部页解释，读回时反向执行；搬运长度不超过各环形轴容量。 -`shape` 表达该外部字段区域的完整范围，而 `axis_slices` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： +`shape` 表达该外部字段区域的完整范围，而 `axis_mappings` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： ```python replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) @@ -1149,15 +1167,15 @@ k_region = TransferRegion( dtype=torch.float16, axes=("token", "layer", "head", "dim"), shape=(8, 2, 4, 128), - axis_slices=( - AxisSlice(axis="head", src_start=0, dst_start=2, length=2), - AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=2, length=2), + AxisMapping(axis="dim", local_start=0, external_start=0, length=128), ), ) v_region = replace(k_region, field="v", byte_start=16384, byte_end=32768) layout = StateTransferLayout( - resource_id="main_gqa_fp16_state", + state_name="main_gqa_fp16_state", tokens_per_page=8, byte_start=0, byte_end=32768, @@ -1169,11 +1187,11 @@ layout = StateTransferLayout( 例如当前 rank 的本地 head 0、token 在页内位置 0、物理层 0、dim 0，对应 K 区域地址偏移为 `0 + 2 * 256 = 512` 字节，V 为 `16384 + 512 = 16896` 字节。CPU buffer 每行 `65536` 字节，选择下一行即进入下一物理页；本 State 的 `byte_end - byte_start = 32768` 只表示自身占用。 -D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸，把本地 head 映射改成 `AxisSlice(axis="head", src_start=0, dst_start=3, length=1)`。P/D 根据各自的 TP 配置生成切片，完整 layout 对象可以不同。 +D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸，把本地 head 映射改成 `AxisMapping(axis="head", local_start=0, external_start=3, length=1)`。P/D 根据各自的 TP 配置生成切片，完整 layout 对象可以不同。 字段列表按实际物理表示组织。例如 K、V、scale 可分别描述；packed 数据也可整体描述为一个字节区域。别名区域只搬运一次。Conv 的 Q/K 和 V、SSM 可能采用不同 head 映射，应分别描述，不能把一个资源的 `head_start` 套用于所有字段。数据类型、量化组粒度及索引单位由对应格式约束。 -`AttResourceManager`从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 +`AttentionStateManager`从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 FullAttentionState 的 token 页操作要求 `tokens_per_page` 为正整数。CPU 参数 `cpu_page_buffers` 为完整二维页池；PD 参数 `pd_page_buffers` 同样为页集合，`pd_page_index` 指定物理页。传入完整页集合，不提前切到单页或 State 区域。对于常规 K/V，一个外部地址可以写成： @@ -1216,13 +1234,13 @@ DP 请求共享信息分为两个数组，均按请求在当前批次中的位 大页列表第 `i` 项对应前缀长度 `(i + 1) * big_page_tokens`。主 TP rank 将 radix 已命中的大页与请求持有的大页按对应边界写入，每次先把当前行清为 `-1`，避免较短请求复用共享行时读到旧页 ID。发布和读取沿用节点内 barrier，源请求引用保持到所有复制完成。 -`TransTask` 使用 `src_token_indexes` / `dst_token_indexes`、`src_big_page_ids` / `dst_big_page_ids` 明确配对关系；`dst_req` 持有目标请求，`src_req_idx` 指定来源运行态，`src_dp_rank` 以节点内 DP 编号定位对应 TP 分片的远端 manager 和请求索引表；节点内进程 rank 只在初始化加载 IPC 资源时使用。目标 token 索引来自 CPU allocator，在调用资源接口前一次性转成 CUDA int64。 +`TransTask` 使用 `src_token_indices` / `dst_token_indices`、`src_big_page_ids` / `dst_big_page_ids` 明确配对关系；`dst_req` 持有目标请求，`src_req_idx` 指定来源运行态，`src_dp_rank` 以节点内 DP 编号定位对应 TP 分片的远端 manager 和请求索引表；节点内进程 rank 只在初始化加载 IPC 资源时使用。目标 token 索引来自 CPU allocator，在调用资源接口前一次性转成 CUDA int64。 目标 DP 只复制从本地已有 KV 长度到来源 KV 长度之间新增的完整大页，不重复复制已有前缀页。先申请目标大页再申请 KV；目标大页分配必须成功，通过断言检查。复制完成后，将目标页按前缀长度登记到 `req.hybrid_len_to_big_page_id`，后续 radix 插入和释放沿用请求的现有归属管理。未跨越新大页边界时，大页列表为空。 -DP 副本之间使用相同的 State 成员名。`copy_from_other_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.resource_id)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_run_time_state_from_other_dp`，再清零目标 tail。调用方保证源 MTP 游标为 0、前缀处于 tail 可重置边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 +DP 副本之间使用相同的 State 成员名。`copy_from_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.state_name)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_runtime_state_from_dp`，再清零目标 tail。调用方保证源 MTP 游标为 0、前缀处于 tail 可重置边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 -full-attention 使用 `copy_tokens_from_other_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_page_state_from_other_dp` 与 `copy_run_time_state_from_other_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 +full-attention 使用 `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_pages_from_dp` 与 `copy_runtime_state_from_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 full-attention DP 的 `dp_transfer.py` 提供显式存储入口：`copy_gqa_tokens` 接收 K/V，供 GQA 和静态 FP8 使用；`copy_group_quantized_tokens` 另接收 K/V 动态 scale；`copy_mla_tokens` 接收完整 latent/packed 存储。三种入口共用四维 token 布局的 Triton kernel，一次 launch 搬运当前 State 的全部字段；编译时展开 Tensor tuple 为指针参数，不创建或上传指针表。按两端 token 索引及各自真实 stride 直接寻址，无 `index_select`、跨设备 payload `.to()` 或 `index_copy_` 中间存储。每个字段独立选择对齐的整数搬运宽度，支持两端不同 token 容量、非连续索引和存储视图；默认 grid=24、BLOCK=4096、warps=8、stages=1，通过跳步循环覆盖全部元素，地址计算使用 int64。静态 FP8 标定参数不搬运；packed MLA 按原 DP 语义复制完整行，包括 padding，区别于 CPU/PD 的有效字段页格式。调用方保证目标索引唯一、源目标区域不重叠，并负责远端 peer/IPC 映射、源数据就绪及存储生命周期；算子在目标设备当前流执行，不在内部建立 peer 映射或同步。 @@ -1262,7 +1280,7 @@ packed MLA 等多字段布局的 Python 准备及逐字段发射开销仍会限 DP 直接复制两端 State 的字段，buffer 基址、容量和 stride 从各自 State 获取。层维度完整、同序复制；字段类型、形状、编码及固定参数的检查仍由具体实现完成。不保存额外的复制映射。 -DP 只复制对应 TP rank 的完整分片，不支持通过 State 接口拼接多个 TP 分片。不同 token 来自不同 DP 时，先按来源分组索引，再逐项调用 `copy_tokens_from_other_dp`；同一批内配对源与目标 token 索引。 +DP 只复制对应 TP rank 的完整分片，不支持通过 State 接口拼接多个 TP 分片。不同 token 来自不同 DP 时，先按来源分组索引，再逐项调用 `copy_tokens_from_dp`；同一批内配对源与目标 token 索引。 源与目标布局完全一致、分片一一对应时，可以采用同构布局的复制 kernel，但须先确认源目标 stride 等条件成立。重分片或编码变化由具体子类明确声明支持并提供实现。 @@ -1286,7 +1304,7 @@ CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索 布局由部署配置和 State 生成；搬运时校验资源名称、字段覆盖、dtype、字节对齐、页内地址边界及层/轴映射范围，TP 切片指定分片范围与副本写入者。各资源的实际写入区域不得相互覆盖；packed 别名按实际物理范围去重。涉及多个布局的描述以实际访问地址为准，不能仅比较 shape。 -数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。`AttResourceManager`校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 +数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。`AttentionStateManager`校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 ### 6.6 子类实现与直接调用示例 @@ -1294,7 +1312,7 @@ CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索 ```python class GQAFullAttentionState(FullAttentionState): - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices, pd_page_buffers, *, pd_page_index: int, layout, ) -> None: gqa_kv_to_pd_buffer( @@ -1305,7 +1323,7 @@ class GQAFullAttentionState(FullAttentionState): layout=layout, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers, token_indices, *, pd_page_index: int, layout, ) -> None: gqa_pd_buffer_to_kv( @@ -1333,10 +1351,10 @@ state.read_from_cpu( # 后续计算与写入建立 stream/event 依赖。 ``` -DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由`AttResourceManager` 组织： +DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由`AttentionStateManager` 组织： ```python -dst_state.copy_tokens_from_other_dp( +dst_state.copy_tokens_from_dp( src_token_indices=src_token_indices, dst_token_indices=dst_token_indices, remote_dp_full_att_state=src_state, @@ -1346,7 +1364,7 @@ dst_state.copy_tokens_from_other_dp( PD 两端分别传入完整发送/接收页集合与本地页索引；两端页索引可以不同，区域与 TP 位置来自各自 layout： ```python -p_state.write_to_pd_buffer( +p_state.write_to_pd_page( src_token_indices, send_pages, pd_page_index=send_page_index, layout=p_layout, @@ -1354,7 +1372,7 @@ p_state.write_to_pd_buffer( # 全部参与资源/分片的打包完成后，调度层发送。 # D 端接收完成并建立 stream 依赖后执行。 -d_state.read_from_pd_buffer( +d_state.read_from_pd_page( recv_pages, dst_token_indices, pd_page_index=recv_page_index, layout=d_layout, @@ -1366,7 +1384,7 @@ d_state.read_from_pd_buffer( ### 6.7 BoundedAttentionState：有界状态的统一操作接口 -`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。`AttResourceManager`直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 +`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。`AttentionStateManager`直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 ```text BoundedAttentionState # 抽象操作契约 @@ -1434,21 +1452,21 @@ class BoundedAttentionState(ABC): def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... @abstractmethod - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: ... @abstractmethod - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: ... @abstractmethod @@ -1461,13 +1479,13 @@ class BoundedAttentionState(ABC): def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_page_state_to_cpu( - self, page_kind: PageKind, slot_ids, cpu_page_buffers, cpu_page_indices, - *, layout: StateTransferLayout, cpu_page_readies: List[bool] | None = None, + def write_state_pages_to_cpu( + self, page_kind: PageKind, page_slot_ids, cpu_page_buffers, cpu_page_indices, + *, layout: StateTransferLayout, cpu_page_ready_flags: List[bool] | None = None, ) -> None: ... @abstractmethod - def read_page_state_from_cpu( + def read_state_pages_from_cpu( self, cpu_page_buffers, cpu_page_indices, *, big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: ... @@ -1476,13 +1494,13 @@ class BoundedAttentionState(ABC): # 四、DP 迁移 # ======================================================================== @abstractmethod - def copy_big_page_state_from_other_dp( + def copy_big_pages_from_dp( self, src_page_ids: List[int], dst_page_ids: List[int], remote_dp_bounded_state: BoundedAttentionState, ) -> None: ... @abstractmethod - def copy_run_time_state_from_other_dp( + def copy_runtime_state_from_dp( self, src_req_indices: List[int], dst_req_indices: List[int], remote_dp_bounded_state: BoundedAttentionState, ) -> None: ... @@ -1494,7 +1512,7 @@ class BoundedAttentionState(ABC): def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -1504,7 +1522,7 @@ class BoundedAttentionState(ABC): ) -> None: ... @abstractmethod - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, @@ -1514,11 +1532,11 @@ class BoundedAttentionState(ABC): ) -> None: ... ``` -bounded State 使用 `write_page_state_to_cpu` / `read_page_state_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 +bounded State 使用 `write_state_pages_to_cpu` / `read_state_pages_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `page_slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 -CPU 页读写的 `cpu_page_indices`、写出的 `slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`state_token_lens`、`slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。`state_token_lens` 表示各快照已经执行到的绝对 token 长度，不包含新采样但尚未执行的 token，也不是本批搬运的 token 数。调用方保证长度与选定快照一致；页池不新增长度记录。KDA/GDN 的 Conv/SSM 布局不依赖长度定位，保留此参数供 ring 等实现使用。 +CPU 页读写的 `cpu_page_indices`、写出的 `page_slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`state_token_lens`、`page_slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。`state_token_lens` 表示各快照已经执行到的绝对 token 长度，不包含新采样但尚未执行的 token，也不是本批搬运的 token 数。调用方保证长度与选定快照一致；页池不新增长度记录。KDA/GDN 的 Conv/SSM 布局不依赖长度定位，保留此参数供 ring 等实现使用。 -PD 状态页按单请求处理，bounded 和 tail 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_other_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 +PD 状态页按单请求处理，bounded 和 tail 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`page_slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 #### 运行态与推理访问 @@ -1533,7 +1551,7 @@ PD 状态页按单请求处理，bounded 和 tail 的读写均接收 `state_toke `accepted_steps` 为逐请求接受的、已经执行过 forward 的输入步数，包含本轮基准输入步，不是“接受的 draft token 数”或新采样输出数。取值范围为 `[1, 本轮实际验证宽度]`；KDA/GDN 的原生索引为 `accepted_steps - 1`。实现不应把不存在的零步回滚当作有效操作。未启用 MTP 时，合法的一步提交可为无数据移动的操作。 -GDN/KDA Backend 不保存 State 引用。prefill/decode 按角色、算法和逻辑层号调用管理器的 `get_att_state()`；MTP decode 的 `init_state()` 在需要游标时局部获取当前角色唯一的对应算法 State，将 `mtp_state_index` 显式传给 fixed/dynamic 初始化函数。此初始化仍要求每个角色只有一个对应算法 State，不假定多个 State 共用游标。 +GDN/KDA Backend 不保存 State 引用。prefill/decode 按角色、算法和逻辑层号调用管理器的 `get_state()`；MTP decode 的 `init_state()` 在需要游标时局部获取当前角色唯一的对应算法 State，将 `mtp_state_index` 显式传给 fixed/dynamic 初始化函数。此初始化仍要求每个角色只有一个对应算法 State，不假定多个 State 共用游标。 KDA 的提交可以只更新 `mtp_state_index`，下一轮 kernel 根据该游标选择输入状态。当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0，直接读取 Conv 的 `[:W]` 窗口和 SSM 的 `req_idx*(S+1)` 基准槽位。导入时写到目标的基准位置，并将游标设为 0。实现中以 TODO 标明：若以后支持 MTP 提交后传输，需要按已接受索引 j 选择 Conv 的 `[j:j+W]` 窗口和 SSM 的 `req_idx*(S+1)+j` 槽位。 @@ -1543,8 +1561,8 @@ prefill 的最终状态、普通 decode 的结果和已提交验证状态都可 | 接口 | 必需语义 | | --- | --- | -| `save_bounded_att_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req、state_token_lens 与 slot 逐项对齐，ring 实现按长度确定窗口；前缀关联由调用方维护 | -| `restore_bounded_att_pages(...)` | 按 state_token_lens 从选定槽位恢复运行态及 ring 位置等控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | +| `save_bounded_att_runtime_to_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req、state_token_lens 与 slot 逐项对齐，ring 实现按长度确定窗口；前缀关联由调用方维护 | +| `restore_bounded_att_runtime_from_pages(...)` | 按 state_token_lens 从选定槽位恢复运行态及 ring 位置等控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | 调用方负责将快照关联到实际执行后的位置；KDA/GDN 的页保存契约只覆盖传入请求的 prefill 完成状态；`state_token_lens` 必须是该实际状态的绝对长度，传入其他长度不能生成历史快照。chunk 跨过中间边界但没有保留对应运行态时，预留槽位不能作为已保存的页面发布。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 @@ -1554,30 +1572,30 @@ prefill 的最终状态、普通 decode 的结果和已提交验证状态都可 #### CPU、DP、PD 的状态单位 -CPU 页与本地大小页池之间使用 `get_cpu_transfer_regions()` 返回的映射：`region.field` 选择页池中的 Conv/SSM 字段，`byte_start/byte_end` 和 `shape` 定位 CPU 全局字段，`axis_slices` 将本地 TP 分片对应到全局 head/channel 范围。`src_start` 固定指本地字段，`dst_start` 固定指外部页，实际读回时只反转搬运方向，不交换两者含义。`read_page_state_from_cpu` 显式准备目标 `pool_page_fields`，共用搬运流程直接展开这些映射，不通过 `_copy_pair` 隐藏切片过程。 +CPU 页与本地大小页池之间使用 `get_cpu_transfer_regions()` 返回的映射：`region.field` 选择页池中的 Conv/SSM 字段，`byte_start/byte_end` 和 `shape` 定位 CPU 全局字段，`axis_mappings` 将本地 TP 分片对应到全局 head/channel 范围。`local_start` 固定指本地字段，`external_start` 固定指外部页，实际读回时只反转搬运方向，不交换两者含义。`read_state_pages_from_cpu` 显式准备目标 `pool_page_fields`，共用搬运流程直接展开这些映射，不通过 `_copy_pair` 隐藏切片过程。 窗口顺序的约定为：请求 ring → 按长度保存为逻辑时间顺序的大小页快照 → CPU 页；恢复时为 CPU 页 → 大页快照（完整页用普通槽位，碎页用 LOAD 槽位）→ 按绝对长度恢复请求 ring。CPU 页与页池之间只处理 TP 映射，不执行环形展开或恢复。 | 路径 | 传输单位与接口职责 | | --- | --- | -| CPU 写出 | `write_page_state_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，长度由缓存任务携带，不写入页内；每个待发布页的参与状态必须完整；CPU pool 到 CPU 页同样是独立的格式搬运 | -| CPU 读回 | State 的 `read_page_state_from_cpu` 按 `big_page_slot_ids` 将快照写入大页池；管理器补上碎页 LOAD 槽位，并调用 `restore_bounded_att_pages` 恢复最后一页运行态 | -| DP 大页复制 | `copy_big_page_state_from_other_dp(src_page_ids, dst_page_ids, remote_dp_bounded_state)` 在同一 TP 分片的 DP 副本之间复制大页 Conv/SSM 快照；不分配页、不修改请求运行态或小页，前缀关联与源页生命周期由调用方维护 | -| DP 运行态复制 | `copy_run_time_state_from_other_dp(src_req_indices, dst_req_indices, remote_dp_bounded_state)` 复制同一 TP 分片的 prefill 基准 Conv/SSM；KDA/GDN 源游标必须为 0，目标游标归零，不修改大小页或其他候选区 | -| PD 写出 | `write_to_pd_buffer` 将一个请求的 prefill 基准状态写入 `pd_page_index` 指定的状态页区域，按协议关联其前缀长度 | -| PD 读回 | `read_from_pd_buffer` 按 state_token_len 将完整快照恢复到指定请求槽位及对应 ring 位置，不读取页内长度、版本字段；不自动生成历史 checkpoint | +| CPU 写出 | `write_state_pages_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，长度由缓存任务携带，不写入页内；每个待发布页的参与状态必须完整；CPU pool 到 CPU 页同样是独立的格式搬运 | +| CPU 读回 | State 的 `read_state_pages_from_cpu` 按 `big_page_slot_ids` 将快照写入大页池；管理器补上碎页 LOAD 槽位，并调用 `restore_bounded_att_runtime_from_pages` 恢复最后一页运行态 | +| DP 大页复制 | `copy_big_pages_from_dp(src_page_ids, dst_page_ids, remote_dp_bounded_state)` 在同一 TP 分片的 DP 副本之间复制大页 Conv/SSM 快照；不分配页、不修改请求运行态或小页，前缀关联与源页生命周期由调用方维护 | +| DP 运行态复制 | `copy_runtime_state_from_dp(src_req_indices, dst_req_indices, remote_dp_bounded_state)` 复制同一 TP 分片的 prefill 基准 Conv/SSM；KDA/GDN 源游标必须为 0，目标游标归零，不修改大小页或其他候选区 | +| PD 写出 | `write_to_pd_page` 将一个请求的 prefill 基准状态写入 `pd_page_index` 指定的状态页区域，按协议关联其前缀长度 | +| PD 读回 | `read_from_pd_page` 按 state_token_len 将完整快照恢复到指定请求槽位及对应 ring 位置，不读取页内长度、版本字段；不自动生成历史 checkpoint | -CPU 尾页若需要保存 prefill 完成后的运行态，管理器先使用保留槽位调用 `save_bounded_att_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 +CPU 尾页若需要保存 prefill 完成后的运行态，管理器先使用保留槽位调用 `save_bounded_att_runtime_to_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字段映射，直接复制同一算法、同一配置及 TP 分片的对应存储。Conv 区域可用 `(layer, channel, window)` 轴，SSM 用 `(layer, head, key_dim, value_dim)` 轴。它们没有历史 KV 的 token 轴。含 KV 的 CPU 页仍有 token 容量，checkpoint 组件不按该容量重复展开；纯 PD 状态页的 `tokens_per_page` 为 None。各资源偏移、TP 范围和副本写入者由管理器组合，State 直接执行方法。 当前 DP/PD 导出 prefill 基准状态，调用方保证 KDA/GDN 源游标为 0，不传 MTP 候选区，也不支持 MTP 提交后导出。目标可按自己的 MTP 步数初始化工作区；DP 两端配置及 TP 分片一致，按各自 MTP 步数计算基准槽位。PD 则通过全局布局与 TP 切片处理两端分片差异。不能因为 KDA/GDN 的 Conv/SSM 形状相同就跨算法恢复。 -多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；`AttResourceManager`完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 +多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；`AttentionStateManager`完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 #### 依赖、同步与使用示例 -资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttResourceManager`单独协调 tail 在恢复边界上的重置或迁移。 +资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttentionStateManager`单独协调 tail 在恢复边界上的重置或迁移。 GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 @@ -1585,19 +1603,19 @@ GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制 # ReqManager 已分配请求槽位；调用方已从管理器持有的 allocator 分配状态页槽位。 state.reset_bounded_att_state(req_indices) # backend 经层视图完成 prefill；调用方维护实际 chunk 终点。 -state.save_bounded_att_pages( +state.save_bounded_att_runtime_to_pages( req_indices=req_indices, state_token_lens=state_token_lens, page_kind="big", - slot_ids=slot_ids, + page_slot_ids=page_slot_ids, ) # 确认所有参与状态已按指定长度保存并完成传输后，缓存调用方发布联合页。 -state.restore_bounded_att_pages( +state.restore_bounded_att_runtime_from_pages( req_indices=dst_req_indices, state_token_lens=state_token_lens, page_kind="big", - slot_ids=slot_ids, + page_slot_ids=page_slot_ids, ) # target 请求在相关 KV、linear、tail 均就绪后继续执行。 @@ -1656,12 +1674,12 @@ class AttentionTailState(ABC): def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx, pd_page_buffers, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, ) -> None: ... @abstractmethod - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers, req_idx, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, ) -> None: ... ``` @@ -1685,14 +1703,14 @@ class AttentionTailState(ABC): | 接口 | 数据路径及约束 | | --- | --- | -| `write_to_pd_buffer(...)` | 按 `state_token_len` 提取指定请求的逻辑窗口 → `pd_page_index` 指定的 PD 状态页区域，使用 layout 的资源偏移与 TP 写入规则 | -| `read_from_pd_buffer(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；窗口外的 MTP 槽位保持原值，调用方建立接收完成依赖 | +| `write_to_pd_page(...)` | 按 `state_token_len` 提取指定请求的逻辑窗口 → `pd_page_index` 指定的 PD 状态页区域，使用 layout 的资源偏移与 TP 写入规则 | +| `read_from_pd_page(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；窗口外的 MTP 槽位保持原值，调用方建立接收完成依赖 | 这些方法传输 tail 自身的 prefill 逻辑窗口，不承担历史 KV、池化 key/scale、bounded 状态或 MTP proposal 的迁移。发送端取最近 `P` 个位置并按时间顺序排列；接收端同样按逻辑窗口遍历，单次 kernel 恢复窗口数据及短前缀补零，窗口之外的 MTP 槽位与其他请求保持不变。后续 MTP 候选历史由 D 端执行过程中产生。 `StateTransferLayout` 描述外部页位置与切片。K-pool region 的轴为 `("layer", "window", "channel")`，形状为 `[layers, P, 2*head_dim]`，与本地 `mtp_step` 无关。`window` 按时间顺序对应 `[state_token_len-P, state_token_len)`；短前缀中小于 0 的位置在发送和恢复时均补零，窗口外槽位保持原值。独立状态页的 `tokens_per_page=None`，该窗口不是 full-attention 的页内 token 轴。 -K-pool 的 `window` 切片设置 `src_circular=True`。`kpool_transfer.py` 从 `state_token_len` 计算本次本地窗口起点；可通过 `axis_slice_overrides` 覆盖本次轴映射，不修改缓存的 layout。窗口起点及短前缀的无效位置数量作为运行时参数传入定制 Triton kernel，不随请求位置编译新变体。写出和读回各一次 launch，处理量均为 `layers * pool_size * channels`，不遍历窗口之外的 MTP 槽位；连续 channel 在对齐允许时使用最大 8 字节整数视图搬运，保留原始位模式。 +K-pool 的 `window` 切片设置 `local_circular=True`。`kpool_transfer.py` 从 `state_token_len` 计算本次本地窗口起点；可通过 `axis_mapping_overrides` 覆盖本次轴映射，不修改缓存的 layout。窗口起点及短前缀的无效位置数量作为运行时参数传入定制 Triton kernel，不随请求位置编译新变体。写出和读回各一次 launch，处理量均为 `layers * pool_size * channels`，不遍历窗口之外的 MTP 槽位；连续 channel 在对齐允许时使用最大 8 字节整数视图搬运，保留原始位模式。 P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_step`。位置 `pos` 在 P 端从 `pos % (P+S_p)` 读取，在 D 端写入 `pos % (P+S_d)`；不能直接截取 ring 的前 P 个物理槽位。模型算法、池化步长、层顺序、字段精度和逻辑维度仍须兼容。现有 main/draft prefill 保持相同的序列位置，因此可共用任务的 `end_kv_index`；位置变化不影响缓存的布局描述。 @@ -1702,7 +1720,7 @@ P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_s #### 与 linear 页接口和 MTP 提交的区别 -- 不提供 `save_bounded_att_pages`、`restore_bounded_att_pages`：tail 没有独立大小页 pool。CPU/DP 缓存前缀在对齐边界通过 reset 恢复；非对齐的 PD 交接使用 prefill 逻辑窗口。 +- 不提供 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages`：tail 没有独立大小页 pool。CPU/DP 缓存前缀在对齐边界通过 reset 恢复；非对齐的 PD 交接使用 prefill 逻辑窗口。 - 不要求 `commit_bounded_att_state`：K-pool 没有类似 KDA 的 SSM 选择游标，按绝对位置访问 ring，MTP 拒绝后由执行流程选择后续位置。增加空的提交方法不能替代窗口容量与重访范围的正确性；未来出现确实需要提交控制状态的实现时，再扩展对应契约。 ## 7. 资源初始化与调用约定 @@ -1715,17 +1733,17 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 | --- | --- | | `get_req_manager_max_seq_len()` | 返回请求映射所需长度容量，取本模型 `batch_max_tokens`、`max_seq_length` 与 0 的最大值；未配置的值按 0 处理，不表示模型最大上下文能力 | | `bind_req_manager(req_manager)` | 绑定调用方创建的共享实例，不判断 main/draft，也不通过 `main_model` 获取资源 | -| `prepare_att_resources()` | main 直接创建空的 `AttResourceManager`，draft 引用 main 的实例；在 KV 内存分区内调用三类 State 工厂并统一注册返回的列表 | +| `prepare_attention_states()` | main 直接创建空的 `AttentionStateManager`，draft 引用 main 的实例；在 KV 内存分区内调用三类 State 工厂并统一注册返回的列表 | | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 -2. backend 按 main、各 draft 的顺序分别调用 `prepare_att_resources()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `create_full_att_states()`、`create_bounded_att_states()`、`create_att_tail_states()`，再通过 `set_att_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 -3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_creation()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 -4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 必须有 K/V 标定，Q 标定可选：有 Q 标定时使用静态量化，未提供时按本次请求的 Q 分组动态量化，prefill/decode 使用同一规则。随后 backend 调用 `initialize_resources(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 +2. backend 按 main、各 draft 的顺序分别调用 `prepare_attention_states()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `create_full_att_states()`、`create_bounded_att_states()`、`create_att_tail_states()`，再通过 `register_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 +3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_registration()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 +4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 必须有 K/V 标定，Q 标定可选：有 Q 标定时使用静态量化，未提供时按本次请求的 Q 分组动态量化，prefill/decode 使用同一规则。随后 backend 调用 `initialize_storage(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 -BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `resource_id`。 +BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `state_name`。 | 创建接口 | 返回值与职责 | | --- | --- | @@ -1737,14 +1755,14 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 `Qwen3NextResourceConfig` 及其 GLM 配置只描述当前模型，保留 `from_model_config()`、`load_from_args()` 和 `get_full_att_layer_indices()` 三个入口。层查询返回模型内部序号，由模型工厂转换为实际逻辑层 ID；不追加 draft 层，不计算物理层映射、Conv/SSM 存储形状或字节数。配置只保存全局 head 数，本地 TP 分片由 State 或后端校验调用方计算。full-attention 维度保存原始特征宽度，GLM 的 indexer 字段和 packed padding 由具体 State 的创建入口处理。 -三个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_att_resources()` 为 main 直接创建空的 `AttResourceManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttResourceManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 +三个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_attention_states()` 为 main 直接创建空的 `AttentionStateManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttentionStateManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 -需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_att_resources()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 +需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_attention_states()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 `create_full_att_states()` 由具体模型直接构造 `FullAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 -`AttResourceManager.initialize_resources()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `initialize_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 +`AttentionStateManager.initialize_storage()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `allocate_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 自动容量规划使用所有模型实际加载后的剩余显存，不再按 main 权重比例预估 draft 占用。静态量化标定须在规划前加载和校验；KV、动态 scale、packed padding 及 HOLD 页计入预算，静态 KV/Q 标定参数占比较小，不在容量规划中单独扣减。按完整物理页选择可容纳的最大 token 容量，各 rank 再统一取最小容量。共享 allocator 的不同 State 存储分别计算，不能因共享索引而漏算 draft buffer。 @@ -1758,7 +1776,7 @@ PD master 负责校验连接参数，长度通过传输任务的 `end_kv_index` P/D 传输进程为每次 KV 任务创建普通 CPU token 索引 Tensor，在 `copy_cuda_stream` 中调用 `cuda(non_blocking=True)` 上传，后续 State 拷贝沿用同一流；预热使用相同路径。不预分配或保留 pinned 索引页缓冲，普通 CPU 内存也不保证 H2D 全程异步。`att_state` 任务按请求槽位搬运，不创建或上传 token 索引。 -CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_transfer_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 +CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_page_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 State 的全局存储字节数查询继续描述存储开销，不用于推导 CPU/PD 传输容量。CPU 联合页不预留页头、长度或版本字段。布局及其派生页大小均按需缓存，实际存储分配后一起清理，避免继续使用 meta 准备阶段的布局尺寸。 @@ -1788,7 +1806,7 @@ flowchart LR GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 -CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_transfer_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 +CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_page_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 @@ -1796,9 +1814,9 @@ CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 fu 没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)` 和 `reset_attention_tail_state(req_indices)`，分别初始化 bounded 与 tail 运行态。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 -调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_pages`。`AttResourceManager` 将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 +调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_runtime_to_pages`。`AttentionStateManager` 将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 -radix 沿用原有判断：大页节点具有状态，小页节点通过 `small_page_buffer_idx is not None` 判断是否可恢复；图片边界仍通过 `clamp_mm_cache_len` 独立检查。匹配得到节点引用后，调用方直接调用 `restore_bounded_att_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_attention_tail_state(req_indices)` 清零 tail，保留已恢复的 linear 状态。调用方还需确认历史 KV、tail 和 bounded 快照均对应恢复前缀，并保证数据就绪。 +radix 沿用原有判断：大页节点具有状态，小页节点通过 `small_page_buffer_idx is not None` 判断是否可恢复；图片边界仍通过 `clamp_mm_cache_len` 独立检查。匹配得到节点引用后，调用方直接调用 `restore_bounded_att_runtime_from_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_attention_tail_state(req_indices)` 清零 tail，保留已恢复的 linear 状态。调用方还需确认历史 KV、tail 和 bounded 快照均对应恢复前缀，并保证数据就绪。 使用 K-pool 的 tail 在完整池化边界上通过 `reset_attention_tail_state` 重置；非对齐恢复需要原始 tail 数据。CPU 读回后，管理器在最终前缀数据恢复后调用 `reset_attention_tail_state`，不读取 tail 快照。DP fetch 在对齐的缓存前缀上重置 tail；PD 按第 6.8 节传输固定逻辑窗口并重排到本地 ring。前缀恢复的清零规则不用于 MTP 回退。 @@ -1814,23 +1832,23 @@ MTP 接受位置由模型执行流程选择，供后续计算使用；当前大 ## 8. CPU、DP 与 PD 传输设计 -### 8.1 具体 State 实现传输，`AttResourceManager` 组织组合 +### 8.1 具体 State 实现传输，`AttentionStateManager` 组织组合 CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合缓存辅助文件： - `read_req_from_cpu(req, cpu_page_indices, page_lengths, *, start_token_len)` 为新增 KV 分配索引，补齐首个 CPU 页已有命中部分，为完整 bounded 页分配大页槽位，再调用资源管理器 `read_req_from_cpu`。成功提交后登记请求的大页引用并更新 KV 长度；CPU cache 加载按必须成功的流程执行，异常直接向上传播，不在此处捕获、同步回收或继续执行。 -- `write_req_to_cpu(req, cpu_kv_cache_stream)` 先按 KV 长度和碎页规则申请 CPU 页；实际页数确定后，合并 radix 共享大页和请求自有页，仅查找本批页面所需槽位。完整大页必须存在，直接索引获取，缺失时暴露错误，不截断缓存前缀。将整批页面、完整 token 索引、大小页槽位和 `cpu_page_readies` 一次传给资源管理器 `write_req_to_cpu`，由各 State 跳过 ready 页面。返回含完成事件的 `TransTask`，由已有任务完成流程发布 ready 并释放引用。 +- `write_req_to_cpu(req, cpu_kv_cache_stream)` 先按 KV 长度和碎页规则申请 CPU 页；实际页数确定后，合并 radix 共享大页和请求自有页，仅查找本批页面所需槽位。完整大页必须存在，直接索引获取，缺失时暴露错误，不截断缓存前缀。将整批页面、完整 token 索引、大小页槽位和 `cpu_page_ready_flags` 一次传给资源管理器 `write_req_to_cpu`，由各 State 跳过 ready 页面。返回含完成事件的 `TransTask`，由已有任务完成流程发布 ready 并释放引用。 写出入口仅通过 `cpu_kv_cache_stream.wait_stream(producer)` 等待当前生产流已提交的写入，不额外等待 overlap 流，也不对 hybrid 模型执行主机同步。TODO：bounded CPU→CPU 快照搬运仍需改为当前流上的 kernel；当前主机 copy 不受流等待约束，kernel 落地前源快照必须预先就绪。 -命中边界、容量准入和 CPU 页引用释放仍由外层缓存调度负责。以上入口不把 `InferReq` 传入 `AttResourceManager`；资源层只接收显式索引与槽位，并组合各 State 的数据操作。 +命中边界、容量准入和 CPU 页引用释放仍由外层缓存调度负责。以上入口不把 `InferReq` 传入 `AttentionStateManager`；资源层只接收显式索引与槽位，并组合各 State 的数据操作。 -外部通过 6.2 节的管理器接口执行 CPU/DP/PD 操作。`AttResourceManager` 将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 +外部通过 6.2 节的管理器接口执行 CPU/DP/PD 操作。`AttentionStateManager` 将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 - full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 - `KDAState`、`GDNState` 实现 `BoundedAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移 prefill 基准请求状态；管理器组织参与对象及位置映射。 - `AttentionTailState` 提供层访问、按请求清零及 PD 直接读写方法，详见 6.8 节；CPU cache 和 DP fetch 通过重置恢复对齐前缀的 tail。GLM K-pool indexer 在 TP ranks 间为副本，其他实现自行描述其布局。 -- 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttResourceManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 +- 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttentionStateManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 跨进程导入的管理器持有借用的运行态存储和索引容量描述，不持有 allocator 或 CPU 状态池。请求映射由 server 层作为 ReqManager 所有的独立存储导出，不能附加到资源管理器。CUDA IPC 句柄按实际读取进程数发布，每个读取者消费独立引用；原进程在全部读取者结束使用前保持存储存活。 @@ -1841,11 +1859,11 @@ PD 任务分为 `kv` 和 `att_state` 两类；涉及 main/draft 时按模型角 | 任务 | 组成 | 处理方式 | | --- | --- | --- | | `kv` | full-attention 历史 KV，包含布局所需的 packed indexer 数据 | 按 token 页复制 | -| `att_state` | linear Conv/SSM + 可选 indexer tail | 模型专属 `AttResourceManager` 按请求恢复位置调用对应 State 方法 | +| `att_state` | linear Conv/SSM + 可选 indexer tail | 模型专属 `AttentionStateManager` 按请求恢复位置调用对应 State 方法 | P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，indexer tail 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将 tail 写入目标 TP 各 rank 的本地请求槽位。 -模型专属 `AttResourceManager` 实现状态页的布局组织和容量校验，调用 linear State 与 indexer tail State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 +模型专属 `AttentionStateManager` 实现状态页的布局组织和容量校验，调用 linear State 与 indexer tail State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 网络层只接收固定容量的字节页，不按 KV 的 head 数或 dtype 推断组合布局。两端在创建传输句柄前核对页容量，PD master 在连接时校验任务类别、资源身份、算法与编码、层映射和全局字段形状；页内不附加格式指纹、长度或版本。布局使用全局字段描述以支持不同 TP 大小，tail 布局只描述固定逻辑窗口，不包含本地 MTP 候选容量。运行态的前缀长度由传输任务提供。 @@ -1857,7 +1875,7 @@ CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布 GPU→CPU 保存状态页完成后，才能在 CPU 上重排或复制这些快照；主机侧读取必须等待保存完成，不能只向另一个 CUDA stream 排入依赖。CPU 页的 ready 发布仍由调用流程在全部参与资源与 TP 分片传输完成后执行。 -模型专属 `AttResourceManager` 实现 CPU 联合页的布局与读写组织，调用所需 main/draft State 的 CPU 方法，组合 full-attention KV 和 linear 页；按调用方提供的槽位写出或读回快照，pool 和 CPU buffer 均不保存长度副本。CPU 缓存调度提供外部页和任务生命周期；每个已发布 CPU 页都具备该模型组合所需的有效 linear 状态。一次读回处理同一前缀的页面，并通过 `reset_attention_tail_state` 清零目标 tail；恢复终点必须满足具体 tail 的可重置边界。 +模型专属 `AttentionStateManager` 实现 CPU 联合页的布局与读写组织，调用所需 main/draft State 的 CPU 方法，组合 full-attention KV 和 linear 页；按调用方提供的槽位写出或读回快照，pool 和 CPU buffer 均不保存长度副本。CPU 缓存调度提供外部页和任务生命周期；每个已发布 CPU 页都具备该模型组合所需的有效 linear 状态。一次读回处理同一前缀的页面，并通过 `reset_attention_tail_state` 清零目标 tail；恢复终点必须满足具体 tail 的可重置边界。 CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大页 pool 分别提供存储。各参与方的暂存数据消费完成后，后续请求才能复用该槽位。磁盘存储使用 CPU 页格式；main/draft 异构组合需定义各自资源段的位置、字段形状及 TP 切片。 @@ -1868,7 +1886,7 @@ DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并 | 问题 | 设计约定 | 待确定内容 | | --- | --- | --- | | main/draft 管理器装配 | 模型组合实现自己的管理器；各角色和具体方案拥有唯一 State，同一分配域共享唯一 allocator 实例 | 新模型的 main/draft 装配时序及独立存储分配；调用方覆盖全部参与资源的引用解除及单次归还责任 | -| AttResourceManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、清零指定请求的 tail、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | +| AttentionStateManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、清零指定请求的 tail、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | | 内部 attention 资源的暴露 | 具体模型在初始化阶段绑定 State 或层视图，不要求统一模型视图获取接口 | 各模型/backend 的具体绑定方式与按层访问结构 | | FullAttentionState 接口与继承 | 必需语义：容量计算、初始化、层视图、KV 写入、完整 token 复制及固定字段传输布局；CPU/DP/PD 由子类按支持路径直接实现读写方法 | 新布局的 TensorField 覆盖范围、StateTransferLayout 映射约束及具体搬运实现 | | BoundedAttentionState 接口与继承 | KDAState、GDNState 实现统一抽象操作；已接受运行态、checkpoint 和传输分别定义，allocator 仍归管理器 | 按 6.7 节细化各子类的描述字段和 kernel，核对 MTP 状态选择、跨 TP/CPU 格式兼容性 | diff --git a/lightllm/common/att_resources/__init__.py b/lightllm/common/att_resources/__init__.py deleted file mode 100644 index 428c8aa2a1..0000000000 --- a/lightllm/common/att_resources/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Attention data resources, independent of request and cache policy.""" - -from .att_resource_manager import AttResourceManager - -__all__ = ["AttResourceManager"] diff --git a/lightllm/common/attention_state/__init__.py b/lightllm/common/attention_state/__init__.py new file mode 100644 index 0000000000..6c3f3107c1 --- /dev/null +++ b/lightllm/common/attention_state/__init__.py @@ -0,0 +1,5 @@ +"""Attention data resources, independent of request and cache policy.""" + +from .manager import AttentionStateManager + +__all__ = ["AttentionStateManager"] diff --git a/lightllm/common/att_resources/allocators/__init__.py b/lightllm/common/attention_state/allocators/__init__.py similarity index 100% rename from lightllm/common/att_resources/allocators/__init__.py rename to lightllm/common/attention_state/allocators/__init__.py diff --git a/lightllm/common/att_resources/allocators/full_att_kv_allocator.py b/lightllm/common/attention_state/allocators/full_att_kv_allocator.py similarity index 100% rename from lightllm/common/att_resources/allocators/full_att_kv_allocator.py rename to lightllm/common/attention_state/allocators/full_att_kv_allocator.py diff --git a/lightllm/common/att_resources/allocators/state_page_allocator.py b/lightllm/common/attention_state/allocators/state_page_allocator.py similarity index 67% rename from lightllm/common/att_resources/allocators/state_page_allocator.py rename to lightllm/common/attention_state/allocators/state_page_allocator.py index 24bc53bbec..9096bd9a96 100644 --- a/lightllm/common/att_resources/allocators/state_page_allocator.py +++ b/lightllm/common/attention_state/allocators/state_page_allocator.py @@ -29,11 +29,11 @@ def alloc(self, num_pages: int) -> list[int] | None: self._allocated[slot] = 1 return slots - def free(self, slot_ids: List[int]) -> None: - assert isinstance(slot_ids, list), "slot indices must be a list" - assert len(set(slot_ids)) == len(slot_ids), "cannot free duplicate slots" - assert all(0 <= slot < self.capacity for slot in slot_ids), "slot outside the allocatable domain" - assert all(self._allocated[slot] for slot in slot_ids), "cannot free an unallocated slot" - for slot in slot_ids: + def free(self, page_slot_ids: List[int]) -> None: + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert len(set(page_slot_ids)) == len(page_slot_ids), "cannot free duplicate slots" + assert all(0 <= slot < self.capacity for slot in page_slot_ids), "slot outside the allocatable domain" + assert all(self._allocated[slot] for slot in page_slot_ids), "cannot free an unallocated slot" + for slot in page_slot_ids: self._allocated[slot] = 0 - self._free.extend(slot_ids) + self._free.extend(page_slot_ids) diff --git a/lightllm/common/att_resources/att_tail/__init__.py b/lightllm/common/attention_state/att_tail/__init__.py similarity index 100% rename from lightllm/common/att_resources/att_tail/__init__.py rename to lightllm/common/attention_state/att_tail/__init__.py diff --git a/lightllm/common/att_resources/att_tail/base_state.py b/lightllm/common/attention_state/att_tail/base_state.py similarity index 94% rename from lightllm/common/att_resources/att_tail/base_state.py rename to lightllm/common/attention_state/att_tail/base_state.py index 2aa784e123..e3e300cd87 100644 --- a/lightllm/common/att_resources/att_tail/base_state.py +++ b/lightllm/common/attention_state/att_tail/base_state.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any from abc import ABC, abstractmethod -from lightllm.common.att_resources.types import ( +from lightllm.common.attention_state.types import ( StateTransferLayout, TransferRegion, ) @@ -25,7 +25,7 @@ def get_state_type(self) -> str: 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 @@ -68,7 +68,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -80,7 +80,7 @@ def write_to_pd_buffer( ... @abstractmethod - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, diff --git a/lightllm/common/att_resources/att_tail/create_kpool_state.py b/lightllm/common/attention_state/att_tail/create_kpool_state.py similarity index 100% rename from lightllm/common/att_resources/att_tail/create_kpool_state.py rename to lightllm/common/attention_state/att_tail/create_kpool_state.py diff --git a/lightllm/common/att_resources/att_tail/kpool_state.py b/lightllm/common/attention_state/att_tail/kpool_state.py similarity index 90% rename from lightllm/common/att_resources/att_tail/kpool_state.py rename to lightllm/common/attention_state/att_tail/kpool_state.py index fdb4330f0c..6c50c7cc0b 100644 --- a/lightllm/common/att_resources/att_tail/kpool_state.py +++ b/lightllm/common/attention_state/att_tail/kpool_state.py @@ -6,12 +6,12 @@ import torch from ..utils import dtype_name -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import AttentionTailState from .kpool_transfer import transfer_kpool_pages if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout @dataclass(frozen=True) @@ -60,7 +60,7 @@ def __init__( self.config = config self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step # ring 额外保留的 MTP 历史槽位数。 - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) @@ -116,18 +116,18 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: dtype=config.dtype, axes=("layer", "window", "channel"), shape=(len(config.logical_layer_ids), window_size, channels), - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="window", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=window_size, - src_circular=True, + local_circular=True, ), - AxisSlice( + AxisMapping( axis="channel", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=channels, ), ), @@ -135,7 +135,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ) return (tail_region,) - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -144,7 +144,7 @@ def write_to_pd_buffer( layout: StateTransferLayout, ) -> None: """将 prefill 末尾的逻辑窗口写入 PD 页；state_token_len 为已执行的前缀长度。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" transfer_kpool_pages( tail=self.tail_buffer[:, req_idx], page_buffers=pd_page_buffers, @@ -154,7 +154,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, @@ -163,7 +163,7 @@ def read_from_pd_buffer( layout: StateTransferLayout, ) -> None: """按前缀长度将逻辑窗口放回本地 ring，保留窗口外槽位，允许两端 MTP 步数不同。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" transfer_kpool_pages( tail=self.tail_buffer[:, req_idx], page_buffers=pd_page_buffers, diff --git a/lightllm/common/att_resources/att_tail/kpool_transfer.py b/lightllm/common/attention_state/att_tail/kpool_transfer.py similarity index 85% rename from lightllm/common/att_resources/att_tail/kpool_transfer.py rename to lightllm/common/attention_state/att_tail/kpool_transfer.py index 06aed5d7a7..490774cc53 100644 --- a/lightllm/common/att_resources/att_tail/kpool_transfer.py +++ b/lightllm/common/attention_state/att_tail/kpool_transfer.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from ..types import AxisSlice, StateTransferLayout +from ..types import AxisMapping, StateTransferLayout # ======================================================================== @@ -19,14 +19,14 @@ def transfer_kpool_pages( state_token_len: int, layout: StateTransferLayout, write: bool, - axis_slice_overrides: tuple[AxisSlice, ...] = (), + axis_mapping_overrides: tuple[AxisMapping, ...] = (), grid_num: int = 24, ) -> None: """在当前 CUDA stream 交换一个请求的 tail 和选定 PD 页，不创建中间 Tensor。 本地 tail 为 [layer, ring_slot, channel]，ring 容量为 pool_size + 本地 mtp_step； 外部页保存最近 pool_size 个逻辑位置，按时间顺序排列。P/D 的 ring 容量可以不同。 - 窗口起点由 state_token_len 推导；axis_slice_overrides 可覆盖本次 window 映射， + 窗口起点由 state_token_len 推导；axis_mapping_overrides 可覆盖本次 window 映射， 其余轴沿用缓存布局。K-pool 只支持 window 环形寻址，channel 保持完整、同序对应。 两个方向都只遍历逻辑窗口，短前缀对应的窗口内位置补零； 本地窗口之外的 MTP 槽位保持原值，不在传输时清理。 @@ -48,34 +48,34 @@ def transfer_kpool_pages( channel_count = tail.shape[2] # 阶段二：确定本次窗口。显式覆盖的起点直接使用；否则从 prefill 长度推导。 - # 只计算整数位置，不修改 layout，也不创建新的 AxisSlice 对象。 - layout_slices = {item.axis: item for item in region.axis_slices} - slice_overrides = {item.axis: item for item in axis_slice_overrides} - window_slice = slice_overrides.get("window", layout_slices["window"]) - channel_slice = slice_overrides.get("channel", layout_slices["channel"]) - window_size = window_slice.length - local_window_start = window_slice.src_start - external_window_start = window_slice.dst_start - - if window_slice.src_circular: - if "window" not in slice_overrides: + # 只计算整数位置，不修改 layout，也不创建新的 AxisMapping 对象。 + layout_mappings = {item.axis: item for item in region.axis_mappings} + mapping_overrides = {item.axis: item for item in axis_mapping_overrides} + window_mapping = mapping_overrides.get("window", layout_mappings["window"]) + channel_mapping = mapping_overrides.get("channel", layout_mappings["channel"]) + window_size = window_mapping.length + local_window_start = window_mapping.local_start + external_window_start = window_mapping.external_start + + if window_mapping.local_circular: + if "window" not in mapping_overrides: # 外部 window[i] 对应绝对位置 state_token_len - window_size + i。 local_window_start += state_token_len - window_size local_window_start %= local_ring_size - if window_slice.dst_circular: + if window_mapping.external_circular: external_window_start %= external_ring_size invalid_prefix = max(window_size - state_token_len, 0) # 定制 kernel 完整处理 key/gate；窗口不能重复覆盖或越过线性存储边界。 assert ( - channel_slice.src_start == channel_slice.dst_start == 0 - and channel_slice.length == channel_count - and not channel_slice.src_circular and not channel_slice.dst_circular + channel_mapping.local_start == channel_mapping.external_start == 0 + and channel_mapping.length == channel_count + and not channel_mapping.local_circular and not channel_mapping.external_circular ), "K-pool transfers require the complete channel axis" assert ( 0 < window_size <= min(local_ring_size, external_ring_size) - and (window_slice.src_circular or 0 <= local_window_start <= local_ring_size - window_size) - and (window_slice.dst_circular or 0 <= external_window_start <= external_ring_size - window_size) + and (window_mapping.local_circular or 0 <= local_window_start <= local_ring_size - window_size) + and (window_mapping.external_circular or 0 <= external_window_start <= external_ring_size - window_size) ), "window exceeds storage capacity" # 阶段三：只扩大连续 channel 的搬运单位，layer/window 维度和环形语义保持不变。 @@ -127,8 +127,8 @@ def transfer_kpool_pages( WINDOW_SIZE=window_size, CHANNELS=local_tail.shape[2], ELEMENT_BYTES=local_tail.element_size(), - SRC_CIRCULAR=window_slice.src_circular, - DST_CIRCULAR=window_slice.dst_circular, + LOCAL_CIRCULAR=window_mapping.local_circular, + EXTERNAL_CIRCULAR=window_mapping.external_circular, WRITE=write, BLOCK_ELEMENTS=block_elements, num_warps=4, @@ -155,8 +155,8 @@ def _transfer_kpool_pages( WINDOW_SIZE: tl.constexpr, CHANNELS: tl.constexpr, ELEMENT_BYTES: tl.constexpr, - SRC_CIRCULAR: tl.constexpr, - DST_CIRCULAR: tl.constexpr, + LOCAL_CIRCULAR: tl.constexpr, + EXTERNAL_CIRCULAR: tl.constexpr, WRITE: tl.constexpr, BLOCK_ELEMENTS: tl.constexpr, ): @@ -185,10 +185,10 @@ def _transfer_kpool_pages( window = element // CHANNELS % WINDOW_SIZE layer = element // (CHANNELS * WINDOW_SIZE) local_slot = local_window_start + window - if SRC_CIRCULAR: + if LOCAL_CIRCULAR: local_slot %= LOCAL_RING_SIZE external_slot = external_window_start + window - if DST_CIRCULAR: + if EXTERNAL_CIRCULAR: external_slot %= EXTERNAL_RING_SIZE has_history = (window >= invalid_prefix) & (window < WINDOW_SIZE) diff --git a/lightllm/common/att_resources/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/__init__.py rename to lightllm/common/attention_state/bounded_att/__init__.py diff --git a/lightllm/common/att_resources/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py similarity index 89% rename from lightllm/common/att_resources/bounded_att/base_state.py rename to lightllm/common/attention_state/bounded_att/base_state.py index 635b68edb8..d65d6bba83 100644 --- a/lightllm/common/att_resources/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -8,7 +8,7 @@ import torch -from lightllm.common.att_resources.types import ( +from lightllm.common.attention_state.types import ( StateTransferLayout, PageKind, TransferRegion, @@ -28,7 +28,7 @@ def get_state_type(self) -> str: 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 @@ -71,12 +71,12 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... @abstractmethod - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 @@ -86,16 +86,16 @@ def save_bounded_att_pages( ... @abstractmethod - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 - req_indices、state_token_lens、slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， + req_indices、state_token_lens、page_slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， 不表示本批搬运的 token 数，调用方保证与所选快照的前缀一致。 """ ... @@ -123,27 +123,27 @@ def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_page_state_to_cpu( + def write_state_pages_to_cpu( self, page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: - """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_readies 中为 True 的页面。""" + """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_ready_flags 中为 True 的页面。""" ... @abstractmethod - def read_page_state_from_cpu( + def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" ... # ======================================================================== @@ -151,7 +151,7 @@ def read_page_state_from_cpu( # ======================================================================== @abstractmethod - def copy_big_page_state_from_other_dp( + def copy_big_pages_from_dp( self, src_page_ids: List[int], dst_page_ids: List[int], @@ -161,7 +161,7 @@ def copy_big_page_state_from_other_dp( ... @abstractmethod - def copy_run_time_state_from_other_dp( + def copy_runtime_state_from_dp( self, src_req_indices: List[int], dst_req_indices: List[int], @@ -183,7 +183,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -198,7 +198,7 @@ def write_to_pd_buffer( ... @abstractmethod - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, diff --git a/lightllm/common/att_resources/bounded_att/conv_ssm_types.py b/lightllm/common/attention_state/bounded_att/conv_ssm_types.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/conv_ssm_types.py rename to lightllm/common/attention_state/bounded_att/conv_ssm_types.py diff --git a/lightllm/common/att_resources/bounded_att/copy_kernel.py b/lightllm/common/attention_state/bounded_att/copy_kernel.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/copy_kernel.py rename to lightllm/common/attention_state/bounded_att/copy_kernel.py diff --git a/lightllm/common/att_resources/bounded_att/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/create_gdn_state.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/create_gdn_state.py rename to lightllm/common/attention_state/bounded_att/create_gdn_state.py diff --git a/lightllm/common/att_resources/bounded_att/create_kda_state.py b/lightllm/common/attention_state/bounded_att/create_kda_state.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/create_kda_state.py rename to lightllm/common/attention_state/bounded_att/create_kda_state.py diff --git a/lightllm/common/att_resources/bounded_att/gdn_pages_pool.py b/lightllm/common/attention_state/bounded_att/gdn_pages_pool.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/gdn_pages_pool.py rename to lightllm/common/attention_state/bounded_att/gdn_pages_pool.py diff --git a/lightllm/common/att_resources/bounded_att/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn_state.py similarity index 87% rename from lightllm/common/att_resources/bounded_att/gdn_state.py rename to lightllm/common/attention_state/bounded_att/gdn_state.py index 03f0a9360c..61506911dd 100644 --- a/lightllm/common/att_resources/bounded_att/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn_state.py @@ -8,7 +8,7 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion +from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig from .gdn_pages_pool import GDNPagesPool @@ -16,7 +16,7 @@ from .copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: - from lightllm.common.att_resources.types import PageKind + from lightllm.common.attention_state.types import PageKind logger = init_logger(__name__) @@ -42,7 +42,7 @@ def __init__( """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config - self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step self._logical_to_storage_layer = { @@ -125,39 +125,39 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, slot_ids): + for req_idx, slot_id in zip(req_indices, page_slot_ids): conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, slot_ids): + for req_idx, slot_id in zip(req_indices, page_slot_ids): conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) copy_bounded_att_tensors(copies=conv_copies) @@ -183,34 +183,34 @@ def commit_bounded_att_state( def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - region.shape 描述 CPU 页中的全局字段；axis_slices 将本地页池的 TP 分片映射到该字段。 + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 """ return self._get_transfer_regions() - def write_page_state_to_cpu( + def write_state_pages_to_cpu( self, page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """将指定页池槽位的快照写入 CPU 页。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" pool = self._pool(page_kind) - assert isinstance(slot_ids, list), "slot indices must be a list" + assert isinstance(page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_readies is not None: - assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_readies is not None and cpu_page_readies[index]: + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: continue page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) page_indices.append(page) @@ -223,15 +223,15 @@ def write_page_state_to_cpu( ): copy_bounded_att_tensors(copies=copies) - def read_page_state_from_cpu( + def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" pool = self.big_pages_pool @@ -241,7 +241,7 @@ def read_page_state_from_cpu( # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_pages 按长度恢复。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 for copies in bounded_att_buffer_copies( page_fields=pool_page_fields, page_buffers=cpu_page_buffers, @@ -255,7 +255,7 @@ def read_page_state_from_cpu( # 四、DP 迁移：大页快照与基准运行态 # ======================================================================== - def copy_big_page_state_from_other_dp( + def copy_big_pages_from_dp( self, src_page_ids: List[int], dst_page_ids: List[int], @@ -279,7 +279,7 @@ def copy_big_page_state_from_other_dp( copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def copy_run_time_state_from_other_dp( + def copy_runtime_state_from_dp( self, src_req_indices: List[int], dst_req_indices: List[int], @@ -318,7 +318,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" return self._get_transfer_regions() - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -327,7 +327,7 @@ def write_to_pd_buffer( layout: StateTransferLayout, ) -> None: """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) for copies in bounded_att_buffer_copies( @@ -339,7 +339,7 @@ def write_to_pd_buffer( ): copy_bounded_att_tensors(copies=copies) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, @@ -348,7 +348,7 @@ def read_from_pd_buffer( layout: StateTransferLayout, ) -> None: """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) for copies in bounded_att_buffer_copies( @@ -377,7 +377,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, conv_shape, @@ -404,17 +404,17 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dtype=config.conv_dtype, axes=("layer", "channel", "window"), shape=conv_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="channel", - src_start=local_channel_start, - dst_start=target_channel_start, + local_start=local_channel_start, + external_start=target_channel_start, length=local_channel_count, ), - AxisSlice( + AxisMapping( axis="window", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.conv_width, ), ), @@ -423,7 +423,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions.append(conv_region) logger.debug( "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", - self.resource_id, + self.state_name, name, local_channel_start, target_channel_start, @@ -448,23 +448,23 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dtype=config.ssm_dtype, axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="head", - src_start=0, - dst_start=head_start, + local_start=0, + external_start=head_start, length=config.num_v_heads, ), - AxisSlice( + AxisMapping( axis="key_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.key_dim, ), - AxisSlice( + AxisMapping( axis="value_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.value_dim, ), ), @@ -473,7 +473,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions.append(ssm_region) logger.debug( "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.resource_id, + self.state_name, head_start, config.num_v_heads, write_enabled, @@ -486,7 +486,7 @@ def _pool(self, page_kind: PageKind): """根据大小页类型取得已启用的页池，不访问 allocator。""" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.resource_id}" + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" return pool def _runtime_fields(self, req: int): diff --git a/lightllm/common/att_resources/bounded_att/kda_pages_pool.py b/lightllm/common/attention_state/bounded_att/kda_pages_pool.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/kda_pages_pool.py rename to lightllm/common/attention_state/bounded_att/kda_pages_pool.py diff --git a/lightllm/common/att_resources/bounded_att/kda_state.py b/lightllm/common/attention_state/bounded_att/kda_state.py similarity index 87% rename from lightllm/common/att_resources/bounded_att/kda_state.py rename to lightllm/common/attention_state/bounded_att/kda_state.py index f61d97fadc..6f8cab3326 100644 --- a/lightllm/common/att_resources/bounded_att/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda_state.py @@ -8,7 +8,7 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, StateTransferLayout, TensorField, TransferRegion +from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig from .kda_pages_pool import KDAPagesPool @@ -16,7 +16,7 @@ from .copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: - from lightllm.common.att_resources.types import PageKind + from lightllm.common.attention_state.types import PageKind logger = init_logger(__name__) @@ -42,7 +42,7 @@ def __init__( """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config - self.resource_id: str | None = None # 由 AttResourceManager.set_att_state() 绑定时设置。 + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step self._logical_to_storage_layer = { @@ -125,39 +125,39 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: if self.mtp_state_index is not None: self.mtp_state_index[req_idx].fill_(0) - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, slot_ids): + for req_idx, slot_id in zip(req_indices, page_slot_ids): conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(slot_ids), "request, length and page counts differ" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, slot_ids): + for req_idx, slot_id in zip(req_indices, page_slot_ids): conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) copy_bounded_att_tensors(copies=conv_copies) @@ -183,34 +183,34 @@ def commit_bounded_att_state( def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - region.shape 描述 CPU 页中的全局字段；axis_slices 将本地页池的 TP 分片映射到该字段。 + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 """ return self._get_transfer_regions() - def write_page_state_to_cpu( + def write_state_pages_to_cpu( self, page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """将指定页池槽位的快照写入 CPU 页。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" pool = self._pool(page_kind) - assert isinstance(slot_ids, list), "slot indices must be a list" + assert isinstance(page_slot_ids, list), "slot indices must be a list" assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_readies is not None: - assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(slot_ids, cpu_page_indices)): + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_readies is not None and cpu_page_readies[index]: + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: continue page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) page_indices.append(page) @@ -223,15 +223,15 @@ def write_page_state_to_cpu( ): copy_bounded_att_tensors(copies=copies) - def read_page_state_from_cpu( + def read_state_pages_from_cpu( self, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], big_page_slot_ids: List[int], layout: StateTransferLayout, ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_pages 恢复。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" pool = self.big_pages_pool @@ -241,7 +241,7 @@ def read_page_state_from_cpu( # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_pages 按长度恢复。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 for copies in bounded_att_buffer_copies( page_fields=pool_page_fields, page_buffers=cpu_page_buffers, @@ -255,7 +255,7 @@ def read_page_state_from_cpu( # 四、DP 迁移：大页快照与基准运行态 # ======================================================================== - def copy_big_page_state_from_other_dp( + def copy_big_pages_from_dp( self, src_page_ids: List[int], dst_page_ids: List[int], @@ -279,7 +279,7 @@ def copy_big_page_state_from_other_dp( copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - def copy_run_time_state_from_other_dp( + def copy_runtime_state_from_dp( self, src_req_indices: List[int], dst_req_indices: List[int], @@ -318,7 +318,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" return self._get_transfer_regions() - def write_to_pd_buffer( + def write_to_pd_page( self, req_idx: int, pd_page_buffers: torch.Tensor, @@ -327,7 +327,7 @@ def write_to_pd_buffer( layout: StateTransferLayout, ) -> None: """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) for copies in bounded_att_buffer_copies( @@ -339,7 +339,7 @@ def write_to_pd_buffer( ): copy_bounded_att_tensors(copies=copies) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, req_idx: int, @@ -348,7 +348,7 @@ def read_from_pd_buffer( layout: StateTransferLayout, ) -> None: """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.resource_id == self.resource_id, "transfer layout does not describe this resource" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" fields = self._runtime_fields(req_idx) for copies in bounded_att_buffer_copies( @@ -377,7 +377,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, conv_shape, @@ -404,17 +404,17 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dtype=config.conv_dtype, axes=("layer", "channel", "window"), shape=conv_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="channel", - src_start=local_channel_start, - dst_start=target_channel_start, + local_start=local_channel_start, + external_start=target_channel_start, length=local_channel_count, ), - AxisSlice( + AxisMapping( axis="window", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.conv_width, ), ), @@ -423,7 +423,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions.append(conv_region) logger.debug( "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", - self.resource_id, + self.state_name, name, local_channel_start, target_channel_start, @@ -448,23 +448,23 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: dtype=config.ssm_dtype, axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="head", - src_start=0, - dst_start=head_start, + local_start=0, + external_start=head_start, length=config.num_v_heads, ), - AxisSlice( + AxisMapping( axis="key_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.key_dim, ), - AxisSlice( + AxisMapping( axis="value_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=config.value_dim, ), ), @@ -473,7 +473,7 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: regions.append(ssm_region) logger.debug( "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.resource_id, + self.state_name, head_start, config.num_v_heads, write_enabled, @@ -486,7 +486,7 @@ def _pool(self, page_kind: PageKind): """根据大小页类型取得已启用的页池，不访问 allocator。""" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.resource_id}" + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" return pool def _runtime_fields(self, req: int): diff --git a/lightllm/common/att_resources/bounded_att/transfer_utils.py b/lightllm/common/attention_state/bounded_att/transfer_utils.py similarity index 87% rename from lightllm/common/att_resources/bounded_att/transfer_utils.py rename to lightllm/common/attention_state/bounded_att/transfer_utils.py index 54a0117c73..112358c06a 100644 --- a/lightllm/common/att_resources/bounded_att/transfer_utils.py +++ b/lightllm/common/attention_state/bounded_att/transfer_utils.py @@ -23,7 +23,7 @@ def bounded_att_buffer_copies( page_fields 与 page_indices 一一对应，调用方已排除 ready 页。 CPU 路径的 page_fields 来自本地大小页池，PD 路径来自请求运行态；均为本地 TP 分片。 layout.regions 描述外部页中的全局字段，以及本地 TP 分片在该字段中的位置。 - axis_slices 的 src_start 固定指本地，dst_start 固定指外部页，读回时也不交换含义。 + axis_mappings 的 local_start 固定指本地，external_start 固定指外部页，读回时也不交换含义。 CPU 页和大小页快照的窗口均按逻辑时间排列，此处只处理普通轴切片，不展开或恢复 ring。 Conv 的 Q/K/V 切片与 SSM 分别成批，不将不同字段或不同 TP 切片混入同一次拷贝。 page_buffers 的第一维为物理页；这里只创建视图，由调用方逐批提交 Triton 算子。 @@ -49,10 +49,10 @@ def bounded_att_buffer_copies( # 例如 rank 1 的本地 head[0:4] 对应外部页 head[4:8]，写出和读回使用相同映射。 local_slices = [slice(None)] * local_tensor.ndim external_slices = [slice(None)] * external_tensor.ndim - for axis_slice in region.axis_slices: - axis = local_field.axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + for axis_mapping in region.axis_mappings: + axis = local_field.axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local_view = local_tensor[tuple(local_slices)] external_view = external_tensor[tuple(external_slices)] assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" diff --git a/lightllm/common/att_resources/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py similarity index 100% rename from lightllm/common/att_resources/bounded_att/utils.py rename to lightllm/common/attention_state/bounded_att/utils.py diff --git a/lightllm/common/att_resources/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py similarity index 100% rename from lightllm/common/att_resources/full_att/__init__.py rename to lightllm/common/attention_state/full_att/__init__.py diff --git a/lightllm/common/att_resources/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py similarity index 96% rename from lightllm/common/att_resources/full_att/base_state.py rename to lightllm/common/attention_state/full_att/base_state.py index 372d226de4..a4772e52fa 100644 --- a/lightllm/common/att_resources/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -7,7 +7,7 @@ import torch -from lightllm.common.att_resources.types import ( +from lightllm.common.attention_state.types import ( StateTransferLayout, TransferRegion, ) @@ -89,7 +89,7 @@ def get_state_type(self) -> str: 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 resource_id = f"{model_name}_{get_state_type()}_state"，与成员名完全一致。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 @@ -115,7 +115,7 @@ def __init__( raise NotImplementedError() @abstractmethod - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" raise NotImplementedError() @@ -176,9 +176,9 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: - """按布局写出 token 数据及配套字段；cpu_page_readies 中为 True 的页跳过写入。""" + """按布局写出 token 数据及配套字段；cpu_page_ready_flags 中为 True 的页跳过写入。""" raise NotImplementedError() @abstractmethod @@ -197,7 +197,7 @@ def read_from_cpu( # ======================================================================== @abstractmethod - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -220,7 +220,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... @abstractmethod - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -231,7 +231,7 @@ def write_to_pd_buffer( raise NotImplementedError() @abstractmethod - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, diff --git a/lightllm/common/att_resources/full_att/create_gqa_state.py b/lightllm/common/attention_state/full_att/create_gqa_state.py similarity index 100% rename from lightllm/common/att_resources/full_att/create_gqa_state.py rename to lightllm/common/attention_state/full_att/create_gqa_state.py diff --git a/lightllm/common/att_resources/full_att/create_mla_state.py b/lightllm/common/attention_state/full_att/create_mla_state.py similarity index 100% rename from lightllm/common/att_resources/full_att/create_mla_state.py rename to lightllm/common/attention_state/full_att/create_mla_state.py diff --git a/lightllm/common/att_resources/full_att/dp_transfer.py b/lightllm/common/attention_state/full_att/dp_transfer.py similarity index 100% rename from lightllm/common/att_resources/full_att/dp_transfer.py rename to lightllm/common/attention_state/full_att/dp_transfer.py diff --git a/lightllm/common/att_resources/full_att/gqa_state.py b/lightllm/common/attention_state/full_att/gqa_state.py similarity index 95% rename from lightllm/common/att_resources/full_att/gqa_state.py rename to lightllm/common/attention_state/full_att/gqa_state.py index 6221800ac9..35583f276f 100644 --- a/lightllm/common/att_resources/full_att/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa_state.py @@ -9,13 +9,13 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .gqa_transfer import transfer_gqa_pages from .dp_transfer import copy_gqa_tokens if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout logger = init_logger(name=__name__) @@ -48,7 +48,7 @@ def __init__( """只记录布局配置，KV 与 HOLD 页由后续统一容量初始化分配。""" self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化方案显式设置，区别于 scale 精度。 - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index @@ -65,7 +65,7 @@ def get_state_type(self) -> str: """返回 GQA KV 精度对应的存储方案标记。""" return "gqa_" + dtype_name(dtype=self.config.dtype) - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity shape = self._storage_shape() @@ -158,7 +158,7 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """按布局将 token 数据写到指定 CPU 页，覆盖本 State 负责的字段和 TP 分片。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" @@ -170,7 +170,7 @@ def write_to_cpu( page_indices=cpu_page_indices, layout=layout, write=True, - page_readies=cpu_page_readies, + page_readies=cpu_page_ready_flags, ) def read_from_cpu( @@ -196,7 +196,7 @@ def read_from_cpu( # 四、DP 迁移 # ======================================================================== - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -226,7 +226,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: page_tokens = get_env_start_args().pd_kv_page_size return self._get_transfer_regions(page_tokens=page_tokens) - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -245,7 +245,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, @@ -275,7 +275,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, page_tokens, @@ -305,11 +305,11 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: head_start = config.tp_rank // replica_count * local_head_count # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 write_enabled = config.tp_rank % replica_count == 0 - axis_slices = tuple( - AxisSlice( + axis_mappings = tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=head_start if axis == "head" else 0, + local_start=0, + external_start=head_start if axis == "head" else 0, length=dimension, ) for axis, dimension in zip(axes, local_shape) @@ -326,14 +326,14 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: dtype=field_dtype, axes=axes, shape=tuple(global_shape), - axis_slices=axis_slices, + axis_mappings=axis_mappings, write_enabled=write_enabled, ) regions.append(field_region) byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, + self.state_name, field_name, head_start, local_head_count, @@ -344,7 +344,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", - self.resource_id, + self.state_name, len(regions), byte_offset, ) diff --git a/lightllm/common/att_resources/full_att/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa_transfer.py similarity index 97% rename from lightllm/common/att_resources/full_att/gqa_transfer.py rename to lightllm/common/attention_state/full_att/gqa_transfer.py index 335ecc3d1f..1782fb2672 100644 --- a/lightllm/common/att_resources/full_att/gqa_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa_transfer.py @@ -93,12 +93,12 @@ def transfer_gqa_pages( external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) local_slices = [slice(None)] * 4 external_slices = [slice(None)] * 5 - for axis_slice in region.axis_slices: - # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 - axis = axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 diff --git a/lightllm/common/att_resources/full_att/group_quantized_state.py b/lightllm/common/attention_state/full_att/group_quantized_state.py similarity index 96% rename from lightllm/common/att_resources/full_att/group_quantized_state.py rename to lightllm/common/attention_state/full_att/group_quantized_state.py index 58d432ee11..f72db0940a 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/group_quantized_state.py @@ -10,13 +10,13 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .group_quantized_transfer import transfer_group_quantized_pages from .dp_transfer import copy_group_quantized_tokens if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout logger = init_logger(name=__name__) @@ -57,7 +57,7 @@ def __init__( ) -> None: """按分组量化配置准备 KV 与动态 scale 存储。""" self.config = layout_config - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index @@ -77,7 +77,7 @@ def get_state_type(self) -> str: """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" return f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.dtype)}_scale" - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" self.capacity = token_capacity shape = self._storage_shape() @@ -207,7 +207,7 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" @@ -221,7 +221,7 @@ def write_to_cpu( page_indices=cpu_page_indices, layout=layout, write=True, - page_readies=cpu_page_readies, + page_readies=cpu_page_ready_flags, ) def read_from_cpu( @@ -249,7 +249,7 @@ def read_from_cpu( # 四、DP 迁移 # ======================================================================== - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -289,7 +289,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: page_tokens = get_env_start_args().pd_kv_page_size return self._get_transfer_regions(page_tokens=page_tokens) - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -310,7 +310,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, @@ -342,7 +342,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, page_tokens, @@ -374,11 +374,11 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: head_start = config.tp_rank // replica_count * local_head_count # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 write_enabled = config.tp_rank % replica_count == 0 - axis_slices = tuple( - AxisSlice( + axis_mappings = tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=head_start if axis == "head" else 0, + local_start=0, + external_start=head_start if axis == "head" else 0, length=dimension, ) for axis, dimension in zip(axes, local_shape) @@ -395,14 +395,14 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: dtype=field_dtype, axes=axes, shape=tuple(global_shape), - axis_slices=axis_slices, + axis_mappings=axis_mappings, write_enabled=write_enabled, ) regions.append(field_region) byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, + self.state_name, field_name, head_start, local_head_count, @@ -413,7 +413,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", - self.resource_id, + self.state_name, len(regions), byte_offset, ) diff --git a/lightllm/common/att_resources/full_att/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/group_quantized_transfer.py similarity index 97% rename from lightllm/common/att_resources/full_att/group_quantized_transfer.py rename to lightllm/common/attention_state/full_att/group_quantized_transfer.py index 402c254e86..371913e758 100644 --- a/lightllm/common/att_resources/full_att/group_quantized_transfer.py +++ b/lightllm/common/attention_state/full_att/group_quantized_transfer.py @@ -103,12 +103,12 @@ def transfer_group_quantized_pages( external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) local_slices = [slice(None)] * 4 external_slices = [slice(None)] * 5 - for axis_slice in region.axis_slices: - # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 - axis = axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 diff --git a/lightllm/common/att_resources/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla_state.py similarity index 94% rename from lightllm/common/att_resources/full_att/mla_state.py rename to lightllm/common/attention_state/full_att/mla_state.py index 5395c9ca06..aabd08eadc 100644 --- a/lightllm/common/att_resources/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla_state.py @@ -9,13 +9,13 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .mla_transfer import transfer_mla_pages from .dp_transfer import copy_mla_tokens if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout logger = init_logger(name=__name__) @@ -38,7 +38,7 @@ def __init__( """按 MLA 布局配置准备压缩 KV 存储。""" self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index @@ -54,7 +54,7 @@ def get_state_type(self) -> str: """返回 MLA 压缩 KV 精度对应的存储方案标记。""" return "mla_" + dtype_name(dtype=self.config.dtype) - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity shape = self._storage_shape() @@ -136,7 +136,7 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" @@ -147,7 +147,7 @@ def write_to_cpu( page_indices=cpu_page_indices, layout=layout, write=True, - page_readies=cpu_page_readies, + page_readies=cpu_page_ready_flags, ) def read_from_cpu( @@ -172,7 +172,7 @@ def read_from_cpu( # 四、DP 迁移 # ======================================================================== - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -200,7 +200,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: page_tokens = get_env_start_args().pd_kv_page_size return self._get_transfer_regions(page_tokens=page_tokens) - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -218,7 +218,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, @@ -247,7 +247,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, page_tokens, @@ -274,11 +274,11 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: head_start = config.tp_rank // replica_count * local_head_count # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 write_enabled = config.tp_rank % replica_count == 0 - axis_slices = tuple( - AxisSlice( + axis_mappings = tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=head_start if axis == "head" else 0, + local_start=0, + external_start=head_start if axis == "head" else 0, length=dimension, ) for axis, dimension in zip(axes, local_shape) @@ -295,14 +295,14 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: dtype=field_dtype, axes=axes, shape=tuple(global_shape), - axis_slices=axis_slices, + axis_mappings=axis_mappings, write_enabled=write_enabled, ) regions.append(field_region) byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, + self.state_name, field_name, head_start, local_head_count, @@ -313,7 +313,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", - self.resource_id, + self.state_name, len(regions), byte_offset, ) diff --git a/lightllm/common/att_resources/full_att/mla_transfer.py b/lightllm/common/attention_state/full_att/mla_transfer.py similarity index 97% rename from lightllm/common/att_resources/full_att/mla_transfer.py rename to lightllm/common/attention_state/full_att/mla_transfer.py index 6b666ec0e3..2855077ff3 100644 --- a/lightllm/common/att_resources/full_att/mla_transfer.py +++ b/lightllm/common/attention_state/full_att/mla_transfer.py @@ -90,12 +90,12 @@ def transfer_mla_pages( external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) local_slices = [slice(None)] * 4 external_slices = [slice(None)] * 5 - for axis_slice in region.axis_slices: - # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 - axis = axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 diff --git a/lightllm/common/att_resources/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla_state.py similarity index 96% rename from lightllm/common/att_resources/full_att/packed_mla_state.py rename to lightllm/common/attention_state/full_att/packed_mla_state.py index 20bfc19e20..95d9c7344f 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla_state.py @@ -10,13 +10,13 @@ from lightllm.utils.log_utils import init_logger from ..utils import dtype_name -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .packed_mla_transfer import transfer_packed_mla_pages from .dp_transfer import copy_mla_tokens if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout logger = init_logger(name=__name__) @@ -79,7 +79,7 @@ def __init__( """按 packed MLA 配置准备 KV 与 indexer 存储。""" self.config = layout_config self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index @@ -96,7 +96,7 @@ def get_state_type(self) -> str: encoding = f"fp8_group{self.config.quant_group_size}" if self.config.fp8_nope_dim else dtype_name(dtype=self.config.dtype) return "mla_packed_" + encoding - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity shape = self._storage_shape() @@ -209,7 +209,7 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" @@ -223,7 +223,7 @@ def write_to_cpu( indexer_offset=self.indexer_offset, indexer_dim=self.config.indexer_dim, write=True, - page_readies=cpu_page_readies, + page_readies=cpu_page_ready_flags, ) def read_from_cpu( @@ -251,7 +251,7 @@ def read_from_cpu( # 四、DP 迁移 # ======================================================================== - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -288,7 +288,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: page_tokens = get_env_start_args().pd_kv_page_size return self._get_transfer_regions(page_tokens=page_tokens) - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -309,7 +309,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, @@ -341,7 +341,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, page_tokens, @@ -379,11 +379,11 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: head_start = config.tp_rank // replica_count * local_head_count # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 write_enabled = config.tp_rank % replica_count == 0 - axis_slices = tuple( - AxisSlice( + axis_mappings = tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=head_start if axis == "head" else 0, + local_start=0, + external_start=head_start if axis == "head" else 0, length=dimension, ) for axis, dimension in zip(axes, local_shape) @@ -400,14 +400,14 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: dtype=field_dtype, axes=axes, shape=tuple(global_shape), - axis_slices=axis_slices, + axis_mappings=axis_mappings, write_enabled=write_enabled, ) regions.append(field_region) byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, + self.state_name, field_name, head_start, local_head_count, @@ -418,7 +418,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", - self.resource_id, + self.state_name, len(regions), byte_offset, ) diff --git a/lightllm/common/att_resources/full_att/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla_transfer.py similarity index 97% rename from lightllm/common/att_resources/full_att/packed_mla_transfer.py rename to lightllm/common/attention_state/full_att/packed_mla_transfer.py index b26854c91e..c3d43d6b93 100644 --- a/lightllm/common/att_resources/full_att/packed_mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla_transfer.py @@ -112,12 +112,12 @@ def transfer_packed_mla_pages( external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) local_slices = [slice(None)] * 4 external_slices = [slice(None)] * 5 - for axis_slice in region.axis_slices: - # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 - axis = axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 diff --git a/lightllm/common/att_resources/full_att/static_fp8_state.py b/lightllm/common/attention_state/full_att/static_fp8_state.py similarity index 96% rename from lightllm/common/att_resources/full_att/static_fp8_state.py rename to lightllm/common/attention_state/full_att/static_fp8_state.py index 61398c4e7e..14350cde90 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/static_fp8_state.py @@ -10,13 +10,13 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..types import AxisSlice, TransferRegion +from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .static_fp8_transfer import transfer_static_fp8_pages from .dp_transfer import copy_gqa_tokens if TYPE_CHECKING: - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout logger = init_logger(name=__name__) @@ -53,7 +53,7 @@ def __init__( self.q_scales: torch.Tensor | None = None self.cpu_scales = None # 实际存储初始化时准备计算 backend 使用的 CPU 标定值。 self.config = layout_config - self.resource_id: str | None = None # 由 manager 绑定模型角色时设置。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index @@ -131,7 +131,7 @@ def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: scales = scale_tensor(calibration, (calibrated_layer_count, 2)) q_scales = None - # 阶段三：绑定标定数值；后续 initialize_storage 将其迁移到实际设备。 + # 阶段三：绑定标定数值；后续 allocate_storage 将其迁移到实际设备。 self.scales = scales self.q_scales = q_scales @@ -139,7 +139,7 @@ def get_state_type(self) -> str: """返回包含 per-tensor 或 per-head 标定方式的 FP8 方案标记。""" return "gqa_fp8_static_" + self.config.scale_mode - def initialize_storage(self, token_capacity: int) -> None: + def allocate_storage(self, token_capacity: int) -> None: """绑定 KV 与真实标定参数，并预备 FlashInfer 使用的 CPU scale。""" assert self.scales is not None, "static FP8 calibration must be bound before allocation" self.capacity = token_capacity @@ -245,7 +245,7 @@ def write_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layout: StateTransferLayout, - cpu_page_readies: list[bool] | None = None, + cpu_page_ready_flags: list[bool] | None = None, ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" @@ -257,7 +257,7 @@ def write_to_cpu( page_indices=cpu_page_indices, layout=layout, write=True, - page_readies=cpu_page_readies, + page_readies=cpu_page_ready_flags, ) def read_from_cpu( @@ -283,7 +283,7 @@ def read_from_cpu( # 四、DP 迁移 # ======================================================================== - def copy_tokens_from_other_dp( + def copy_tokens_from_dp( self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, @@ -313,7 +313,7 @@ def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: page_tokens = get_env_start_args().pd_kv_page_size return self._get_transfer_regions(page_tokens=page_tokens) - def write_to_pd_buffer( + def write_to_pd_page( self, token_indices: torch.Tensor, pd_page_buffers: torch.Tensor, @@ -332,7 +332,7 @@ def write_to_pd_buffer( write=True, ) - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, token_indices: torch.Tensor, @@ -362,7 +362,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: regions = [] logger.debug( "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.resource_id, + self.state_name, config.tp_rank, config.tp_world_size, page_tokens, @@ -392,11 +392,11 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: head_start = config.tp_rank // replica_count * local_head_count # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 write_enabled = config.tp_rank % replica_count == 0 - axis_slices = tuple( - AxisSlice( + axis_mappings = tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=head_start if axis == "head" else 0, + local_start=0, + external_start=head_start if axis == "head" else 0, length=dimension, ) for axis, dimension in zip(axes, local_shape) @@ -413,14 +413,14 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: dtype=field_dtype, axes=axes, shape=tuple(global_shape), - axis_slices=axis_slices, + axis_mappings=axis_mappings, write_enabled=write_enabled, ) regions.append(field_region) byte_offset = field_region.byte_end logger.debug( "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.resource_id, + self.state_name, field_name, head_start, local_head_count, @@ -431,7 +431,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", - self.resource_id, + self.state_name, len(regions), byte_offset, ) diff --git a/lightllm/common/att_resources/full_att/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/static_fp8_transfer.py similarity index 97% rename from lightllm/common/att_resources/full_att/static_fp8_transfer.py rename to lightllm/common/attention_state/full_att/static_fp8_transfer.py index bcbf55f871..13ec1cba1c 100644 --- a/lightllm/common/att_resources/full_att/static_fp8_transfer.py +++ b/lightllm/common/attention_state/full_att/static_fp8_transfer.py @@ -95,12 +95,12 @@ def transfer_static_fp8_pages( external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) local_slices = [slice(None)] * 4 external_slices = [slice(None)] * 5 - for axis_slice in region.axis_slices: - # src_start/dst_start 固定表示本地/外部起点，不随读写方向改变。 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 - axis = axes.index(axis_slice.axis) - local_slices[axis] = slice(axis_slice.src_start, axis_slice.src_start + axis_slice.length) - external_slices[axis + 1] = slice(axis_slice.dst_start, axis_slice.dst_start + axis_slice.length) + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 diff --git a/lightllm/common/att_resources/att_resource_manager.py b/lightllm/common/attention_state/manager.py similarity index 84% rename from lightllm/common/att_resources/att_resource_manager.py rename to lightllm/common/attention_state/manager.py index e449fc2a65..b9af2e2c73 100644 --- a/lightllm/common/att_resources/att_resource_manager.py +++ b/lightllm/common/attention_state/manager.py @@ -9,7 +9,7 @@ import torch -from lightllm.common.att_resources.types import ( +from lightllm.common.attention_state.types import ( StateTransferLayout, PageKind, ) @@ -32,7 +32,7 @@ logger = init_logger(__name__) -class AttResourceManager: +class AttentionStateManager: """管理模型组合的注意力资源，提供缓存复用及 CPU/DP/PD 传输所需的聚合操作。 通过聚合查询操作各类 State；模型通过三类创建接口决定资源组合，main/draft 共享本管理器。 @@ -47,7 +47,7 @@ class AttResourceManager: full_att_allocator: FullAttKVAllocator | None bounded_att_big_page_allocator: StatePageAllocator | None bounded_att_small_page_allocator: StatePageAllocator | None - _states_created: bool = False + _state_registration_finished: bool = False def __init__(self, request_capacity: int): """记录共享请求容量；State 和 allocator 由后续初始化流程设置。""" @@ -57,57 +57,57 @@ def __init__(self, request_capacity: int): self._request_capacity = request_capacity self.cpu_tail_restore = False - def set_att_state( - self, model_name: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState + def register_state( + self, model_role: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState ) -> None: """按模型角色和具体方案注册 State，保留原对象及其存储。 - 成员名和 resource_id 均为 {model_name}_{state.get_state_type()}_state。 - model_name 为 main 或 draft0..draft35；State 只返回中间的算法/存储方案标记， + 成员名和 state_name 均为 {model_role}_{state.get_state_type()}_state。 + model_role 为 main 或 draft0..draft35；State 只返回中间的算法/存储方案标记， 不自行添加角色或 _state。各 State 基类的 get_state_type 注释规定具体命名格式。 同角色、同标记不能重复注册；成员名是资源身份，不代表完整布局兼容性。 """ - assert not self._states_created, "State creation is already finished" - assert model_name in MODEL_ROLES, f"unknown model role {model_name}" + assert not self._state_registration_finished, "State registration is already finished" + assert model_role in MODEL_ROLES, f"unknown model role {model_role}" state_type = state.get_state_type() assert state_type.isidentifier() and state_type.islower(), f"invalid State type {state_type}" - name = f"{model_name}_{state_type}_state" + name = f"{model_role}_{state_type}_state" assert name not in self.__dict__, f"State {name} is already bound" assert all(state is not value for value in self.__dict__.values()), "a State cannot belong to multiple model roles" # config 只描述布局，资源身份仅保存在 State；绑定时保留原配置对象。 - state.resource_id = name + state.state_name = name setattr(self, name, state) @lru_cache(maxsize=None) - def get_att_state( + def get_state( self, - model_name: str, - att_type: str, + model_role: str, + state_type_filter: str, logical_layer_id: int, state_class: type[FullAttentionState | BoundedAttentionState | AttentionTailState] | None = None, ) -> FullAttentionState | BoundedAttentionState | AttentionTailState: """按模型角色、算法/方案字符串和逻辑层选择唯一 State，缓存查询结果。 - model_name 为 main 或 draft0..draft35，按完整角色前缀匹配，避免 draft1 匹配 draft10。 - att_type 在 get_state_type() 返回的字符串中做包含匹配，如 gqa、gqa_int8、mla_packed。 + model_role 为 main 或 draft0..draft35，按完整角色前缀匹配，避免 draft1 匹配 draft10。 + state_type_filter 在 get_state_type() 返回的字符串中做包含匹配，如 gqa、gqa_int8、mla_packed。 通用推理模板可传空字符串，并通过 state_class 指定资源类别；具体算法传自己的标记。 同层存在多类资源时，state_class 可进一步区分；最终必须恰好匹配一个 State。 全部模型 State 创建完成后才允许查询；后续原地分配存储不改变缓存对象身份。 """ - assert self._states_created, "State creation must finish before querying" - assert model_name in MODEL_ROLES, f"unknown model role {model_name}" + assert self._state_registration_finished, "State registration must finish before querying" + assert model_role in MODEL_ROLES, f"unknown model role {model_role}" matches = [ state for name, state in self.__dict__.items() - if name.startswith(model_name + "_") + if name.startswith(model_role + "_") and isinstance(state, (FullAttentionState, BoundedAttentionState, AttentionTailState)) - and att_type in state.get_state_type() + and state_type_filter in state.get_state_type() and (state_class is None or isinstance(state, state_class)) and state.contains_logical_layer(logical_layer_id) ] assert len(matches) == 1, ( - f"expected exactly one attention State: {model_name=}, {att_type=}, {logical_layer_id=}, " - f"matches={[state.resource_id for state in matches]}" + f"expected exactly one attention State: {model_role=}, {state_type_filter=}, {logical_layer_id=}, " + f"matches={[state.state_name for state in matches]}" ) return matches[0] @@ -145,7 +145,7 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: states = sorted( self.get_full_attention_states(role), key=lambda state: ( - min(layer for group in state.config.logical_layer_groups for layer in group), state.resource_id, + min(layer for group in state.config.logical_layer_groups for layer in group), state.state_name, ), ) for state in states: @@ -157,7 +157,7 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: storage_layer_start += len(state.config.logical_layer_groups) logger.info("Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales") - def initialize_resources( + def initialize_storage( self, args: StartArgs | EasyDict, token_capacity: int | None, @@ -170,7 +170,7 @@ def initialize_resources( 已完成创建标记和所需的静态标定绑定，大小页池尚未分配。调用方在本方法返回后， 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 """ - assert self._states_created, "State creation must finish before resource initialization" + assert self._state_registration_finished, "State registration must finish before resource initialization" # 阶段一：规划共享 token 容量，并向下对齐到完整 KV 页。 # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 @@ -209,7 +209,7 @@ def initialize_resources( # 阶段二：管理器为已创建的 State 分配实际存储与共享 allocator。 # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - self.initialize_storage( + self.allocate_storage( token_capacity=token_capacity, page_size=args.page_size, big_page_tokens=big_page_tokens, @@ -218,12 +218,12 @@ def initialize_resources( ) # 创建完成只表示对象集合就绪；实际容量和存储在后续阶段初始化。 - def finish_state_creation(self) -> None: + def finish_state_registration(self) -> None: """确认全部模型的 State 对象已创建并装配完成，允许查询 meta 定义并开始容量规划。""" - assert not self._states_created, "State creation is already finished" - self._states_created = True + assert not self._state_registration_finished, "State registration is already finished" + self._state_registration_finished = True - def initialize_storage( + def allocate_storage( self, token_capacity: int, page_size: int, @@ -235,7 +235,7 @@ def initialize_storage( 请求容量对应的 bounded/indexer 运行态已由各模型创建，这里仅补齐缓存存储。 """ - assert self._states_created, "State creation must finish before storage initialization" + assert self._state_registration_finished, "State registration must finish before storage initialization" assert self.full_att_allocator is None, "resources already initialized" full_states = self.get_full_attention_states() assert full_states, "the shared token index domain requires a full-attention State" @@ -243,7 +243,7 @@ def initialize_storage( # 管理器统一设置实际页规格和设备；State 的分配接口只接收 token 容量。 state.page_size = page_size state.device = torch.device(device) - state.initialize_storage(token_capacity) + state.allocate_storage(token_capacity) self.full_att_allocator = FullAttKVAllocator(token_capacity, page_size) # bounded-attention 大页池末尾保留 CPU 加载/卸载暂存槽位，普通分配不使用它们。 @@ -265,13 +265,13 @@ def initialize_storage( self.cpu_tail_restore = bool(self.get_attention_tail_states()) # 存储分配可能改变 State 的物理页规格，清除准备阶段的布局及其派生 CPU/PD 页大小缓存。 - self.get_cpu_transfer_layouts.cache_clear() + self.get_cpu_page_layouts.cache_clear() self.get_cpu_page_bytes.cache_clear() - self.get_pd_full_att_transfer_layouts.cache_clear() - self.get_pd_bounded_att_or_tail_transfer_layouts.cache_clear() + self.get_pd_token_page_layouts.cache_clear() + self.get_pd_runtime_page_layouts.cache_clear() self.get_pd_page_bytes.cache_clear() - def create_pd_ipc_copy(self) -> AttResourceManager: + def create_pd_ipc_copy(self) -> AttentionStateManager: """创建 PD 导出副本，保留运行态 Tensor，排除 allocator 和全部大小页池。""" ipc_copy = copy(self) ipc_copy.full_att_allocator = None @@ -283,10 +283,10 @@ def create_pd_ipc_copy(self) -> AttResourceManager: state_copy = copy(state) state_copy.big_pages_pool = None state_copy.small_pages_pool = None - setattr(ipc_copy, state.resource_id, state_copy) + setattr(ipc_copy, state.state_name, state_copy) return ipc_copy - def create_dp_ipc_copy(self) -> AttResourceManager: + def create_dp_ipc_copy(self) -> AttentionStateManager: """创建 DP shared 导出副本，保留运行态和共享大小页池，仅排除 allocator。""" # DP fetch 初始化已选择共享 CPU 页池；禁止导出私有页池时隐式分配、复制整块存储。 for state in self.get_bounded_attention_states(): @@ -337,7 +337,7 @@ def get_full_att_token_bytes(self) -> int: def get_cpu_page_bytes(self) -> int: """从 CPU 联合页布局的最大终点计算物理页大小，包含字段及 State 间的对齐空间。""" page_bytes = 0 - for layout in self.get_cpu_transfer_layouts().values(): + for layout in self.get_cpu_page_layouts().values(): page_bytes = max(page_bytes, layout.byte_end) # State 区间按 8 字节对齐，整个 CPU 物理页再按 16 字节对齐。 return (page_bytes + 15) // 16 * 16 @@ -345,8 +345,8 @@ def get_cpu_page_bytes(self) -> int: @lru_cache(maxsize=None) def get_pd_page_bytes(self) -> int: """取 KV 页与完整运行态页占用的最大值，保证两类 PD 任务都能容纳。""" - full_layouts = self.get_pd_full_att_transfer_layouts() - runtime_layouts = self.get_pd_bounded_att_or_tail_transfer_layouts() + full_layouts = self.get_pd_token_page_layouts() + runtime_layouts = self.get_pd_runtime_page_layouts() full_att_page_bytes = 0 for layout in full_layouts.values(): full_att_page_bytes = max(full_att_page_bytes, layout.byte_end) @@ -387,49 +387,49 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: # Bounded attention 状态页：供 radix/CPU cache 复用；未启用对应页池时操作报错。 # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 - def save_bounded_att_pages( + def save_bounded_att_runtime_to_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" bounded_states = self.get_bounded_attention_states() - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" + assert len(page_slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" for state in bounded_states: - state.save_bounded_att_pages( + state.save_bounded_att_runtime_to_pages( req_indices=req_indices, state_token_lens=state_token_lens, page_kind=page_kind, - slot_ids=slot_ids, + page_slot_ids=page_slot_ids, ) - def restore_bounded_att_pages( + def restore_bounded_att_runtime_from_pages( self, req_indices: List[int], state_token_lens: List[int], page_kind: PageKind, - slot_ids: List[int], + page_slot_ids: List[int], ) -> None: """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" bounded_states = self.get_bounded_attention_states() - assert isinstance(req_indices, list) and isinstance(slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" - assert not slot_ids or bounded_states, "state restore requires participating bounded states" + assert len(page_slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" + assert not page_slot_ids or bounded_states, "state restore requires participating bounded states" for state in bounded_states: - state.restore_bounded_att_pages( + state.restore_bounded_att_runtime_from_pages( req_indices=req_indices, state_token_lens=state_token_lens, page_kind=page_kind, - slot_ids=slot_ids, + page_slot_ids=page_slot_ids, ) # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 @@ -466,8 +466,8 @@ def reset_attention_tail_state( # ======================================================================== @lru_cache(maxsize=None) - def get_cpu_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: - """聚合各 State 的 CPU 字段区间，返回按 resource_id 索引的联合页布局。 + def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: + """聚合各 State 的 CPU 字段区间，返回按 state_name 索引的联合页布局。 使用全局 cpu_cache_token_page_size，按需缓存；不依赖页大小查询。 """ @@ -486,8 +486,8 @@ def get_cpu_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ) # 以实际字段终点确定 State 占用，保留内部 padding，并对齐下一个 State 的起点。 state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.resource_id] = StateTransferLayout( - resource_id=state.resource_id, + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, tokens_per_page=page_tokens, byte_start=offset, byte_end=state_end, @@ -504,7 +504,7 @@ def write_req_to_cpu( cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int], layouts: Mapping[str, StateTransferLayout], - cpu_page_readies: List[bool] | None = None, + cpu_page_ready_flags: List[bool] | None = None, big_page_slot_ids: List[int] | None = None, small_page_slot_ids: List[int] | None = None, ) -> None: @@ -512,11 +512,11 @@ def write_req_to_cpu( cpu_page_indices 前段对应 big_page_slot_ids，后段对应 small_page_slot_ids， 各段按列表顺序一一对应；未参与的页池可传 None 或空列表。 - cpu_page_readies 与 CPU 页一一对应，True 跳过整页；None 表示全部写出。 + cpu_page_ready_flags 与 CPU 页一一对应，True 跳过整页；None 表示全部写出。 """ - if cpu_page_readies is not None: - assert len(cpu_page_readies) == len(cpu_page_indices), "CPU page and ready counts differ" - if all(cpu_page_readies): + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + if all(cpu_page_ready_flags): return full_states = self.get_full_attention_states() bounded_states = self.get_bounded_attention_states() @@ -527,8 +527,8 @@ def write_req_to_cpu( assert big_page_count + len(small_page_slot_ids) == len(cpu_page_indices), "one state slot is required per CPU page" big_cpu_page_indices = cpu_page_indices[:big_page_count] small_cpu_page_indices = cpu_page_indices[big_page_count:] - big_page_readies = cpu_page_readies[:big_page_count] if cpu_page_readies is not None else None - small_page_readies = cpu_page_readies[big_page_count:] if cpu_page_readies is not None else None + big_page_readies = cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None + small_page_readies = cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None # 阶段一：统一写出本批 full-attention token 数据。 for state in full_states: @@ -536,30 +536,30 @@ def write_req_to_cpu( token_indices=token_indices, cpu_page_buffers=cpu_page_buffers, cpu_page_indices=cpu_page_indices, - layout=layouts[state.resource_id], - cpu_page_readies=cpu_page_readies, + layout=layouts[state.state_name], + cpu_page_ready_flags=cpu_page_ready_flags, ) # 阶段二：分别从大小页池取快照，写入对应 CPU 页面，不修改请求运行态。 for state in bounded_states: - layout = layouts[state.resource_id] + layout = layouts[state.state_name] if big_page_slot_ids: - state.write_page_state_to_cpu( + state.write_state_pages_to_cpu( page_kind="big", - slot_ids=big_page_slot_ids, + page_slot_ids=big_page_slot_ids, cpu_page_buffers=cpu_page_buffers, cpu_page_indices=big_cpu_page_indices, layout=layout, - cpu_page_readies=big_page_readies, + cpu_page_ready_flags=big_page_readies, ) if small_page_slot_ids: - state.write_page_state_to_cpu( + state.write_state_pages_to_cpu( page_kind="small", - slot_ids=small_page_slot_ids, + page_slot_ids=small_page_slot_ids, cpu_page_buffers=cpu_page_buffers, cpu_page_indices=small_cpu_page_indices, layout=layout, - cpu_page_readies=small_page_readies, + cpu_page_ready_flags=small_page_readies, ) def read_req_from_cpu( @@ -586,7 +586,7 @@ def read_req_from_cpu( if bounded_states: assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" # token 索引只用于判断本批末页是否完整，不读取 GPU 上的索引值。 - page_tokens = layouts[bounded_states[0].resource_id].tokens_per_page + page_tokens = layouts[bounded_states[0].state_name].tokens_per_page has_partial_page = token_indices.numel() % page_tokens != 0 assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), "slots must cover complete CPU pages only" @@ -601,25 +601,25 @@ def read_req_from_cpu( cpu_page_buffers=cpu_page_buffers, cpu_page_indices=cpu_page_indices, token_indices=token_indices, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) # 阶段二：完整页和末尾碎页统一读入大页池，碎页使用追加的 LOAD 槽位。 for state in bounded_states: - state.read_page_state_from_cpu( + state.read_state_pages_from_cpu( cpu_page_buffers=cpu_page_buffers, cpu_page_indices=cpu_page_indices, big_page_slot_ids=load_big_page_slot_ids, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) # 阶段三：从末页恢复请求运行态；纯碎页也走 LOAD 槽位。 if bounded_states: - self.restore_bounded_att_pages( + self.restore_bounded_att_runtime_from_pages( req_indices=[req_idx], state_token_lens=[state_token_len], page_kind="big", - slot_ids=[load_big_page_slot_ids[-1]], + page_slot_ids=[load_big_page_slot_ids[-1]], ) # 阶段四：清零目标请求的 tail；调用方保证前缀终点满足重置边界。 @@ -630,11 +630,11 @@ def read_req_from_cpu( # 六、DP 迁移 # ======================================================================== - def copy_from_other_dp( + def copy_from_dp( self, - remote_att_manager: AttResourceManager, - src_token_indexes: torch.Tensor, - dst_token_indexes: torch.Tensor, + remote_att_manager: AttentionStateManager, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, src_big_page_ids: List[int], dst_big_page_ids: List[int], src_req_idx: int, @@ -647,28 +647,28 @@ def copy_from_other_dp( 调用方保证源 MTP 游标为 0、目标前缀位于 tail 可重置边界，以及源页就绪和引用有效。 """ assert type(remote_att_manager) is type(self), "DP replicas must use the same model resource manager" - assert src_token_indexes.numel() == dst_token_indexes.numel(), "source and destination token counts differ" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" assert len(src_big_page_ids) == len(dst_big_page_ids), "source and destination page counts differ" - if src_token_indexes.numel(): + if src_token_indices.numel(): for state in self.get_full_attention_states(): - state.copy_tokens_from_other_dp( - src_token_indices=src_token_indexes, - dst_token_indices=dst_token_indexes, - remote_dp_full_att_state=getattr(remote_att_manager, state.resource_id), + state.copy_tokens_from_dp( + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + remote_dp_full_att_state=getattr(remote_att_manager, state.state_name), ) if src_big_page_ids: for state in self.get_bounded_attention_states(): - state.copy_big_page_state_from_other_dp( + state.copy_big_pages_from_dp( src_page_ids=src_big_page_ids, dst_page_ids=dst_big_page_ids, - remote_dp_bounded_state=getattr(remote_att_manager, state.resource_id), + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), ) # 运行态不由大页复制隐式恢复；按指定请求复制基准位置，随后重置目标 tail。 for state in self.get_bounded_attention_states(): - state.copy_run_time_state_from_other_dp( + state.copy_runtime_state_from_dp( src_req_indices=[src_req_idx], dst_req_indices=[dst_req_idx], - remote_dp_bounded_state=getattr(remote_att_manager, state.resource_id), + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), ) self.reset_attention_tail_state([dst_req_idx]) @@ -677,7 +677,7 @@ def copy_from_other_dp( # ======================================================================== @lru_cache(maxsize=None) - def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: + def get_pd_token_page_layouts(self) -> Mapping[str, StateTransferLayout]: """按实际字段区间生成 kv 布局，各 State 记录自身在页内的起点与终点。 使用全局 pd_kv_page_size，按需缓存；不依赖页大小查询或运行态布局。 @@ -692,8 +692,8 @@ def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: ) # 字段保留自身数据区间，State 终点包含留给下一个资源的对齐空间。 state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.resource_id] = StateTransferLayout( - resource_id=state.resource_id, + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, tokens_per_page=page_tokens, byte_start=offset, byte_end=state_end, @@ -703,7 +703,7 @@ def get_pd_full_att_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: return layouts @lru_cache(maxsize=None) - def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTransferLayout]: + def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: """独立生成 att_state 布局，记录各运行态 State 的完整字节区间。 每个 State 保存一份运行态，tokens_per_page 为 None；无 State 时返回空字典。 @@ -715,8 +715,8 @@ def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTrans for region in state.get_pd_transfer_regions() ) state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.resource_id] = StateTransferLayout( - resource_id=state.resource_id, + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, tokens_per_page=None, byte_start=offset, byte_end=state_end, @@ -727,7 +727,7 @@ def get_pd_bounded_att_or_tail_transfer_layouts(self) -> Mapping[str, StateTrans # PD 分离：kind="kv" 传输 token 数据，kind="att_state" 传输 bounded 运行态及 tail。 # pd_page_buffers 第一维为页，pd_page_index 选择物理页；网络收发与同步由传输流程组织。 - def write_to_pd_buffer( + def write_to_pd_page( self, pd_page_buffers: torch.Tensor, pd_page_index: int, @@ -743,11 +743,11 @@ def write_to_pd_buffer( # KV 页只写出 full-attention 数据；各 State 负责自身布局与容量约束。 assert token_indices is not None, "KV PD pages require token indices" for state in self.get_full_attention_states(): - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=token_indices, pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) case "att_state": @@ -758,28 +758,28 @@ def write_to_pd_buffer( # 阶段一：写出 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): - state.write_to_pd_buffer( + state.write_to_pd_page( req_idx=req_idx, pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, state_token_len=state_token_len, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) # 阶段二：按前缀长度写出 tail 逻辑窗口，适配两端各自的 MTP ring。 for state in tail_states: - state.write_to_pd_buffer( + state.write_to_pd_page( req_idx=req_idx, pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, state_token_len=state_token_len, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) case _: raise AssertionError(f"unknown PD page kind {kind}") - def read_from_pd_buffer( + def read_from_pd_page( self, pd_page_buffers: torch.Tensor, pd_page_index: int, @@ -795,11 +795,11 @@ def read_from_pd_buffer( # KV 页只恢复 full-attention 数据；各 State 负责自身布局与容量约束。 assert token_indices is not None, "KV PD pages require token indices" for state in self.get_full_attention_states(): - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pd_page_buffers, token_indices=token_indices, pd_page_index=pd_page_index, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) case "att_state": @@ -810,22 +810,22 @@ def read_from_pd_buffer( # 阶段一：恢复 bounded 的 prefill 基准运行态。 for state in self.get_bounded_attention_states(): - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pd_page_buffers, req_idx=req_idx, pd_page_index=pd_page_index, state_token_len=state_token_len, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) # 阶段二：按前缀长度恢复 tail 逻辑窗口，适配两端各自的 MTP ring。 for state in tail_states: - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pd_page_buffers, req_idx=req_idx, pd_page_index=pd_page_index, state_token_len=state_token_len, - layout=layouts[state.resource_id], + layout=layouts[state.state_name], ) case _: @@ -843,11 +843,11 @@ def _get_states( role: str | None, ): """按 State 类型和模型角色遍历规范成员，只缓存对象引用。""" - assert self._states_created, "State creation must finish before querying" + assert self._state_registration_finished, "State registration must finish before querying" assert role is None or role in MODEL_ROLES, f"unknown model role {role}" states = ( state for state in self.__dict__.values() - if isinstance(state, state_class) and (role is None or state.resource_id.startswith(role + "_")) + if isinstance(state, state_class) and (role is None or state.state_name.startswith(role + "_")) ) - return tuple(sorted(states, key=lambda state: state.resource_id)) + return tuple(sorted(states, key=lambda state: state.state_name)) diff --git a/lightllm/common/att_resources/types.py b/lightllm/common/attention_state/types.py similarity index 82% rename from lightllm/common/att_resources/types.py rename to lightllm/common/attention_state/types.py index b2eaf0d226..e2fe4c8bd6 100644 --- a/lightllm/common/att_resources/types.py +++ b/lightllm/common/attention_state/types.py @@ -1,6 +1,6 @@ """注意力存储字段与传输位置的只读描述，不执行数据搬运。 -TensorField 描述本地 Tensor；AxisSlice 描述某个轴的源、目标范围。 +TensorField 描述本地 Tensor；AxisMapping 描述某个轴的本地、外部范围。 TransferRegion / StateTransferLayout 用于 State 与 CPU/PD 页之间的搬运； DP 在对应 TP 分片间直接复制，不需要字段映射类型。 @@ -27,8 +27,8 @@ # "big" 对应 big_pages_pool 与 bounded_att_big_page_allocator； # "small" 对应 small_pages_pool 与 bounded_att_small_page_allocator。 # 两个域中的同一个 slot_id 指向不同状态页。调用方先在对应 allocator 分配槽位， -# 再将类别传给 save_bounded_att_pages / restore_bounded_att_pages 等操作，例如： -# manager.save_bounded_att_pages(req_indices=[0], state_token_lens=[256], page_kind="big", slot_ids=[3]) +# 再将类别传给 save_bounded_att_runtime_to_pages / restore_bounded_att_runtime_from_pages 等操作，例如： +# manager.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[256], page_kind="big", page_slot_ids=[3]) # 上例假定请求槽位 0 的已接受长度为 256，且大页槽位 3 已由调用方分配。 PageKind = Literal["big", "small"] @@ -65,36 +65,36 @@ class TensorField: @dataclass(frozen=True) -class AxisSlice: +class AxisMapping: """某一轴上的逻辑连续范围映射，起点与长度均以元素为单位。 TransferRegion 中的方向固定为本地 State 到外部页，读取时反向应用。 环形轴按对应 Tensor 的实际轴长度取模，不要求两侧容量相同。 示例：本地 head[0:2] 写到全局传输区域 head[2:4]，读回时反向复制。 - >>> head_slice = AxisSlice(axis="head", src_start=0, dst_start=2, length=2) - >>> assert head_slice.dst_start + head_slice.length == 4 + >>> head_mapping = AxisMapping(axis="head", local_start=0, external_start=2, length=2) + >>> assert head_mapping.external_start + head_mapping.length == 4 环形示例：本地 window 容量为 4，从槽位 3 起搬运 4 个位置到线性外部页。 本地位置依次为 [3, 0, 1, 2]，外部位置依次为 [0, 1, 2, 3]。 - >>> window_slice = AxisSlice( - ... axis="window", src_start=3, dst_start=0, length=4, src_circular=True, + >>> window_mapping = AxisMapping( + ... axis="window", local_start=3, external_start=0, length=4, local_circular=True, ... ) - >>> assert window_slice.src_circular and not window_slice.dst_circular + >>> assert window_mapping.local_circular and not window_mapping.external_circular """ - # 要切片的轴名，必须同时存在于源、目标字段中，例如 head、channel、dim。 - # layer 固定完整复制，token 由本次传输数量控制，这两个轴不放入 AxisSlice。 + # 要映射的轴名，必须同时存在于本地和外部字段中，例如 head、channel、dim。 + # layer 固定完整复制，token 由本次传输数量控制，这两个轴不放入 AxisMapping。 axis: str - # 源轴的起始元素索引，包含该位置且必须非负。 - # TransferRegion 中本地字段的位置。 - src_start: int + # 本地 State 轴的起始元素索引，包含该位置且必须非负。 + # 写出时从此处读取，读回时向此处写入。 + local_start: int - # 目标轴的起始元素索引，包含该位置且必须非负。 - # TransferRegion 中外部页区域的位置。 + # 外部页轴的起始元素索引，包含该位置且必须非负。 + # 写出时向此处写入，读回时从此处读取。 # 用于将本地 TP 分片放到全局 head/channel 的对应位置。 - dst_start: int + external_start: int # 搬运的元素数量，必须为正；普通轴使用半开区间 [start, start + length)。 # 环形轴对 start + i 取模，length 不得超过该侧 Tensor 的轴容量，避免重复覆盖。 @@ -103,11 +103,11 @@ class AxisSlice: # 本地 State 的该轴采用环形寻址；容量取自本地 Tensor，而不是外部 region.shape。 # 例如 tail/SWA 的 window 轴。读回时仍作用于本地侧，不随实际搬运方向交换含义。 - src_circular: bool = False + local_circular: bool = False # 外部区域的该轴采用环形寻址；容量取自外部 Tensor 的对应轴。 # 常规 CPU/PD 页按逻辑顺序线性保存，使用默认 False。 - dst_circular: bool = False + external_circular: bool = False # ======================================================================== @@ -120,7 +120,7 @@ class TransferRegion: """描述 State 内一个字段与 CPU/PD 外部页区域之间的拷贝映射。 由 StateTransferLayout.regions 组织，同一字段可以对应多个区域映射。 - 区域的 shape 覆盖外部完整数据，axis_slices 选择当前 rank 的子范围。 + 区域的 shape 覆盖外部完整数据，axis_mappings 选择当前 rank 的子范围。 本地和外部区域的层数必须一致，并按相同存储顺序完整复制。 示例：外部 K 区域有 4 个全局 heads，head_dim=128，当前 rank 持有后 2 个 heads。 @@ -131,9 +131,9 @@ class TransferRegion: ... dtype=torch.float16, ... axes=("layer", "token", "head", "dim"), ... shape=(2, 8, 4, 128), - ... axis_slices=( - ... AxisSlice(axis="head", src_start=0, dst_start=2, length=2), - ... AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + ... axis_mappings=( + ... AxisMapping(axis="head", local_start=0, external_start=2, length=2), + ... AxisMapping(axis="dim", local_start=0, external_start=0, length=128), ... ), ... write_enabled=True, ... ) @@ -169,10 +169,10 @@ class TransferRegion: # 描述本地字段到外部区域的轴范围，每个非 layer/token 轴必须恰好出现一次。 # 无分片的轴也显式填写完整范围，例如 dim: 0 -> 0、length=head_dim。 - # 读回时交换源、目标起点；同一组切片应用于全部层和本次有效 token。 + # 两个方向共用同一组本地/外部起点，仅反转数据流向；映射应用于全部层和本次有效 token。 # 环形寻址由支持它的具体 State 搬运实现处理；请求窗口起点通过本次调用的 - # axis_slice_overrides 覆盖，不回写缓存布局，起点也不作为 kernel 的编译期常量。 - axis_slices: tuple[AxisSlice, ...] + # axis_mapping_overrides 覆盖，不回写缓存布局，起点也不作为 kernel 的编译期常量。 + axis_mappings: tuple[AxisMapping, ...] # 当前 TP rank 是否负责向外部区域写入；False 不影响从该区域读回。 # KV head 有 TP 副本时，每组只选一个 rank 写入；复制的 indexer tail 只由 rank 0 写。 @@ -184,8 +184,8 @@ class TransferRegion: class StateTransferLayout: """描述一个 State 的 CPU/PD 传输布局，按字段组织 TransferRegion。 - 每个 layout 通过 resource_id 绑定一个 State；regions 描述该 State 各字段 - 与外部页之间的映射，具体轴范围由各 region 的 AxisSlice 描述。 + 每个 layout 通过 state_name 绑定一个 State；regions 描述该 State 各字段 + 与外部页之间的映射，具体轴范围由各 region 的 AxisMapping 描述。 byte_start / byte_end 记录本 State 在联合页中的完整预留区间，包含自身对齐空间。 CPU/PD 根据页集合的行选择物理页；PD 通过显式 pd_page_index 指定目标页。 实际页容量由 manager 统一计算，可以大于当前 State 的终点。 @@ -195,13 +195,13 @@ class StateTransferLayout: >>> k_region = TransferRegion( ... field="k", byte_start=16384, byte_end=32768, dtype=torch.float16, ... axes=("layer", "token", "head", "dim"), shape=(2, 8, 4, 128), - ... axis_slices=( - ... AxisSlice(axis="head", src_start=0, dst_start=2, length=2), - ... AxisSlice(axis="dim", src_start=0, dst_start=0, length=128), + ... axis_mappings=( + ... AxisMapping(axis="head", local_start=0, external_start=2, length=2), + ... AxisMapping(axis="dim", local_start=0, external_start=0, length=128), ... ), ... ) >>> layout = StateTransferLayout( - ... resource_id="main_gqa_fp16_state", + ... state_name="main_gqa_fp16_state", ... tokens_per_page=8, ... byte_start=16384, byte_end=49152, ... regions=(k_region, replace(k_region, field="v", byte_start=32768, byte_end=49152)), @@ -209,9 +209,9 @@ class StateTransferLayout: >>> assert layout.regions[1].byte_end == layout.byte_end == 49152 """ - # 当前 State 的规范身份，与管理器成员名、state.resource_id 一致。 + # 当前 State 的规范身份，与管理器成员名、state.state_name 一致。 # 用于从联合页布局字典中选择对应资源，区分 main/draft 及不同存储方案。 - resource_id: str + state_name: str # 整张外部页可容纳的 token 数，不是本次实际传输数量；含 KV 的页使用正整数。 # CPU/PD 分别使用各自的全局页配置，页容量不随本次传输长度改变。 diff --git a/lightllm/common/att_resources/utils.py b/lightllm/common/attention_state/utils.py similarity index 100% rename from lightllm/common/att_resources/utils.py rename to lightllm/common/attention_state/utils.py diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index f2233c8128..4e33d187a0 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -43,8 +43,8 @@ def init_state(self): super().init_state() batch_size = self.infer_state.batch_size state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", ) offline_scales: torch.Tensor = state.scales @@ -99,8 +99,8 @@ def _fp8_prefill_att( size=(-1, 1, k_head_num, k_head_dim), ).view(torch.float8_e4m3fn) state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", ) storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) q = q.reshape(q.shape[0], k_head_num, -1) @@ -144,8 +144,8 @@ def init_state(self): self.backend: Fp8Fa3AttBackend = self.backend state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", ) super().init_state() @@ -206,8 +206,8 @@ def _fp8_decode_att( ).view(torch.float8_e4m3fn) state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_head_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", ) storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index e16e87ab97..260333e191 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -73,8 +73,8 @@ def _fp8_prefill_att( k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_tensor_state", ) storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) k_descale = state.cpu_scales[storage_layer_index][0] @@ -131,8 +131,8 @@ def _fp8_decode_att( k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) state = getattr( - self.infer_state.att_resource_manager, - f"{self.infer_state.att_resource_role}_gqa_fp8_static_per_tensor_state", + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_tensor_state", ) storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 7e73bf58fe..81a740ffa9 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.bounded_att import GDNState +from lightllm.common.attention_state.bounded_att import GDNState import dataclasses import weakref import torch @@ -140,8 +140,8 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: LinearAttBackend = self.backend - state = self.infer_state.att_resource_manager.get_att_state( - model_name=self.infer_state.att_resource_role, att_type="gdn", + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, state_type_filter="gdn", logical_layer_id=layer_num, state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) @@ -220,8 +220,8 @@ def init_state(self): # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 (state,) = ( - state for state in self.infer_state.att_resource_manager.get_bounded_attention_states( - self.infer_state.att_resource_role + state for state in self.infer_state.att_state_manager.get_bounded_attention_states( + self.infer_state.att_state_role ) if isinstance(state, GDNState) ) @@ -296,8 +296,8 @@ def decode_att( backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - state = self.infer_state.att_resource_manager.get_att_state( - model_name=self.infer_state.att_resource_role, att_type="gdn", + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, state_type_filter="gdn", logical_layer_id=layer_num, state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 60cf21e96f..0b7f7eaa1d 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,7 +10,7 @@ import torch -from lightllm.common.att_resources.bounded_att import KDAState +from lightllm.common.attention_state.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, @@ -84,8 +84,8 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_resource_manager.get_att_state( - model_name=self.infer_state.att_resource_role, att_type="kda", + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, state_type_filter="kda", logical_layer_id=layer_num, state_class=KDAState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) @@ -144,8 +144,8 @@ def init_state(self): # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 (state,) = ( - state for state in self.infer_state.att_resource_manager.get_bounded_attention_states( - self.infer_state.att_resource_role + state for state in self.infer_state.att_state_manager.get_bounded_attention_states( + self.infer_state.att_state_role ) if isinstance(state, KDAState) ) @@ -219,8 +219,8 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_resource_manager.get_att_state( - model_name=self.infer_state.att_resource_role, att_type="kda", + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, state_type_filter="kda", logical_layer_id=layer_num, state_class=KDAState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 043973a622..b7f1013a2f 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -14,7 +14,7 @@ from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.req_manager import ReqManager -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, @@ -49,9 +49,9 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState - from lightllm.common.att_resources.bounded_att import BoundedAttentionState - from lightllm.common.att_resources.att_tail import AttentionTailState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.att_tail import AttentionTailState logger = init_logger(__name__) @@ -124,9 +124,9 @@ def __init__(self, kvargs): return @property - def att_resource_role(self) -> str: + def att_state_role(self) -> str: """当前模型的资源前缀；每个 draft 独立拥有 State。""" - from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS if not self.is_mtp_draft_model: return "main" @@ -142,16 +142,16 @@ def bind_req_manager(self, req_manager: ReqManager) -> None: """绑定启动流程创建的共享请求管理器；main/draft 共用请求槽位和 token 索引映射。""" self.req_manager = req_manager - def prepare_att_resources(self) -> None: + def prepare_attention_states(self) -> None: """main 创建共享管理器，draft 引用同一实例，再准备当前模型自己的 State。""" - assert getattr(self, "att_resource_manager", None) is None, "resources already initialized" + assert getattr(self, "att_state_manager", None) is None, "resources already initialized" if self.is_mtp_draft_model: - self.att_resource_manager = self.main_model.att_resource_manager - assert not self.att_resource_manager._states_created, "State creation is already finished" - assert self.max_req_num + 1 == self.att_resource_manager._request_capacity, "models must share the request domain" + self.att_state_manager = self.main_model.att_state_manager + assert not self.att_state_manager._state_registration_finished, "State registration is already finished" + assert self.max_req_num + 1 == self.att_state_manager._request_capacity, "models must share the request domain" with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): if not self.is_mtp_draft_model: - self.att_resource_manager = AttResourceManager(request_capacity=self.max_req_num + 1) + self.att_state_manager = AttentionStateManager(request_capacity=self.max_req_num + 1) # 无 attention 的 draft 只共享管理器，不创建自身 State。 if self.is_mtp_draft_model and self.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): return @@ -173,16 +173,16 @@ def prepare_att_resources(self) -> None: device="cuda", ) for state in (*full_states, *bounded_states, *tail_states): - self.att_resource_manager.set_att_state(self.att_resource_role, state) + self.att_state_manager.register_state(self.att_state_role, state) def init_runtime(self) -> None: """Initialize this model's backends, warmup and graphs after shared allocation.""" - self.req_manager.bind_hold_page(self.att_resource_manager.full_att_allocator.hold_token_indices) + self.req_manager.bind_hold_page(self.att_state_manager.full_att_allocator.hold_token_indices) self._check_mem_size() self._init_att_backend() self._init_att_backend1() self._init_hidden_collector() - role = self.att_resource_role + role = self.att_state_role logger.info( f"{role} prefill/decode attention: {type(self.prefill_att_backend).__name__} / " f"{type(self.decode_att_backend).__name__}" @@ -295,7 +295,7 @@ def create_att_tail_states( return [] def _check_mem_size(self): - self.max_total_token_num = self.att_resource_manager.full_att_allocator.size + self.max_total_token_num = self.att_state_manager.full_att_allocator.size assert ( self.max_total_token_num > self.batch_max_tokens @@ -335,8 +335,8 @@ def _init_some_value(self): return def _init_att_backend(self): - role = self.att_resource_role - if not self.att_resource_manager.get_full_attention_states(role): + role = self.att_state_role + if not self.att_state_manager.get_full_attention_states(role): # FFN-only and linear-only models do not need a full-attention backend. self.prefill_att_backend = self.decode_att_backend = None return @@ -432,7 +432,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.is_draft_model = self.is_mtp_draft_model - infer_state.att_resource_role = self.att_resource_role + infer_state.att_state_role = self.att_state_role infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num @@ -456,7 +456,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.multimodal_params = model_input.multimodal_params infer_state.req_manager = self.req_manager - infer_state.att_resource_manager = self.att_resource_manager + infer_state.att_state_manager = self.att_state_manager infer_state.mem_index = self._select_mem_indexes(model_input) infer_state.microbatch_index = microbatch_index @@ -1077,9 +1077,9 @@ def _check_max_len_infer(self): logger.info("begin check max_len infer") dummy_input_ids = torch.ones(self.batch_max_tokens, dtype=torch.int64, device="cuda") b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.att_resource_manager.full_att_allocator.page_size + page_size = self.att_state_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.att_resource_manager.full_att_allocator.alloc(alloc_token_num).cuda() + mem_indexes = self.att_state_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = self.batch_max_tokens @@ -1113,7 +1113,7 @@ def _check_max_len_infer(self): torch.argmax(prob_out, dim=1, keepdim=True) prob_out = None self.req_manager.free_all() - self.att_resource_manager.full_att_allocator.free_all() + self.att_state_manager.full_att_allocator.free_all() logger.info(f"check max_len {self.batch_max_tokens} infer ok") except (RuntimeError, torch.OutOfMemoryError) as e: logger.exception(str(e)) @@ -1158,9 +1158,9 @@ def _autotune_warmup(self): 0, 10000, (input_len,), dtype=torch.int64, device="cuda", generator=rand_gen ) b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.att_resource_manager.full_att_allocator.page_size + page_size = self.att_state_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.att_resource_manager.full_att_allocator.alloc(alloc_token_num).cuda() + mem_indexes = self.att_state_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = input_len @@ -1194,14 +1194,14 @@ def _autotune_warmup(self): ) del model_output self.req_manager.free_all() - self.att_resource_manager.full_att_allocator.free_all() + self.att_state_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() except Exception as e: logger.warning(f"autotune warmup for length {input_len} failed: {str(e)}") logger.exception(str(e)) self.req_manager.free_all() - self.att_resource_manager.full_att_allocator.free_all() + self.att_state_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index d79ef0e245..ab0e0c0721 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -292,7 +292,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.att_resource_manager.full_att_allocator.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -354,7 +354,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1]) - model.att_resource_manager.full_att_allocator.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() del decode_batches diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 08af85b649..82a826fc54 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,7 +1,7 @@ import torch import triton import collections -from lightllm.common.att_resources.att_resource_manager import AttResourceManager +from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup from typing import Tuple, Any, Optional, List @@ -47,8 +47,8 @@ def __init__(self): self.is_prefill: bool = None self.is_draft_model: bool = False - self.att_resource_manager: AttResourceManager = None - self.att_resource_role: str = "main" # main 或 draft0 .. draft35，区分共享管理器中的 State。 + self.att_state_manager: AttentionStateManager = None + self.att_state_role: str = "main" # main 或 draft0 .. draft35，区分共享管理器中的 State。 self.req_manager: ReqManager = None self.mem_index: torch.Tensor = None diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index dedc153180..b65cab258a 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import os import torch import torch.distributed as dist @@ -34,9 +34,9 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[to raise Exception("need to impl") def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight): - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 8ffffb657e..4f21d90cfc 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -226,7 +226,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.att_resource_manager.full_att_allocator.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -289,7 +289,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_prefill(prefill_batches[0], prefill_batches[1]) - model.att_resource_manager.full_att_allocator.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() del prefill_batches diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index 5709657ffd..5ec00d542a 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import torch from typing import Tuple from lightllm.common.basemodel import TransformerLayerInferTpl @@ -29,9 +29,9 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -50,9 +50,9 @@ def _context_attention_kernel( def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 121174d513..0767c9f652 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class BloomTpPartModel(TpPartBaseModel): @@ -37,7 +37,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 61411d4eaf..986b8a703f 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import os import torch import triton @@ -104,9 +104,9 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -126,9 +126,9 @@ def _decompress_kv( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 76c760dc21..21593fe97d 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class Deepseek2TpPartModel(LlamaTpPartModel): @@ -69,7 +69,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index 6d5097a00f..e026d1ea8e 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo @@ -100,9 +100,9 @@ def _context_attention_kernel( }, ) - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -148,9 +148,9 @@ def _token_attention_kernel( }, ) - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -209,9 +209,9 @@ def _get_indices( q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_idx_, state_class=FullAttentionState, ) diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 2fba860e64..e899db92f4 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class Deepseek3_2TpPartModel(Deepseek2TpPartModel): @@ -25,7 +25,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index 5114544b5f..17f6246153 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState logger = init_logger(__name__) @@ -149,7 +149,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 5a3e427360..d1ebd30e30 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import math import torch import torch.nn as nn @@ -184,9 +184,9 @@ def _att_control(self): def _get_layer_kv(self, infer_state: InferStateInfo): # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 322ccf112b..e586d33fad 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -21,7 +21,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class Gemma4TpPartModel(LlamaTpPartModel): @@ -96,7 +96,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index 94c315507e..2e95f10d14 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class Gemma_2bTpPartModel(LlamaTpPartModel): @@ -47,7 +47,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index e06dbc2fe6..8a49767838 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import torch import triton @@ -38,15 +38,15 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.eps) gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) raw = torch.cat((k, gate), -1) - role = infer_state.att_resource_role - tail = infer_state.att_resource_manager.get_att_state( - model_name=role, - att_type=f"kpool_c{self.pool_size}", + role = infer_state.att_state_role + tail = infer_state.att_state_manager.get_state( + model_role=role, + state_type_filter=f"kpool_c{self.pool_size}", logical_layer_id=self.layer_idx, ).get_layer_view(logical_layer_id=self.layer_idx) - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_idx, state_class=FullAttentionState, ) diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 4374e4f51b..be7a2fa781 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -2,7 +2,7 @@ from __future__ import annotations -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import torch @@ -161,9 +161,9 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): layer_weight=layer_weight, ) del infer_state.get_topk_indices_params - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -248,9 +248,9 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="mla", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 658334f123..3b03fc0386 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -20,9 +20,9 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState - from lightllm.common.att_resources.bounded_att import BoundedAttentionState - from lightllm.common.att_resources.att_tail import AttentionTailState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.att_tail import AttentionTailState class Glm5NextTpPartModel(TpPartBaseModel): @@ -45,7 +45,7 @@ def create_bounded_att_states( device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" - from lightllm.common.att_resources.bounded_att import create_kda_state, ConvSSMStateConfig + from lightllm.common.attention_state.bounded_att import create_kda_state, ConvSSMStateConfig config = cls.normalize_model_config(model_config) assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" @@ -116,7 +116,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -153,7 +153,7 @@ def create_att_tail_states( device: str = "meta", ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" - from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + from lightllm.common.attention_state.att_tail import create_kpool_state, KPoolStateConfig from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 9a94863a23..9e96fd74f2 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -14,8 +14,8 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState - from lightllm.common.att_resources.att_tail import AttentionTailState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.att_tail import AttentionTailState class Glm5NextMTPModel(TpPartBaseModel): @@ -81,7 +81,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -116,7 +116,7 @@ def create_att_tail_states( device: str = "meta", ) -> list["AttentionTailState"]: """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" - from lightllm.common.att_resources.att_tail import create_kpool_state, KPoolStateConfig + from lightllm.common.attention_state.att_tail import create_kpool_state, KPoolStateConfig from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index 4fe03c943d..b26025c753 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState import torch from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo @@ -73,9 +73,9 @@ def _context_attention_kernel( window_size = (-1, -1) use_sliding_window = False - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -106,9 +106,9 @@ def _token_attention_kernel( window_size = (-1, -1) use_sliding_window = False - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index d9fca23aff..2bec35f191 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.basemodel.attention import AttControl import torch import triton @@ -46,9 +46,9 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -70,9 +70,9 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index be2c8448b2..2a33eda402 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class LlamaTpPartModel(TpPartBaseModel): @@ -45,7 +45,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 0de2177240..b4dea62592 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class MistralTpPartModel(TpPartBaseModel): @@ -39,7 +39,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 51a5b0b1f2..b90658e289 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class MixtralTpPartModel(TpPartBaseModel): @@ -47,7 +47,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 3b2432c77c..dc7a0f685d 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -91,7 +91,7 @@ def _decode(self, model_input: ModelInput) -> ModelOutput: infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens infer_state.position_cos = torch.index_select(self._cos_cached, 0, position_ids) infer_state.position_sin = torch.index_select(self._sin_cached, 0, position_ids) - infer_state.att_resource_manager = self.att_resource_manager + infer_state.att_state_manager = self.att_state_manager infer_state.is_draft_model = True infer_state.mem_index = self._select_mem_indexes(model_input) diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b4001b4044..7a631622d6 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -1,4 +1,4 @@ -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.basemodel.attention import AttControl from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv @@ -14,9 +14,9 @@ def __init__(self, layer_num, network_config): self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) def _token_attention_kernel(self, q, infer_state, layer_weight): - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) @@ -42,9 +42,9 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): - state = infer_state.att_resource_manager.get_att_state( - model_name=infer_state.att_resource_role, - att_type="gqa", + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState, ) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 79415952fe..899d0ddf8d 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -16,8 +16,8 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState - from lightllm.common.att_resources.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState logger = init_logger(__name__) @@ -51,7 +51,7 @@ def create_bounded_att_states( device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" - from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig + from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig config = cls.normalize_model_config(model_config) assert config["num_attention_heads"] % tp_world_size == 0 @@ -106,7 +106,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -131,8 +131,8 @@ def create_full_att_states( return [state] def _init_att_backend1(self): - role = self.att_resource_role - if not self.att_resource_manager.get_bounded_attention_states(role): + role = self.att_state_role + if not self.att_state_manager.get_bounded_attention_states(role): self.prefill_att_backend1 = None self.decode_att_backend1 = None return diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index ce97049e02..49d295e298 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class StarcoderTpPartModel(TpPartBaseModel): @@ -35,7 +35,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 24a21cf1e0..35ed6d35e2 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.att_resources.full_att import FullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState class Starcoder2TpPartModel(TpPartBaseModel): @@ -38,7 +38,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index bf563ae791..52e27992ec 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -194,7 +194,7 @@ class PDChunckedTransTask: # att_state: 混合注意力模型的请求运行态 buffer（如 linear attention 的 conv/SSM 状态）。 # start_kv_index == end_kv_index 标记 prefill 状态对应的 token 位置，mem_indexes 为空。 # tail PD 读写用该位置在固定逻辑窗口与本地 MTP ring 之间转换，不另存长度字段。 - # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 AttResourceManager 负责。 + # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 AttentionStateManager 负责。 page_kind: str = "kv" # Only valid for the local task owner; remote notify copies may carry the sender-local req_idx. req_idx: Optional[int] = None diff --git a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py index 8f0a8c6a87..29898f122a 100644 --- a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py @@ -116,7 +116,7 @@ def __init__( rank_in_node: int, hash_page_size: int, big_page_num: int, - att_resource_manager, + att_state_manager, ): assert hash_page_size >= 1, "hash_page_size must be >= 1" assert big_page_num >= 1, "big_page_num must be >= 1" @@ -126,7 +126,7 @@ def __init__( self.big_page_tokens = hash_page_size * big_page_num self.total_token_num = total_token_num - self.att_resource_manager = att_resource_manager + self.att_state_manager = att_state_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 @@ -146,14 +146,14 @@ def __init__( self.tree_total_tokens_num = SharedArray(f"tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64) self.tree_total_tokens_num.arr[0] = 0 - def _release_pages(self, page_kind, slot_ids): - resources = self.att_resource_manager + def _release_pages(self, page_kind, page_slot_ids): + resources = self.att_state_manager allocator = ( resources.bounded_att_big_page_allocator if page_kind == "big" else resources.bounded_att_small_page_allocator ) - allocator.free(slot_ids) + allocator.free(page_slot_ids) def _discard_node(self, node: HybridAttPagedTreeNode): if node.is_leaf(): @@ -305,7 +305,7 @@ def _insert_helper( self._add_node(child) else: # 节点已有 checkpoint，释放本次重复申请的槽位。 - self._release_pages("small", slot_ids=[block_state_idxs[0]]) + self._release_pages("small", page_slot_ids=[block_state_idxs[0]]) sub_prefix_len, ans_node = self._insert_helper( child, @@ -587,8 +587,8 @@ def _print_helper(self, node: HybridAttPagedTreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - if need_token_num > self.att_resource_manager.full_att_allocator.can_use_mem_size: - need_evict_token_num = need_token_num - self.att_resource_manager.full_att_allocator.can_use_mem_size + if need_token_num > self.att_state_manager.full_att_allocator.can_use_mem_size: + need_evict_token_num = need_token_num - self.att_state_manager.full_att_allocator.can_use_mem_size release_mems = [] small_page_buffer_ids = [] @@ -599,14 +599,14 @@ def release_mem(mem_index, small_page_buffer_id): self._evict(need_evict_token_num, release_mem) mem_index = torch.concat(release_mems) - self.att_resource_manager.full_att_allocator.free(mem_index) + self.att_state_manager.full_att_allocator.free(mem_index) small_page_buffer_ids = [idx for idx in small_page_buffer_ids if idx is not None] if len(small_page_buffer_ids) > 0: self._release_pages("small", small_page_buffer_ids) return def free_one_small_page_buffer(self): - allocator = self.att_resource_manager.bounded_att_small_page_allocator + allocator = self.att_state_manager.bounded_att_small_page_allocator if allocator is None: return if allocator.free_count > 0: @@ -618,7 +618,7 @@ def free_one_small_page_buffer(self): self._discard_node(node) assert node.small_page_buffer_idx is not None - self._release_pages("small", slot_ids=[node.small_page_buffer_idx]) + self._release_pages("small", page_slot_ids=[node.small_page_buffer_idx]) node.small_page_buffer_idx = None self._add_node(node) diff --git a/lightllm/server/router/dynamic_prompt/radix_cache.py b/lightllm/server/router/dynamic_prompt/radix_cache.py index f0d4b774b3..afc7df82af 100644 --- a/lightllm/server/router/dynamic_prompt/radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/radix_cache.py @@ -110,14 +110,14 @@ def match(t1: torch.Tensor, t2: torch.Tensor) -> int: class RadixCache: - def __init__(self, total_token_num, rank_in_node, att_resource_manager=None, page_size: int = 1): + def __init__(self, total_token_num, rank_in_node, att_state_manager=None, page_size: int = 1): self.total_token_num = total_token_num - self.att_resource_manager = att_resource_manager + self.att_state_manager = att_state_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 if page_size < 1: raise ValueError(f"page_size must be >= 1, got {page_size}") - allocator = None if att_resource_manager is None else att_resource_manager.full_att_allocator + allocator = None if att_state_manager is None else att_state_manager.full_att_allocator if allocator is not None and page_size != allocator.page_size: raise ValueError( f"RadixCache page_size {page_size} must match allocator page_size {allocator.page_size}" @@ -545,7 +545,7 @@ def _print_helper(self, node: TreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - allocator = self.att_resource_manager.full_att_allocator + allocator = self.att_state_manager.full_att_allocator if need_token_num > allocator.can_use_mem_size: need_evict_token_num = need_token_num - allocator.can_use_mem_size release_mems = [] diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index f2c089cc0e..b2de4fa93b 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -27,7 +27,7 @@ from lightllm.utils.profiler import ProfilerCmd from lightllm.server.router.token_load import TokenLoad from lightllm.server.metrics.manager import MetricClient -from lightllm.common.att_resources.allocators import KVCapacityReader +from lightllm.common.attention_state.allocators import KVCapacityReader from lightllm.utils.graceful_utils import graceful_registry from lightllm.utils.process_check import start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index cf04e3be4f..ebb06ce913 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -8,7 +8,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union from lightllm.common.req_manager import ReqManager -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.utils.infer_utils import mark_start, mark_end from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode @@ -39,7 +39,7 @@ class InferenceContext: backend: Optional["ModeBackend"] = None req_manager: ReqManager = None # gpu 请求管理 - att_resource_manager: AttResourceManager = None + att_state_manager: AttentionStateManager = None radix_cache: Union[HybridAttPagedRadixCache, RadixCache] = None shm_req_manager: ShmReqManager = None # 共享内存请求对象管理 requests_mapping: Dict[int, "InferReq"] = None @@ -65,7 +65,7 @@ def register( self.args = get_env_start_args() self.backend: ModeBackend = backend self.req_manager = req_manager - self.att_resource_manager = backend.model.att_resource_manager + self.att_state_manager = backend.model.att_state_manager self.req_sampling_manager = self.req_manager.req_sampling_params_manager self.radix_cache = radix_cache self.shm_req_manager = shm_req_manager @@ -176,11 +176,11 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 if req.tail_small_page_buffer_id is not None: - self.att_resource_manager.bounded_att_small_page_allocator.free([req.tail_small_page_buffer_id]) + self.att_state_manager.bounded_att_small_page_allocator.free([req.tail_small_page_buffer_id]) req.tail_small_page_buffer_id = None # 释放请求执行期间申请、但不再插入 radix cache 的大页状态。 if req.hybrid_len_to_big_page_id: - self.att_resource_manager.bounded_att_big_page_allocator.free( + self.att_state_manager.bounded_att_big_page_allocator.free( list(req.hybrid_len_to_big_page_id.values()) ) req.hybrid_len_to_big_page_id.clear() @@ -296,7 +296,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): # big page 模式下请求在 prefill 跨过 big page 边界后、到达末尾前被 pause / abort。 # 若不释放，会泄漏 big page state slot，并触发 free_a_req_mem 中 dict 为空的断言。 if req.hybrid_len_to_big_page_id: - self.att_resource_manager.bounded_att_big_page_allocator.free( + self.att_state_manager.bounded_att_big_page_allocator.free( list(req.hybrid_len_to_big_page_id.values()) ) req.hybrid_len_to_big_page_id.clear() @@ -340,7 +340,7 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool if free_token_index: free_token_index = custom_cat(free_token_index) - self.att_resource_manager.full_att_allocator.free(free_token_index) + self.att_state_manager.full_att_allocator.free(free_token_index) self.req_manager.free(free_req_index) finished_req_ids_set = set(finished_request_ids) @@ -351,8 +351,8 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool f"free a batch state:\n" f"radix refed token num {self.radix_cache.get_refed_tokens_num()}\n" f"radix hold token num {self.radix_cache.get_tree_total_tokens_num()}\n" - f"KV allocator available tokens {self.att_resource_manager.full_att_allocator.can_use_mem_size}\n" - f"KV allocator total tokens {self.att_resource_manager.full_att_allocator.size}\n" + f"KV allocator available tokens {self.att_state_manager.full_att_allocator.can_use_mem_size}\n" + f"KV allocator total tokens {self.att_state_manager.full_att_allocator.size}\n" ) return @@ -381,7 +381,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if len(free_token_index) != 0: free_token_index = custom_cat(free_token_index) - self.att_resource_manager.full_att_allocator.free(free_token_index) + self.att_state_manager.full_att_allocator.free(free_token_index) return self def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool): @@ -415,7 +415,7 @@ def get_can_alloc_token_num(self): radix_cache_unref_token_num = ( self.radix_cache.get_tree_total_tokens_num() - self.radix_cache.get_refed_tokens_num() ) - return self.att_resource_manager.full_att_allocator.can_use_mem_size + radix_cache_unref_token_num + return self.att_state_manager.full_att_allocator.can_use_mem_size + radix_cache_unref_token_num def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" @@ -434,7 +434,7 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR # 完整图片可能跨过多个大页结束点，例如页长 256、chunk [128, 768) 覆盖 256/512/768。 # 这些页的 KV 都有效，仍需为 radix 插入和 CPU offload 分配对应槽位。 for page_end in range(first_page_end, last_page_end + 1, big_page_token_num): - slots = self.att_resource_manager.bounded_att_big_page_allocator.alloc(1) + slots = self.att_state_manager.bounded_att_big_page_allocator.alloc(1) assert slots is not None big_page_id = slots[0] assert page_end not in req.hybrid_len_to_big_page_id @@ -447,11 +447,11 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR assert len(b_req_idx) == len(big_page_buffer_ids) if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): selected = [(req, slot) for req, slot in zip(reqs, big_page_buffer_ids) if slot != -1] - self.att_resource_manager.save_bounded_att_pages( + self.att_state_manager.save_bounded_att_runtime_to_pages( req_indices=[req.req_idx for req, _ in selected], state_token_lens=[req.get_chuncked_input_token_len() for req, _ in selected], page_kind="big", - slot_ids=[slot for _, slot in selected], + page_slot_ids=[slot for _, slot in selected], ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -463,16 +463,16 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR assert req.tail_small_page_buffer_id is None if req.hybrid_cache_len % big_page_token_num != 0: self.radix_cache.free_one_small_page_buffer() - allocator = self.att_resource_manager.bounded_att_small_page_allocator + allocator = self.att_state_manager.bounded_att_small_page_allocator slots = allocator.alloc(1) if allocator is not None else None req.tail_small_page_buffer_id = slots[0] if slots else None if req.tail_small_page_buffer_id is not None: dst_buffer_idx = req.tail_small_page_buffer_id - self.att_resource_manager.save_bounded_att_pages( + self.att_state_manager.save_bounded_att_runtime_to_pages( req_indices=[req.req_idx], state_token_lens=[req.hybrid_cache_len], page_kind="small", - slot_ids=[dst_buffer_idx], + page_slot_ids=[dst_buffer_idx], ) return @@ -675,13 +675,13 @@ def _match_radix_cache(self): return def _restore_attention_page(self, node, page_kind): - resources = g_infer_context.att_resource_manager + resources = g_infer_context.att_state_manager slot = node.big_page_buffer_idx if page_kind == "big" else node.small_page_buffer_idx - resources.restore_bounded_att_pages( + resources.restore_bounded_att_runtime_from_pages( req_indices=[self.req_idx], state_token_lens=[node.node_prefix_total_len], page_kind=page_kind, - slot_ids=[slot], + page_slot_ids=[slot], ) resources.reset_attention_tail_state([self.req_idx]) @@ -744,7 +744,7 @@ def _hybrid_match_radix_cache(self): # 有充足的token 容量时 assert need_tokens % self.args.page_size == 0 radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) - tail_mems = g_infer_context.att_resource_manager.full_att_allocator.alloc(need_tokens) + tail_mems = g_infer_context.att_state_manager.full_att_allocator.alloc(need_tokens) g_infer_context.req_manager.req_to_token_indexs[ self.req_idx, 0:cur_big_page_tokens ] = value_tensor[0:cur_big_page_tokens] @@ -753,7 +753,7 @@ def _hybrid_match_radix_cache(self): ] = tail_mems # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 - g_infer_context.att_resource_manager.copy_full_att_tokens( + g_infer_context.att_state_manager.copy_full_att_tokens( value_tensor[cur_big_page_tokens:shared_kv_len].to(device="cuda", non_blocking=True), tail_mems.to(device="cuda", non_blocking=True), ) @@ -799,8 +799,8 @@ def _hybrid_match_radix_cache(self): if self.cur_kv_len == 0: # 说明没有任何命中 - g_infer_context.att_resource_manager.reset_bounded_att_state([self.req_idx]) - g_infer_context.att_resource_manager.reset_attention_tail_state([self.req_idx]) + g_infer_context.att_state_manager.reset_bounded_att_state([self.req_idx]) + g_infer_context.att_state_manager.reset_attention_tail_state([self.req_idx]) return def is_master_req(self): diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 98d26db78f..c51668035b 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -149,11 +149,11 @@ def init_model(self, kvargs): self.init_mtp_draft_model(model_kvargs) from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag - from lightllm.common.att_resources import AttResourceManager + from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager models = (self.model, *self.draft_models) - assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" assert all(item.max_req_num == self.model.max_req_num for item in models), "models must share the request domain" # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) @@ -162,13 +162,13 @@ def init_model(self, kvargs): for item in models: item.bind_req_manager(req_manager=req_manager) for item in models: - item.prepare_att_resources() - manager: AttResourceManager = self.model.att_resource_manager - assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + item.prepare_attention_states() + manager: AttentionStateManager = self.model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - manager.finish_state_creation() + manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) - manager.initialize_resources( + manager.initialize_storage( self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) for item in models: @@ -176,24 +176,24 @@ def init_model(self, kvargs): torch.cuda.empty_cache() set_model_init_status(True) set_random_seed(2147483647) - self.is_hybrid_att_model = bool(self.model.att_resource_manager.get_bounded_attention_states()) + self.is_hybrid_att_model = bool(self.model.att_state_manager.get_bounded_attention_states()) if not self.use_dynamic_prompt_cache: self.radix_cache = None else: if self.is_hybrid_att_model: self.radix_cache = HybridAttPagedRadixCache( - total_token_num=self.model.att_resource_manager.full_att_allocator.size, + total_token_num=self.model.att_state_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, hash_page_size=self.args.linear_att_hash_page_size, big_page_num=self.args.linear_att_page_block_num, - att_resource_manager=self.model.att_resource_manager, + att_state_manager=self.model.att_state_manager, ) else: self.radix_cache = RadixCache( - total_token_num=self.model.att_resource_manager.full_att_allocator.size, + total_token_num=self.model.att_state_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, - att_resource_manager=self.model.att_resource_manager, + att_state_manager=self.model.att_state_manager, page_size=self.args.page_size, ) @@ -240,12 +240,12 @@ def init_model(self, kvargs): # PD 和 DP 使用独立导出，避免 PD 读者收到大小页池或 DP 读者丢失页池。 if self.args.run_mode in ("prefill", "decode"): publish_att_resources( - self.model.att_resource_manager, self.model.req_manager.req_to_token_indexs, + self.model.att_state_manager, self.model.req_manager.req_to_token_indexs, kind="pd", reader_count=self.node_world_size, ) if self.args.enable_dp_prompt_cache_fetch: publish_att_resources( - self.model.att_resource_manager, self.model.req_manager.req_to_token_indexs, + self.model.att_state_manager, self.model.req_manager.req_to_token_indexs, kind="dp", reader_count=self.dp_size_in_node - 1, ) dist.barrier(group=self.node_nccl_group) @@ -254,7 +254,7 @@ def init_model(self, kvargs): # 非主 rank 不需要分配 buffer，避免重复占用内存。 if self.is_master_in_dp: # Capture 只保存 allocator 管理的真实 KV 槽位，HOLD 页对应的 padding 写入由 kernel 过滤。 - kv_cache_size = self.model.att_resource_manager.full_att_allocator.size + kv_cache_size = self.model.att_state_manager.full_att_allocator.size if self.args.enable_prompt_logprobs: mgr = PromptLogprobsCaptureManager.get_instance() if mgr is not None: @@ -310,24 +310,24 @@ def init_dp_kv_shared(self): ) # 每个 DP 只保留对应 TP 分片，统一以节点内 DP rank 索引。 - self.dp_att_resource_managers = {} + self.dp_att_state_managers = {} self.dp_req_to_token_indices = {} # 取节点内的 TP 起点，兼容单个 TP 组跨节点、rank_in_dp 超过本节点 rank 范围的情况。 for dp_rank, rank_idx in enumerate( range(self.rank_in_node % self.dp_world_size, self.node_world_size, self.dp_world_size) ): if rank_idx != self.rank_in_node: - att_resource_manager, req_to_token_indices = load_att_resources(rank_idx, kind="dp") - att_resource_manager.register_page_pools_pin_memory() + att_state_manager, req_to_token_indices = load_att_resources(rank_idx, kind="dp") + att_state_manager.register_page_pools_pin_memory() else: - att_resource_manager = self.model.att_resource_manager + att_state_manager = self.model.att_state_manager req_to_token_indices = self.model.req_manager.req_to_token_indexs - self.dp_att_resource_managers[dp_rank] = att_resource_manager + self.dp_att_state_managers[dp_rank] = att_state_manager self.dp_req_to_token_indices[dp_rank] = req_to_token_indices return def get_max_total_token_num(self): - return self.model.att_resource_manager.full_att_allocator.size + return self.model.att_state_manager.full_att_allocator.size def infer_loop(self): raise NotImplementedError() @@ -345,7 +345,7 @@ def init_mtp_draft_model(self, main_kvargs: dict): is_chained_draft = spec_mode in ("vanilla_with_att", "vanilla_no_att") draft_model_count = self.max_draft_step if is_chained_draft else 1 - from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS assert 0 < draft_model_count <= MAX_DRAFT_MODELS, "at most 36 draft models are supported" draft_model_dirs = self.args.mtp_draft_model_dir @@ -666,7 +666,7 @@ def _alloc_req_kv_mem( old_hold_kv_len = req_obj.hold_kv_len new_hold_kv_len = old_hold_kv_len + alloc_token_num - mem_indexes = g_infer_context.att_resource_manager.full_att_allocator.alloc(alloc_token_num) + mem_indexes = g_infer_context.att_state_manager.full_att_allocator.alloc(alloc_token_num) # 高频调度路径允许异步写入请求索引表，其他调用方默认保持原有的同步拷贝语义。 g_infer_context.req_manager.req_to_token_indexs[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( mem_indexes, non_blocking=no_blcoking_copy diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index af49f208e4..57d002049a 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -182,7 +182,7 @@ def _master_req_to_radix_cache(self, master_req: InferReq): prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value) old_prefix_len = 0 if master_req.shared_kv_node is None else master_req.shared_kv_node.node_prefix_total_len assert old_prefix_len <= master_req.cur_kv_len - self.model.att_resource_manager.full_att_allocator.free( + self.model.att_state_manager.full_att_allocator.free( self.model.req_manager.req_to_token_indexs[master_req.req_idx][old_prefix_len:prefix_len] ) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 0979697e4a..90fc008b86 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -108,14 +108,14 @@ def build_shared_kv_trans_tasks( req_index, src_dp_rank, first_page:end_page ].tolist() assert all(page_id >= 0 for page_id in src_big_page_ids), "DP source is missing prefix big pages" - allocator = self.backend.model.att_resource_manager.bounded_att_big_page_allocator + allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" - dst_token_indexes = self.backend._alloc_req_kv_mem(req, alloc_token_num) - assert dst_token_indexes is not None + dst_token_indices = self.backend._alloc_req_kv_mem(req, alloc_token_num) + assert dst_token_indices is not None # 只搬运有效 KV；页尾预留槽位已写入请求表，不参与本次传输。 - dst_token_indexes = dst_token_indexes[:trans_size] - src_token_indexes = self.backend.dp_req_to_token_indices[src_dp_rank][ + dst_token_indices = dst_token_indices[:trans_size] + src_token_indices = self.backend.dp_req_to_token_indices[src_dp_rank][ src_req_idx, req.cur_kv_len:target_kv_len ] trans_tasks.append( @@ -123,8 +123,8 @@ def build_shared_kv_trans_tasks( dst_req=req, src_dp_rank=src_dp_rank, src_req_idx=src_req_idx, - src_token_indexes=src_token_indexes, - dst_token_indexes=dst_token_indexes, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, src_big_page_ids=src_big_page_ids, dst_big_page_ids=dst_big_page_ids, ) @@ -133,15 +133,15 @@ def build_shared_kv_trans_tasks( return trans_tasks def kv_trans(self, trans_tasks: List["TransTask"]) -> None: - target = self.backend.model.att_resource_manager + target = self.backend.model.att_state_manager for task in trans_tasks: - source = self.backend.dp_att_resource_managers[task.src_dp_rank] + source = self.backend.dp_att_state_managers[task.src_dp_rank] # 源索引来自远端 CUDA 请求表；allocator 返回 CPU 索引，在此一次性准备目标 CUDA int64 索引。 - dst_token_indexes = task.dst_token_indexes.to(device="cuda", dtype=torch.long, non_blocking=True) - target.copy_from_other_dp( + dst_token_indices = task.dst_token_indices.to(device="cuda", dtype=torch.long, non_blocking=True) + target.copy_from_dp( remote_att_manager=source, - src_token_indexes=task.src_token_indexes, - dst_token_indexes=dst_token_indexes, + src_token_indices=task.src_token_indices, + dst_token_indices=dst_token_indices, src_big_page_ids=task.src_big_page_ids, dst_big_page_ids=task.dst_big_page_ids, src_req_idx=task.src_req_idx, @@ -159,7 +159,7 @@ def kv_trans(self, trans_tasks: List["TransTask"]) -> None: for page_offset, page_id in enumerate(task.dst_big_page_ids): token_len = first_page_end + page_offset * self.big_page_tokens req.hybrid_len_to_big_page_id[token_len] = page_id - req.cur_kv_len += len(task.dst_token_indexes) + req.cur_kv_len += len(task.dst_token_indices) assert req.cur_kv_len <= req.hold_kv_len if self.backend.is_master_in_dp: req.shm_req.shm_cur_kv_len = req.cur_kv_len @@ -170,7 +170,7 @@ class TransTask: dst_req: InferReq # 本 DP 接收前缀的请求，持有目标请求槽位和新分配的资源。 src_dp_rank: int # 来源 DP 在当前节点内的编号。 src_req_idx: int # 远端请求槽位，用于复制 prefill 基准运行态。 - src_token_indexes: torch.Tensor # 来源 CUDA 请求表中的有效 KV 索引。 - dst_token_indexes: torch.Tensor # allocator 返回的目标 CPU 索引，提交时统一转为 CUDA int64。 + src_token_indices: torch.Tensor # 来源 CUDA 请求表中的有效 KV 索引。 + dst_token_indices: torch.Tensor # allocator 返回的目标 CPU 索引，提交时统一转为 CUDA int64。 src_big_page_ids: List[int] # 按前缀长度递增排列的源大页槽位。 dst_big_page_ids: List[int] # 与源列表一一配对的新分配目标大页槽位。 diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index ae4edf83c7..8723445951 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -34,7 +34,7 @@ def __init__(self, backend): self.cpu_cache_handle_queue: Deque[TransTask] = deque() self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False) - resources = self.backend.model.att_resource_manager + resources = self.backend.model.att_state_manager meta = self.cpu_cache_client.kv_cache_tensor_meta if meta.page_bytes != resources.get_cpu_page_bytes(): raise ValueError("CPU cache service and model resource page sizes differ") @@ -147,7 +147,7 @@ def read_req_from_cpu( """ if not cpu_page_indices: return - resources = self.backend.model.att_resource_manager + resources = self.backend.model.att_state_manager req_manager = self.backend.model.req_manager need_token_num = page_lengths[-1] - req.cur_kv_len @@ -181,7 +181,7 @@ def read_req_from_cpu( cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, cpu_page_indices=cpu_page_indices, token_indices=token_indices, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=big_page_slot_ids, req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, state_token_len=page_lengths[-1], @@ -321,7 +321,7 @@ def write_req_to_cpu( token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num] # 阶段三：只查找实际分配的 CPU 页所需快照；完整大页必须存在，不做缺失回退。 - resources = self.backend.model.att_resource_manager + resources = self.backend.model.att_state_manager big_page_slot_ids, small_page_slot_ids = None, None if self.backend.is_hybrid_att_model: page_tokens = self.args.cpu_cache_token_page_size @@ -343,8 +343,8 @@ def write_req_to_cpu( token_indices=token_indexes, cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, cpu_page_indices=page_list, - layouts=resources.get_cpu_transfer_layouts(), - cpu_page_readies=ready_list, + layouts=resources.get_cpu_page_layouts(), + cpu_page_ready_flags=ready_list, big_page_slot_ids=big_page_slot_ids, small_page_slot_ids=small_page_slot_ids, ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py index fa2183182f..6f1e90590e 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py @@ -111,7 +111,7 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: .detach() .cpu() ) - self.model.att_resource_manager.full_att_allocator.free(mem_indexes) + self.model.att_state_manager.full_att_allocator.free(mem_indexes) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len if self.is_master_in_dp: diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index b61c3bac6a..bffd16ea24 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -8,7 +8,7 @@ import pickle from typing import List, Dict, Union, Optional from lightllm.utils.log_utils import init_logger -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import ( @@ -87,7 +87,7 @@ def _init_env( # 从共享内存读取所有 rank 的 attention 资源 node_world_size = args.tp // args.nnodes - resource_managers: List[AttResourceManager] = [ + resource_managers: List[AttentionStateManager] = [ load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] @@ -121,7 +121,7 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - resource_managers: List[AttResourceManager], + resource_managers: List[AttentionStateManager], up_status_in_queue: Optional[mp.SimpleQueue], ): self.args = args @@ -177,7 +177,7 @@ def _warmup(self): copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, - write_to_pd_buffer=False, + write_to_pd_page=False, token_indices=token_indices, pd_page_index=0, dp_rank_in_node=dp_index, @@ -410,7 +410,7 @@ def read_page_to_mems_loop(self): copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, - write_to_pd_buffer=False, + write_to_pd_page=False, token_indices=token_indices, pd_page_index=trans_task.dst_page_index, dp_rank_in_node=trans_task.decode_dp_index, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index e59a2201b6..1f2544b8fe 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -8,7 +8,7 @@ import pickle from typing import List, Dict, Optional, Union from lightllm.utils.log_utils import init_logger -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask @@ -62,7 +62,7 @@ def _init_env( # 从共享内存读取所有 rank 的 attention 资源 node_world_size = args.tp // args.nnodes - resource_managers: List[AttResourceManager] = [ + resource_managers: List[AttentionStateManager] = [ load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] @@ -95,7 +95,7 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - resource_managers: List[AttResourceManager], + resource_managers: List[AttentionStateManager], ) -> None: self.args = args self.dp_world_size = self.args.tp // self.args.dp @@ -152,7 +152,7 @@ def _warmup(self): copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, - write_to_pd_buffer=True, + write_to_pd_page=True, token_indices=token_indices, pd_page_index=0, dp_rank_in_node=dp_index, @@ -217,7 +217,7 @@ def local_copy_kv_loop(self): copy_page( node_att_managers=self.resource_managers, pd_page_buffers=self.kv_move_buffer, - write_to_pd_buffer=True, + write_to_pd_page=True, token_indices=token_indices, pd_page_index=trans_task.src_page_index, dp_rank_in_node=trans_task.prefill_dp_index, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index d26e56c099..370a17694a 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -7,24 +7,24 @@ if TYPE_CHECKING: from collections.abc import Sequence - from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from lightllm.common.attention_state.manager import AttentionStateManager from typing import Literal -def allocate_transfer_buffer(resources: AttResourceManager, page_count: int, device: str | torch.device | int): +def allocate_transfer_buffer(resources: AttentionStateManager, page_count: int, device: str | torch.device | int): if page_count <= 0: raise ValueError("PD buffer must contain pages") return torch.empty((page_count, resources.get_pd_page_bytes()), dtype=torch.uint8, device=device) def copy_page( - node_att_managers: Sequence[AttResourceManager], + node_att_managers: Sequence[AttentionStateManager], pd_page_buffers: torch.Tensor, pd_page_index: int, dp_rank_in_node: int, tp_world_size: int, token_indices: torch.Tensor | None, - write_to_pd_buffer: bool, + write_to_pd_page: bool, task_kind: Literal["kv", "att_state"] = "kv", req_idx: int | None = None, state_token_len: int | None = None, @@ -34,7 +34,7 @@ def copy_page( node_att_managers 按节点内 DP 编号、组内 TP rank 顺序平铺。 pd_page_buffers 是 PD 页集合，pd_page_index 选择本次传输的页。 token_indices 在 kv 任务中为 CUDA 索引；att_state 任务传 None，不进行索引上传。 - write_to_pd_buffer=True 将 State 写入发送页；False 从接收页恢复 State。 + write_to_pd_page=True 将 State 写入发送页；False 从接收页恢复 State。 """ tp_att_managers = node_att_managers[dp_rank_in_node * tp_world_size : (dp_rank_in_node + 1) * tp_world_size] assert ( @@ -47,9 +47,9 @@ def copy_page( for att_manager in tp_att_managers: match task_kind: case "kv": - state_layouts = att_manager.get_pd_full_att_transfer_layouts() + state_layouts = att_manager.get_pd_token_page_layouts() case "att_state": - state_layouts = att_manager.get_pd_bounded_att_or_tail_transfer_layouts() + state_layouts = att_manager.get_pd_runtime_page_layouts() case _: assert False, f"unknown PD task kind {task_kind}" assert state_layouts and att_manager.get_pd_page_bytes() == pd_page_buffers.shape[1], ( @@ -60,8 +60,8 @@ def copy_page( for att_manager, state_layouts in zip(tp_att_managers, tp_transfer_layouts): match task_kind: case "kv": - if write_to_pd_buffer: - att_manager.write_to_pd_buffer( + if write_to_pd_page: + att_manager.write_to_pd_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -69,7 +69,7 @@ def copy_page( token_indices=token_indices, ) else: - att_manager.read_from_pd_buffer( + att_manager.read_from_pd_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -77,8 +77,8 @@ def copy_page( token_indices=token_indices, ) case "att_state": - if write_to_pd_buffer: - att_manager.write_to_pd_buffer( + if write_to_pd_page: + att_manager.write_to_pd_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -87,7 +87,7 @@ def copy_page( state_token_len=state_token_len, ) else: - att_manager.read_from_pd_buffer( + att_manager.read_from_pd_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 898a4e08fd..a32cedd17d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -32,7 +32,7 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - backend.model.att_resource_manager.commit_bounded_att_state( + backend.model.att_state_manager.commit_bounded_att_state( b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" ) return accept_lengths, accepted_index diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py index 88db8df336..dc03199294 100644 --- a/lightllm/server/router/model_infer/shared_att_resources_utils.py +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -14,11 +14,11 @@ if TYPE_CHECKING: import torch - from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from lightllm.common.attention_state.manager import AttentionStateManager def publish_att_resources( - resources: AttResourceManager, + resources: AttentionStateManager, request_mapping: torch.Tensor, kind: Literal["pd", "dp"], reader_count: int, @@ -56,7 +56,7 @@ def publish_att_resources( def load_att_resources( rank_in_node: int, kind: Literal["pd", "dp"], -) -> tuple[AttResourceManager, torch.Tensor]: +) -> tuple[AttentionStateManager, torch.Tensor]: """读取指定用途的资源；DP 接收方随后通过 manager 统一注册共享页池。""" assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index ecc24a1a9a..e2aaa921a6 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -3,19 +3,19 @@ from __future__ import annotations from typing import TYPE_CHECKING -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from collections.abc import Mapping - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.types import StateTransferLayout def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: """使用启动流程相同的 meta State 定义，由管理器统一生成 CPU 服务页布局。""" from transformers.configuration_utils import PretrainedConfig from lightllm.models import get_draft_model_class - from lightllm.common.att_resources.utils import MAX_DRAFT_MODELS + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS from lightllm.utils.llm_utils import get_llm_model_class # 阶段一：读取 main 和启用 attention 的 draft 配置，不构造实际模型。 @@ -31,7 +31,7 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) # 阶段二：复用模型定义入口创建 meta State，并累计逻辑层号；CPU 页不包含静态标定参数。 - manager = AttResourceManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 + manager = AttentionStateManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 logical_layer_start = 0 for role, model_class, config in models: config = model_class.normalize_model_config(config) @@ -49,13 +49,13 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: ) # tail 不占 CPU 页；无需创建 tail，也不需要实际模型或临时 CUDA 运行态。 for state in (*full_states, *bounded_states): - manager.set_att_state(role, state) + manager.register_state(role, state) logical_layer_start += count # 阶段三：确认混合模型的页边界一致，由管理器聚合 CPU 联合页布局。 - manager.finish_state_creation() + manager.finish_state_registration() if manager.get_bounded_attention_states(): assert ( args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num ), "hybrid CPU pages must match the large state-page token boundary" - return manager.get_cpu_transfer_layouts() + return manager.get_cpu_page_layouts() diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index ab62f5c4a7..961e65042e 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -58,10 +58,10 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP draft_models=draft_models, token_source=TokenSource(args), ) - resources = model.att_resource_manager + resources = model.att_state_manager (full_state,) = resources.get_full_attention_states("main") for draft in draft_models: - assert draft.att_resource_manager is resources + assert draft.att_state_manager is resources assert not hasattr(draft, "full_att_state") (draft_state,) = resources.get_full_attention_states("draft0") assert draft_state is not full_state @@ -96,7 +96,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP first_next_ids=executor._argmax_ids(output.logits), ) if hybrid: - resources.save_bounded_att_pages(source_req.tolist(), page_kind="big", slot_ids=[source_slots[page]]) + resources.save_bounded_att_runtime_to_pages(source_req.tolist(), page_kind="big", page_slot_ids=[source_slots[page]]) torch.cuda.synchronize() next_token = output.logits.argmax(dim=-1).cpu() source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() @@ -106,18 +106,18 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=source_indices, cpu_page_buffers=cpu, cpu_page_indices=pages, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), page_kind="big" if hybrid else None, - slot_ids=source_slots, + page_slot_ids=source_slots, ) # Host-side checkpoint repacking consumes the CPU page on the next call. torch.cuda.synchronize() - full_att_layouts = resources.get_pd_full_att_transfer_layouts() - runtime_layouts = resources.get_pd_bounded_att_or_tail_transfer_layouts() + full_att_layouts = resources.get_pd_token_page_layouts() + runtime_layouts = resources.get_pd_runtime_page_layouts() pd = {"kv": torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda")} if runtime_layouts: pd["att_state"] = torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - resources.write_to_pd_buffer( + resources.write_to_pd_page( pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", @@ -125,7 +125,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=source_indices, ) if "att_state" in pd: - resources.write_to_pd_buffer( + resources.write_to_pd_page( pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", @@ -145,7 +145,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP if kind == "gpu_pages": resources.copy_full_att_tokens(source_indices, indices) if hybrid: - resources.restore_bounded_att_pages(request.tolist(), page_kind="big", slot_ids=[source_slots[-1]]) + resources.restore_bounded_att_runtime_from_pages(request.tolist(), page_kind="big", page_slot_ids=[source_slots[-1]]) resources.reset_attention_tail_state(request.tolist()) elif kind == "cpu_pages": slots = big.alloc(len(pages)) if hybrid else None @@ -155,12 +155,12 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP cpu_page_buffers=cpu, cpu_page_indices=pages, token_indices=indices, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=slots, req_idx=int(request[0]) if getattr(resources, "cpu_tail_restore", False) else None, ) if hybrid: - resources.restore_bounded_att_pages(request.tolist(), page_kind="big", slot_ids=[slots[-1]]) + resources.restore_bounded_att_runtime_from_pages(request.tolist(), page_kind="big", page_slot_ids=[slots[-1]]) elif kind == "dp_copy": resources.copy_from_dp( resources, @@ -171,7 +171,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP dst_req_indices=request.tolist() if hybrid else None, ) else: - resources.read_from_pd_buffer( + resources.read_from_pd_page( pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", @@ -179,7 +179,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=indices, ) if "att_state" in pd: - resources.read_from_pd_buffer( + resources.read_from_pd_page( pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", @@ -312,11 +312,11 @@ def main(): draft_models = init_mtp_draft_models(args, kvargs, model) from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag - from lightllm.common.att_resources import AttResourceManager + from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager models = (model, *draft_models) - assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) @@ -325,13 +325,13 @@ def main(): for item in models: item.bind_req_manager(req_manager=req_manager) for item in models: - item.prepare_att_resources() - manager: AttResourceManager = model.att_resource_manager - assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + item.prepare_attention_states() + manager: AttentionStateManager = model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - manager.finish_state_creation() + manager.finish_state_registration() manager.load_full_att_calibration(args=model.args) - manager.initialize_resources( + manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) for item in models: diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index a3ccd92fc0..980c1f20e1 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -568,7 +568,7 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T req_id = int(req_id) max_target_len_by_req[req_id] = max(max_target_len_by_req.get(req_id, 0), int(target_len)) - allocator = self.model.att_resource_manager.full_att_allocator + allocator = self.model.att_state_manager.full_att_allocator page_size = allocator.page_size for req_id, target_len in max_target_len_by_req.items(): new_hold_kv_len = align_up(target_len, page_size) @@ -729,7 +729,7 @@ def _alloc_req_indexes(self, batch_size: int) -> torch.Tensor: return torch.tensor(req_indexes, dtype=torch.int32, device="cpu") def _reset_model_cache(self): - self.model.att_resource_manager.full_att_allocator.free_all() + self.model.att_state_manager.full_att_allocator.free_all() self.model.req_manager.free_all() self._hold_kv_len_by_req.clear() torch.cuda.synchronize() @@ -1273,11 +1273,11 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) draft_models = init_mtp_draft_models(args, model_kvargs, model) from lightllm.utils.envs_utils import set_model_init_status from lightllm.utils.torch_memory_saver_utils import MemoryTag - from lightllm.common.att_resources import AttResourceManager + from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager models = (model, *draft_models) - assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) @@ -1286,21 +1286,21 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) for item in models: item.bind_req_manager(req_manager=req_manager) for item in models: - item.prepare_att_resources() - manager: AttResourceManager = model.att_resource_manager - assert all(item.att_resource_manager is manager for item in models), "models must share one resource manager" + item.prepare_attention_states() + manager: AttentionStateManager = model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - manager.finish_state_creation() + manager.finish_state_registration() manager.load_full_att_calibration(args=model.args) - manager.initialize_resources( + manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) for item in models: item.init_runtime() torch.cuda.empty_cache() set_model_init_status(True) - cases = resolve_batch_max_prefill_cases(args, cases, model.att_resource_manager.full_att_allocator.size) - cases = resolve_profile_decode_cases(args, cases, model.att_resource_manager.full_att_allocator.size) + cases = resolve_batch_max_prefill_cases(args, cases, model.att_state_manager.full_att_allocator.size) + cases = resolve_profile_decode_cases(args, cases, model.att_state_manager.full_att_allocator.size) if defer_cudagraph: init_deferred_cudagraph(args, cases, model_kvargs, model) for draft in draft_models: diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 405a78e82f..db67d9aebb 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -11,15 +11,15 @@ import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, FullAttentionConfig, ) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import StatePageAllocator -from lightllm.common.att_resources.bounded_att import create_gdn_state, ConvSSMStateConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import set_env_start_args @@ -59,8 +59,8 @@ def main(): tp_world_size=8, ), ) - full.initialize_storage(tokens) - resources = AttResourceManager(request_capacity=1) + full.allocate_storage(tokens) + resources = AttentionStateManager(request_capacity=1) bounded_state = create_gdn_state( config=ConvSSMStateConfig( logical_layer_ids=tuple(layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval), @@ -79,27 +79,27 @@ def main(): device="cuda", ) bounded_state.initialize_page_pools(big_page_capacity=args.pages + 2, small_page_capacity=0) - resources.set_att_state("main", full) - resources.set_att_state("main", bounded_state) + resources.register_state("main", full) + resources.register_state("main", bounded_state) resources.full_att_allocator = FullAttKVAllocator(tokens) resources.bounded_att_big_page_allocator = StatePageAllocator(args.pages + 2, reserved_slots=2) resources.bounded_att_cpu_load_big_page_slot = args.pages resources.bounded_att_cpu_offload_big_page_slot = args.pages + 1 - resources.finish_state_creation() + resources.finish_state_registration() pages = list(range(args.pages)) slots = resources.bounded_att_big_page_allocator.alloc(args.pages) cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) # Populate every TP region through the production API before measuring restores. full.k_storage.zero_() full.v_storage.zero_() - resources.save_bounded_att_pages([0] * args.pages, page_kind="big", slot_ids=slots) + resources.save_bounded_att_runtime_to_pages([0] * args.pages, page_kind="big", page_slot_ids=slots) torch.cuda.synchronize() token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") resources.write_req_to_cpu( token_indices=token_indices, cpu_page_buffers=cpu_page_buffers, cpu_page_indices=pages, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=slots, ) torch.cuda.synchronize() @@ -109,7 +109,7 @@ def restore(): cpu_page_buffers=cpu_page_buffers, cpu_page_indices=pages, token_indices=token_indices, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=slots, req_idx=0, ) diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index b82a563c17..a252410165 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -5,11 +5,11 @@ from typing import TYPE_CHECKING from dataclasses import replace from functools import lru_cache -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.att_resources.full_att import MLAFullAttentionState -from lightllm.common.att_resources.bounded_att import GDNState -from lightllm.common.att_resources.att_tail import KPoolAttentionTailState +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.attention_state.full_att import MLAFullAttentionState +from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.attention_state.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state from unit_tests.common.att_resources.test_bounded_state import make_state from unit_tests.common.att_resources.test_tail_state import make_tail @@ -18,10 +18,10 @@ import torch -class MainDraftTestResourceManager(AttResourceManager): +class MainDraftTestResourceManager(AttentionStateManager): def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): super().__init__(request_capacity=3) - self.set_att_state("main", make_state(device=device)) + self.register_state("main", make_state(device=device)) draft_state = GDNState( config=replace(self.main_kda_conv_fp32_ssm_fp32_state.config, global_v_heads=2), request_capacity=3, @@ -29,14 +29,14 @@ def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = T device=device, ) draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") - self.set_att_state("draft0", draft_state) - self.set_att_state("main", full_state(device=device)) + self.register_state("draft0", draft_state) + self.register_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( self.main_gqa_fp32_state.config, page_size=self.main_gqa_fp32_state.page_size, device=device, ) - prepared_full_state.initialize_storage(self.main_gqa_fp32_state.capacity) - self.set_att_state("draft0", prepared_full_state) + prepared_full_state.allocate_storage(self.main_gqa_fp32_state.capacity) + self.register_state("draft0", prepared_full_state) self.att_tail_type = f"kpool_c{tail_pool_size}_fp32" if with_tail: tail = make_tail(device=device) @@ -46,20 +46,20 @@ def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = T mtp_step=tail.mtp_step, device=tail.tail_buffer.device, ) - self.set_att_state("main", state) + self.register_state("main", state) self.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) self.bounded_att_big_page_allocator = StatePageAllocator(3) self.bounded_att_small_page_allocator = StatePageAllocator(2) - self.finish_state_creation() + self.finish_state_registration() @lru_cache(maxsize=None) - def get_cpu_transfer_layouts(self): + def get_cpu_page_layouts(self): raise NotImplementedError("This test manager uses explicitly supplied CPU layouts") @lru_cache(maxsize=None) - def get_pd_full_att_transfer_layouts(self): + def get_pd_token_page_layouts(self): raise NotImplementedError("This test manager uses explicitly supplied PD layouts") @lru_cache(maxsize=None) - def get_pd_bounded_att_or_tail_transfer_layouts(self): + def get_pd_runtime_page_layouts(self): raise NotImplementedError("This test manager uses explicitly supplied PD layouts") diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 81d8201b35..54c19df896 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -8,8 +8,8 @@ from types import SimpleNamespace import torch -from lightllm.common.att_resources.types import TensorField -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.types import TensorField +from lightllm.common.attention_state.full_att import ( FullAttentionState, FullAttentionConfig, StaticFP8FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, ) @@ -17,9 +17,9 @@ if TYPE_CHECKING: from lightllm.models.qwen3next.config import Qwen3NextResourceConfig - from lightllm.common.att_resources.allocators import FullAttKVAllocator + from lightllm.common.attention_state.allocators import FullAttKVAllocator from collections.abc import Mapping, Sequence - from lightllm.common.att_resources.att_resource_manager import AttResourceManager + from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -73,12 +73,12 @@ def rebuild_full_state( scales=template.scales if scales is None else scales, q_scales=template.q_scales if q_scales is None else q_scales, ) - prepared_full_state.initialize_storage(token_capacity) + prepared_full_state.allocate_storage(token_capacity) return prepared_full_state prepared_full_state = type(template)( config, page_size=page_size, device=device, ) - prepared_full_state.initialize_storage(token_capacity) + prepared_full_state.allocate_storage(token_capacity) return prepared_full_state @@ -87,12 +87,12 @@ def owned_states(attributes: Mapping[str, object], state_class: type | tuple[typ return tuple( sorted( (state for state in attributes.values() if isinstance(state, state_class)), - key=lambda state: state.resource_id, + key=lambda state: state.state_name, ) ) -def main_linear_state(resources: AttResourceManager): +def main_linear_state(resources: AttentionStateManager): (state,) = resources.get_bounded_attention_states("main") return state @@ -104,7 +104,7 @@ def prepare_model_states(model: TpPartBaseModel, logical_layer_ids: Sequence[int model.layers_infer = [SimpleNamespace(layer_num_=i) for i in logical_layer_ids] if not hasattr(model, "torch_memory_saver"): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) - model.prepare_att_resources() + model.prepare_attention_states() def create_full_attention_roles( @@ -196,15 +196,15 @@ def build_hybrid_resources( small_page_capacity: int, draft_full_state: FullAttentionState | None = None, tp_rank: int = 0, -) -> AttResourceManager: +) -> AttentionStateManager: """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.att_resources.allocators import StatePageAllocator - from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state - from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.allocators import StatePageAllocator + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state + from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state from lightllm.models.glm5_next.config import Glm5NextResourceConfig - manager = AttResourceManager(request_capacity=request_capacity) + manager = AttentionStateManager(request_capacity=request_capacity) glm = isinstance(config, Glm5NextResourceConfig) create_state = create_kda_state if glm else create_gdn_state bounded_state = create_state( @@ -229,11 +229,11 @@ def build_hybrid_resources( bounded_state.initialize_page_pools( big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=True, ) - manager.set_att_state("main", bounded_state) + manager.register_state("main", bounded_state) for role, state in (("main", full_state), ("draft0", draft_full_state)): if state is None: continue - manager.set_att_state(role, state) + manager.register_state(role, state) if glm: tail = create_kpool_state( config=KPoolStateConfig( @@ -248,12 +248,12 @@ def build_hybrid_resources( mtp_step=mtp_step, device="cuda", ) - manager.set_att_state(role, tail) + manager.register_state(role, tail) manager.full_att_allocator = full_allocator manager.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) manager.bounded_att_small_page_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None manager.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 manager.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 manager.cpu_tail_restore = glm - manager.finish_state_creation() + manager.finish_state_registration() return manager diff --git a/unit_tests/common/att_resources/test_allocators.py b/unit_tests/common/att_resources/test_allocators.py index da3c705509..3b932d8de5 100644 --- a/unit_tests/common/att_resources/test_allocators.py +++ b/unit_tests/common/att_resources/test_allocators.py @@ -3,7 +3,7 @@ import pytest import torch -from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator def test_slot_allocation_domains_and_reserved_storage(): diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index 0f88538c65..d3d0705d1a 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -127,16 +127,16 @@ def test_backend_allocates_all_roles_before_warmup( mapping = SimpleNamespace(bind_hold_page=lambda indices: None) request_allocations = [] manager_creations = [] - from lightllm.common.att_resources import AttResourceManager + from lightllm.common.attention_state import AttentionStateManager def create_empty_manager(request_capacity: int): - manager = AttResourceManager(request_capacity=request_capacity) - assert not manager._states_created + manager = AttentionStateManager(request_capacity=request_capacity) + assert not manager._state_registration_finished assert not any(name.endswith("_state") for name in manager.__dict__) manager_creations.append(request_capacity) return manager - monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttResourceManager", create_empty_manager) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttentionStateManager", create_empty_manager) def create_req_manager(max_request_num: int, max_sequence_length: int): assert len(models) == 1 + draft_count @@ -149,13 +149,13 @@ def create_req_manager(max_request_num: int, max_sequence_length: int): def mark_backends(model: TpPartBaseModel): # 第一个模型初始化 backend 前，全部 State 已创建完成，可通过缓存查询取得资源。 for item in models: - for state in item.att_resource_manager.get_full_attention_states(item.att_resource_role): + for state in item.att_state_manager.get_full_attention_states(item.att_state_role): logical_layer_ids = {layer_id for group in state.config.logical_layer_groups for layer_id in group} for layer in item.layers_infer: if layer.layer_num_ in logical_layer_ids: - assert item.att_resource_manager.get_att_state( - model_name=item.att_resource_role, - att_type="", + assert item.att_state_manager.get_state( + model_role=item.att_state_role, + state_type_filter="", logical_layer_id=layer.layer_num_, state_class=type(state), ) is state @@ -166,8 +166,8 @@ def mark_backends(model: TpPartBaseModel): def warmup(model: TpPartBaseModel): assert model.backends_bound assert all(item.req_manager is mapping for item in models) - assert all(item.att_resource_manager is models[0].att_resource_manager for item in models) - assert model.att_resource_manager.full_att_allocator.size == 32 + assert all(item.att_state_manager is models[0].att_state_manager for item in models) + assert model.att_state_manager.full_att_allocator.size == 32 assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) assert all(item.max_total_token_num is None for item in models[models.index(model) + 1 :]) assert captured == prefill_captured == {id(item) for item in models[: models.index(model)]} @@ -211,7 +211,7 @@ def check_max_len(model: TpPartBaseModel): def memory_info(device: str | torch.device | int): assert all(item.req_manager is mapping for item in models) - manager = models[0].att_resource_manager + manager = models[0].att_state_manager states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.conv_state.is_cuda @@ -222,7 +222,7 @@ def memory_info(device: str | torch.device | int): monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", ready.append) - assert all(not hasattr(item, "att_resource_manager") for item in models) + assert all(not hasattr(item, "att_state_manager") for item in models) class RuntimeReady(Exception): """Stop after model readiness, before unrelated queue and thread setup.""" @@ -255,7 +255,7 @@ def start_runtime(): ) ) else: - assert all(getattr(item, "att_resource_manager", None) is None for item in models), "resources already initialized" + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" req_manager = create_req_manager( max_request_num=models[0].max_req_num, max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), @@ -263,13 +263,13 @@ def start_runtime(): for item in models: item.bind_req_manager(req_manager=req_manager) for index, item in enumerate(models): - item.prepare_att_resources() - assert all(other.att_resource_manager is models[0].att_resource_manager for other in models[: index + 1]) - assert all(not hasattr(other, "att_resource_manager") for other in models[index + 1 :]) + item.prepare_attention_states() + assert all(other.att_state_manager is models[0].att_state_manager for other in models[: index + 1]) + assert all(not hasattr(other, "att_state_manager") for other in models[index + 1 :]) main = models[0] - main.att_resource_manager.finish_state_creation() - main.att_resource_manager.load_full_att_calibration(args=args) - main.att_resource_manager.initialize_resources( + main.att_state_manager.finish_state_registration() + main.att_state_manager.load_full_att_calibration(args=args) + main.att_state_manager.initialize_storage( args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction ) assert all(item.max_total_token_num is None for item in models) @@ -286,7 +286,7 @@ def start_runtime(): assert captured == {id(model) for model in models} assert checked == [models[0]] assert ready == ([True] if via_backend else []) - manager = models[0].att_resource_manager + manager = models[0].att_state_manager assert len(manager.get_bounded_attention_states()) == draft_count + int(main_linear) assert all(not hasattr(model, "mtp_state_cursors") for model in models) assert all(not hasattr(model, "full_att_state") for model in models) @@ -331,9 +331,9 @@ def start_runtime(): layouts = build_cpu_cache_layouts() assert (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 == manager.get_cpu_page_bytes() - assert layouts.keys() == manager.get_cpu_transfer_layouts().keys() - for resource_id, layout in layouts.items(): - runtime_layout = manager.get_cpu_transfer_layouts()[resource_id] + assert layouts.keys() == manager.get_cpu_page_layouts().keys() + for state_name, layout in layouts.items(): + runtime_layout = manager.get_cpu_page_layouts()[state_name] assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in layout.regions] == [ (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in runtime_layout.regions ] @@ -341,25 +341,25 @@ def start_runtime(): for i, state in enumerate(manager.get_bounded_attention_states()): state.conv_state.fill_(i + 3) state.ssm_state.fill_(i + 7) - manager.save_bounded_att_pages( + manager.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[4], page_kind="big", - slot_ids=[slot], + page_slot_ids=[slot], ) manager.reset_bounded_att_state([0]) assert all(not state.conv_state[:, 0].any() for state in manager.get_bounded_attention_states()) - manager.restore_bounded_att_pages( + manager.restore_bounded_att_runtime_from_pages( req_indices=[0], state_token_lens=[4], page_kind="big", - slot_ids=[slot], + page_slot_ids=[slot], ) for i, state in enumerate(manager.get_bounded_attention_states()): assert (state.conv_state[:, 0, :, : state.config.conv_width] == i + 3).all() assert (state.ssm_state[:, 0] == i + 7).all() with pytest.raises(AssertionError, match="resources already initialized"): start_runtime() - assert all(item.att_resource_manager is manager for item in models) + assert all(item.att_state_manager is manager for item in models) get_env_start_args.cache_clear() get_added_mtp_kv_layer_num.cache_clear() diff --git a/unit_tests/common/att_resources/test_bounded_copy_kernel.py b/unit_tests/common/att_resources/test_bounded_copy_kernel.py index d64b42ce15..0984df7a3d 100644 --- a/unit_tests/common/att_resources/test_bounded_copy_kernel.py +++ b/unit_tests/common/att_resources/test_bounded_copy_kernel.py @@ -4,8 +4,8 @@ import torch import triton.language as tl -from lightllm.common.att_resources.bounded_att.copy_kernel import copy_bounded_att_tensors -from lightllm.common.att_resources.bounded_att.utils import register_page_pool_tensor +from lightllm.common.attention_state.bounded_att.copy_kernel import copy_bounded_att_tensors +from lightllm.common.attention_state.bounded_att.utils import register_page_pool_tensor pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -88,7 +88,7 @@ def test_exact_float_bit_patterns(float_dtype: torch.dtype, int_dtype: torch.dty def test_one_kernel_for_same_dimension_batch(monkeypatch: pytest.MonkeyPatch) -> None: - from lightllm.common.att_resources.bounded_att import copy_kernel + from lightllm.common.attention_state.bounded_att import copy_kernel actual_kernel = copy_kernel._copy_bounded_att_1d launches = [] @@ -136,7 +136,7 @@ def test_same_shape_batch_uses_one_kernel( num_stages: int, monkeypatch: pytest.MonkeyPatch, ) -> None: - from lightllm.common.att_resources.bounded_att import copy_kernel + from lightllm.common.attention_state.bounded_att import copy_kernel launches = [] kernel_name = f"_copy_bounded_att_{ndim}d" @@ -172,7 +172,7 @@ def __getitem__(self, grid: tuple[int]): @pytest.mark.parametrize("layout", ["ssm", "conv_mtp", "dense", "width4", "middle_unaligned", "prefix", "singletons", "wide_singleton", "wide_singleton_prefix", "all_singletons"]) def test_transfer_views_preserve_gaps_and_alignment(layout: str, monkeypatch: pytest.MonkeyPatch) -> None: - from lightllm.common.att_resources.bounded_att import copy_kernel + from lightllm.common.attention_state.bounded_att import copy_kernel launches = [] diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 7a9ffe5aad..9322b32867 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -6,20 +6,20 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.bounded_att import ( +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ( GDNPagesPool, GDNState, KDAPagesPool, KDAState, ConvSSMStateConfig, ) -from lightllm.common.att_resources.types import ( - AxisSlice, StateTransferLayout, TransferRegion, +from lightllm.common.attention_state.types import ( + AxisMapping, StateTransferLayout, TransferRegion, ) if TYPE_CHECKING: - from lightllm.common.att_resources.types import PageKind + from lightllm.common.attention_state.types import PageKind pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for aggregated copies") @@ -46,8 +46,8 @@ def make_state( mtp_step=mtp_step, device=device, ) - manager = AttResourceManager(request_capacity=3) - manager.set_att_state("main", state) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", state) state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") return state @@ -60,7 +60,7 @@ def page_layout(state: KDAState | GDNState): ssm_bytes = len(cfg.logical_layer_ids) * cfg.num_v_heads * cfg.key_dim * cfg.value_dim * 4 end = 16 + conv_bytes + ssm_bytes return StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=None, byte_start=0, byte_end=end, @@ -72,17 +72,17 @@ def page_layout(state: KDAState | GDNState): dtype=torch.float32, axes=("channel", "layer", "window"), shape=conv_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="channel", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=cfg.conv_channels, ), - AxisSlice( + AxisMapping( axis="window", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=cfg.conv_width, ), ), @@ -94,23 +94,23 @@ def page_layout(state: KDAState | GDNState): dtype=torch.float32, axes=("layer", "head", "key_dim", "value_dim"), shape=ssm_shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="head", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=cfg.num_v_heads, ), - AxisSlice( + AxisMapping( axis="key_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=cfg.key_dim, ), - AxisSlice( + AxisMapping( axis="value_dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=cfg.value_dim, ), ), @@ -146,7 +146,7 @@ def test_prefill_state_is_saved_and_restored( mtp_step=2, device=device, ) - assert state.resource_id is None and not hasattr(config, "resource_id") + assert state.state_name is None and not hasattr(config, "state_name") assert state.big_pages_pool is None and state.small_pages_pool is None conv, ssm = state.get_layer_view(3) assert conv.data_ptr() == state.conv_state[1].data_ptr() @@ -164,11 +164,11 @@ def test_prefill_state_is_saved_and_restored( assert type(state.small_pages_pool) is pool_cls assert state.big_pages_pool is not state.small_pages_pool assert state.mtp_state_index.tolist() == [0, 0, 1] - state.save_bounded_att_pages( + state.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind=page_kind, - slot_ids=[1], + page_slot_ids=[1], ) if device == "cuda": torch.cuda.synchronize() @@ -177,11 +177,11 @@ def test_prefill_state_is_saved_and_restored( assert torch.equal(pool.conv_state[1, ..., 1], torch.full_like(pool.conv_state[1, ..., 1], 11)) assert (pool.ssm_state[1] == 20).all() state.mtp_state_index[1] = 2 - state.restore_bounded_att_pages( + state.restore_bounded_att_runtime_from_pages( req_indices=[1], state_token_lens=[8], page_kind=page_kind, - slot_ids=[1], + page_slot_ids=[1], ) assert state.mtp_state_index[1].item() == 0 assert torch.equal(state.conv_state[:, 1, :, :2].cpu(), pool.conv_state[1]) @@ -210,19 +210,19 @@ def test_gdn_dimensions_and_page_restore(): state.conv_state[:, 0, :, :2].copy_(conv) state.ssm_state[:, 0].copy_(ssm) for page_kind in ("big", "small"): - state.save_bounded_att_pages( + state.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind=page_kind, - slot_ids=[0], + page_slot_ids=[0], ) state.reset_bounded_att_state([0]) for req, page_kind in enumerate(("big", "small"), start=1): - state.restore_bounded_att_pages( + state.restore_bounded_att_runtime_from_pages( req_indices=[req], state_token_lens=[8], page_kind=page_kind, - slot_ids=[0], + page_slot_ids=[0], ) assert torch.equal(state.conv_state[:, req, :, :2].cpu(), conv) assert torch.equal(state.ssm_state[:, req * 3].cpu(), ssm) @@ -235,11 +235,11 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): state = make_state(cls, device="cuda") state.conv_state[:, 0].fill_(7) state.ssm_state[:, 0:3].fill_(11) - state.save_bounded_att_pages( + state.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[6], page_kind="small", - slot_ids=[0], + page_slot_ids=[0], ) # 推进候选区和游标后，重置应清零全部运行态，同时保留已保存的页快照。 state.conv_state[:, 0, :, 2:4].fill_(19) @@ -259,17 +259,17 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): source = make_state(cls, device="cuda") source.conv_state[:, 0].fill_(13) source.ssm_state[:, 0].fill_(29) - source.save_bounded_att_pages( + source.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=[0], + page_slot_ids=[0], ) layout = page_layout(source) buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) - source.write_page_state_to_cpu( + source.write_state_pages_to_cpu( page_kind="big", - slot_ids=[0], + page_slot_ids=[0], cpu_page_buffers=buffer, cpu_page_indices=[2], layout=layout, @@ -277,30 +277,30 @@ def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): torch.cuda.synchronize() assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() target = make_state(cls, device="cuda") - target.read_page_state_from_cpu( + target.read_state_pages_from_cpu( cpu_page_buffers=buffer, cpu_page_indices=[2], big_page_slot_ids=[1], layout=layout, ) - target.restore_bounded_att_pages( + target.restore_bounded_att_runtime_from_pages( req_indices=[2], state_token_lens=[8], page_kind="big", - slot_ids=[1], + page_slot_ids=[1], ) assert (target.conv_state[:, 2, :, :2] == 13).all() assert (target.ssm_state[:, 6] == 29).all() # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 - target.read_page_state_from_cpu( + target.read_state_pages_from_cpu( cpu_page_buffers=buffer, cpu_page_indices=[2], big_page_slot_ids=[0], layout=layout, ) - target.write_page_state_to_cpu( + target.write_state_pages_to_cpu( page_kind="big", - slot_ids=[0], + page_slot_ids=[0], cpu_page_buffers=buffer, cpu_page_indices=[0], layout=layout, @@ -331,7 +331,7 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( device=device, ) page = allocation[1] - source.write_to_pd_buffer( + source.write_to_pd_page( req_idx=0, pd_page_buffers=allocation, pd_page_index=1, @@ -339,7 +339,7 @@ def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( layout=layout, ) target = make_state(cls, device=device, mtp_step=1) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=allocation, req_idx=2, pd_page_index=1, @@ -374,7 +374,7 @@ def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp src_conv = source.conv_state.clone() src_ssm = source.ssm_state.clone() - target.copy_run_time_state_from_other_dp( + target.copy_runtime_state_from_dp( src_req_indices=[2, 0], dst_req_indices=[0, 1], remote_dp_bounded_state=source, ) @@ -411,7 +411,7 @@ def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): src_conv = source.big_pages_pool.conv_state.clone() src_ssm = source.big_pages_pool.ssm_state.clone() - target.copy_big_page_state_from_other_dp( + target.copy_big_pages_from_dp( src_page_ids=[2, 0], dst_page_ids=[0, 1], remote_dp_bounded_state=source, ) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index a8d42164fd..c33c9d1eba 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -18,12 +18,12 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNState, KDAState -from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState +from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( FullAttentionState, GQAFullAttentionState, FullAttentionConfig, @@ -111,9 +111,9 @@ def test_budget_matches_token_stores_including_hold_and_padding( device="cpu", ) actual = allocated_bytes(state, token_only=True) - planning = AttResourceManager(request_capacity=1) - planning.set_att_state("main", rebuild_full_state(template, config=template.config)) - planning.finish_state_creation() + planning = AttentionStateManager(request_capacity=1) + planning.register_state("main", rebuild_full_state(template, config=template.config)) + planning.finish_state_registration() token_bytes = planning.get_full_att_token_bytes() pd_page_bytes = planning.get_pd_page_bytes() assert not hasattr(planning, "pd_page_bytes") @@ -125,10 +125,10 @@ def test_budget_matches_token_stores_including_hold_and_padding( minimum = token_bytes * 4 assert minimum > 0 - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", state) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) manager.full_att_allocator = FullAttKVAllocator(20, 4, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() (state,) = manager.get_full_attention_states("main") assert manager.get_pd_page_bytes() == pd_page_bytes args = StartArgs( @@ -140,18 +140,18 @@ def test_budget_matches_token_stores_including_hold_and_padding( reserve = manager.get_pd_page_bytes() * args.pd_kv_page_num monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.get_current_device_id", lambda: 0) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 0) monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) for budget, expected in ((actual, 20), (actual - 1, 16), (minimum, 0), (minimum - 1, None)): - resources = AttResourceManager(request_capacity=1) - resources.set_att_state("main", rebuild_full_state(template, config=template.config)) + resources = AttentionStateManager(request_capacity=1) + resources.register_state("main", rebuild_full_state(template, config=template.config)) monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) - resources.finish_state_creation() + resources.finish_state_registration() if expected is None: with pytest.raises(MemoryError, match="HOLD"): - resources.initialize_resources( + resources.initialize_storage( args=args, token_capacity=None, mem_fraction=1.0, @@ -159,14 +159,14 @@ def test_budget_matches_token_stores_including_hold_and_padding( ) elif expected == 0: with pytest.raises(AssertionError, match="at least one complete page"): - resources.initialize_resources( + resources.initialize_storage( args=args, token_capacity=None, mem_fraction=1.0, device="cpu", ) else: - resources.initialize_resources( + resources.initialize_storage( args=args, token_capacity=None, mem_fraction=1.0, @@ -178,10 +178,10 @@ def test_budget_matches_token_stores_including_hold_and_padding( def test_shared_allocator_accounts_for_distinct_state_stores(): main, draft = definition("int4"), definition("static_head") - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", main) - manager.set_att_state("draft0", draft) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", main) + manager.register_state("draft0", draft) + manager.finish_state_registration() expected = sum( allocated_bytes( rebuild_full_state( @@ -236,19 +236,19 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( ssm = reference.ssm_state[:, 0] bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() tail_bytes = len(tail.config.logical_layer_ids) * tail.config.pool_size * 2 * tail.config.head_dim * tail.config.dtype.itemsize - manager = AttResourceManager(request_capacity=3) - manager.set_att_state("main", reference) - manager.set_att_state("main", tail) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", reference) + manager.register_state("main", tail) + manager.finish_state_registration() # PD 页只包含运行态数据，使用第二页验证 8 字节页跨度下的读写。 assert manager.get_pd_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 reference.conv_state.fill_(2) reference.ssm_state.fill_(3) tail.tail_buffer.fill_(4) pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() + layouts = manager.get_pd_runtime_page_layouts() assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "tail"} - manager.write_to_pd_buffer( + manager.write_to_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -259,7 +259,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( reference.conv_state.zero_() reference.ssm_state.zero_() tail.tail_buffer.zero_() - manager.read_from_pd_buffer( + manager.read_from_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -287,7 +287,7 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( pd_args: EasyDict, ): pd_args.pd_kv_page_size = 3 - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) for state_cls, role, dim, dtype in ( (MLAFullAttentionState, "draft0", 1, torch.bfloat16), (GQAFullAttentionState, "main", 3, torch.bfloat16), @@ -299,14 +299,14 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( head_dim=dim, dtype=dtype, ) - manager.set_att_state(role, state_cls(config, page_size=1, device="meta")) - manager.finish_state_creation() + manager.register_state(role, state_cls(config, page_size=1, device="meta")) + manager.finish_state_registration() # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 assert manager.get_pd_page_bytes() == 72 monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage( + manager.allocate_storage( token_capacity=4, page_size=1, big_page_tokens=4, @@ -319,9 +319,9 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( storage_buffer.fill_(i) pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") page = pages[1] - layouts = manager.get_pd_full_att_transfer_layouts() - assert [min(region.byte_start for region in layouts[state.resource_id].regions) for state in states] == [0, 8, 48] - manager.write_to_pd_buffer( + layouts = manager.get_pd_token_page_layouts() + assert [min(region.byte_start for region in layouts[state.state_name].regions) for state in states] == [0, 8, 48] + manager.write_to_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind="kv", @@ -332,7 +332,7 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( for state in states: for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() - manager.read_from_pd_buffer( + manager.read_from_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind="kv", @@ -372,20 +372,20 @@ def _get_transfer_regions(self, page_tokens: int): bits=8, group_size=128, ) - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") - prepared_full_state.initialize_storage(8) - manager.set_att_state("main", prepared_full_state) + prepared_full_state.allocate_storage(8) + manager.register_state("main", prepared_full_state) prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") - prepared_full_state.initialize_storage(8) - manager.set_att_state("draft0", prepared_full_state) + prepared_full_state.allocate_storage(8) + manager.register_state("draft0", prepared_full_state) manager.full_att_allocator = FullAttKVAllocator(8, 1, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() # 每个 State 的裸数据为 792 字节；scale 间 padding 使实际末尾为 796，页段需 800 字节。 for state in manager.get_full_attention_states(): assert state.get_token_bytes() * 3 == 792 monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("capacity must use field layouts")) - get_layouts = manager.get_cpu_transfer_layouts if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts + get_layouts = manager.get_cpu_page_layouts if transfer == "cpu" else manager.get_pd_token_page_layouts get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes if layout_first: with monkeypatch.context() as patch: @@ -402,7 +402,7 @@ def _get_transfer_regions(self, page_tokens: int): pin_memory=transfer == "cpu", ) for index, state in enumerate(manager.get_full_attention_states()): - layout = layouts[state.resource_id] + layout = layouts[state.state_name] assert (layout.byte_start, layout.byte_end) == (index * 800, (index + 1) * 800) regions = layout.regions assert [region.byte_start for region in regions] == [index * 800 + n for n in (0, 384, 768, 784)] @@ -423,14 +423,14 @@ def _get_transfer_regions(self, page_tokens: int): layouts=layouts, ) else: - manager.write_to_pd_buffer( + manager.write_to_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), ) - manager.read_from_pd_buffer( + manager.read_from_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", @@ -460,7 +460,7 @@ def get_pd_transfer_regions(self): for region in super().get_pd_transfer_regions() ) - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) if full_page_tokens is not None: pd_args.pd_kv_page_size = full_page_tokens config = FullAttentionConfig( @@ -469,7 +469,7 @@ def get_pd_transfer_regions(self): head_dim=1, dtype=torch.float32, ) - manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + manager.register_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) state = PaddedGDNState( ConvSSMStateConfig( logical_layer_ids=(0,), @@ -483,27 +483,27 @@ def get_pd_transfer_regions(self): ), request_capacity=1, device="cuda", ) - manager.set_att_state("main", state) - manager.finish_state_creation() + manager.register_state("main", state) + manager.finish_state_registration() # Conv 为 6 字节，SSM 从偏移 72 开始，占 4 字节；完整运行态需要 80 字节。 # 即使 KV 页不足，布局仍可独立查询；最终块大小查询选择两类占用的最大值。 with monkeypatch.context() as patch: - patch.setattr(manager, "get_pd_full_att_transfer_layouts", lambda: pytest.fail("runtime layout must be independent")) + patch.setattr(manager, "get_pd_token_page_layouts", lambda: pytest.fail("runtime layout must be independent")) patch.setattr(manager, "get_pd_page_bytes", lambda: pytest.fail("layout must not query page bytes")) - layouts = manager.get_pd_bounded_att_or_tail_transfer_layouts() - layout = layouts[state.resource_id] + layouts = manager.get_pd_runtime_page_layouts() + layout = layouts[state.state_name] assert layout.byte_end == 80 and layout.tokens_per_page is None assert layout.regions[-1].byte_start == 72 warnings = [] - monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.logger.warning", warnings.append) + monkeypatch.setattr("lightllm.common.attention_state.manager.logger.warning", warnings.append) expected = max(80, (full_page_tokens or 0) * 8) assert manager.get_pd_page_bytes() == expected - assert manager.get_pd_bounded_att_or_tail_transfer_layouts() is layouts + assert manager.get_pd_runtime_page_layouts() is layouts # 实际传输页可大于运行态布局，选中第二页后仍从该页起点搬运。 pages = torch.full((3, expected), 91, dtype=torch.uint8, device="cuda") state.conv_state.fill_(3) state.ssm_state.fill_(5) - manager.write_to_pd_buffer( + manager.write_to_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -513,7 +513,7 @@ def get_pd_transfer_regions(self): ) state.conv_state.zero_() state.ssm_state.zero_() - manager.read_from_pd_buffer( + manager.read_from_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -549,28 +549,28 @@ def _get_transfer_regions(self, page_tokens: int): ) pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) config = FullAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=3, dtype=torch.bfloat16, ) - manager.set_att_state("main", PaddedState(config, page_size=1, device="meta")) - manager.finish_state_creation() + manager.register_state("main", PaddedState(config, page_size=1, device="meta")) + manager.finish_state_registration() get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes assert get_page_bytes() == 48 monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage( + manager.allocate_storage( token_capacity=8, page_size=4, big_page_tokens=8, device="cpu", ) assert get_page_bytes() == (80 if transfer == "cpu" else 72) - layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() + layouts = manager.get_cpu_page_layouts() if transfer == "cpu" else manager.get_pd_token_page_layouts() assert layouts["main_gqa_bf16_state"].byte_end == 72 @@ -581,7 +581,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( extra_tokens: int, pd_args: EasyDict, ): - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) config = FullAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, @@ -589,7 +589,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( dtype=torch.float32, ) full_state = GQAFullAttentionState(config, page_size=1, device="meta") - manager.set_att_state("main", full_state) + manager.register_state("main", full_state) for cls, role in ((KDAState, "main"), (GDNState, "draft0")): config = ConvSSMStateConfig( logical_layer_ids=(0,), @@ -602,7 +602,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( ssm_dtype=torch.float32, ) state = cls(config, request_capacity=1, device="meta") - manager.set_att_state(role, state) + manager.register_state(role, state) if with_tail: state = KPoolAttentionTailState( config=KPoolStateConfig( @@ -615,12 +615,12 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( mtp_step=0, device="meta", ) - manager.set_att_state("main", state) + manager.register_state("main", state) # 两份运行态各占 16 字节数据，tail 另占 32 字节，无页头及长度、版本字段。 runtime_page_bytes = 32 + 32 * with_tail page_tokens = runtime_page_bytes // 8 + extra_tokens pd_args.pd_kv_page_size = page_tokens - manager.finish_state_creation() + manager.finish_state_registration() assert manager.get_pd_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) assert manager.full_att_allocator is None assert read_kv_storage(manager.main_gqa_fp32_state) is None @@ -644,9 +644,9 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) scales = torch.ones((len(config.logical_layer_groups), 2 * config.head_num)) if kind == "static_head" else None q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None - manager = AttResourceManager(request_capacity=3) + manager = AttentionStateManager(request_capacity=3) for role in ("main", "draft0"): - manager.set_att_state(role, rebuild_full_state( + manager.register_state(role, rebuild_full_state( template=template, config=config, scales=scales, @@ -672,14 +672,14 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( mtp_step=2, device="meta", ) - manager.set_att_state("main", bounded) - manager.finish_state_creation() + manager.register_state("main", bounded) + manager.finish_state_registration() # CPU 容量和布局都只依赖 State 的 CPU 字段描述，不读取 GPU 存储开销或 PD 布局。 for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): if isinstance(state, FullAttentionState): monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("CPU capacity must use field layouts")) monkeypatch.setattr(state, "get_pd_transfer_regions", lambda: pytest.fail("CPU layout must be independent")) - layouts = manager.get_cpu_transfer_layouts() + layouts = manager.get_cpu_page_layouts() assert min(region.byte_start for layout in layouts.values() for region in layout.regions) == 0 assert not {"state_token_len", "format_version"} & { region.field for layout in layouts.values() for region in layout.regions @@ -689,7 +689,7 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): monkeypatch.setattr(state, "get_cpu_transfer_regions", lambda: pytest.fail("cached layout must be reused")) assert manager.get_cpu_page_bytes() == expected_page_bytes - assert manager.get_cpu_transfer_layouts() is layouts + assert manager.get_cpu_page_layouts() is layouts def test_cpu_bounded_capacity_includes_state_field_padding(cpu_args: EasyDict): @@ -701,7 +701,7 @@ def get_cpu_transfer_regions(self): for region in super().get_cpu_transfer_regions() ) - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) state = PaddedGDNState( ConvSSMStateConfig( logical_layer_ids=(0,), @@ -715,11 +715,11 @@ def get_cpu_transfer_regions(self): ), request_capacity=1, device="meta", ) - manager.set_att_state("main", state) - manager.finish_state_creation() + manager.register_state("main", state) + manager.finish_state_registration() # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 assert manager.get_cpu_page_bytes() == 80 - layout = manager.get_cpu_transfer_layouts()[state.resource_id] + layout = manager.get_cpu_page_layouts()[state.state_name] assert (layout.byte_start, layout.byte_end) == (0, 80) assert layout.regions[-1].byte_end == 76 @@ -741,14 +741,14 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.view(torch.uint8).fill_(211) states.append(state) - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", states[0]) - manager.set_att_state("draft0", states[1]) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", states[0]) + manager.register_state("draft0", states[1]) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - manager.finish_state_creation() - layouts = manager.get_cpu_transfer_layouts() + manager.finish_state_registration() + layouts = manager.get_cpu_page_layouts() packed, gqa = manager.get_full_attention_states() - packed_layout, gqa_layout = layouts[packed.resource_id], layouts[gqa.resource_id] + packed_layout, gqa_layout = layouts[packed.state_name], layouts[gqa.state_name] data_end = max(r.byte_end for r in packed_layout.regions) gqa_start = gqa_layout.byte_start # 后一个 State 紧跟实际字段终点，仅保留对齐空间，不为 GPU packed 行 padding 预留位置。 @@ -788,7 +788,7 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: EasyDict, page_tokens: int): cpu_args.cpu_cache_token_page_size = page_tokens - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) manager.bounded_att_big_page_allocator = StatePageAllocator(2) manager.bounded_att_small_page_allocator = StatePageAllocator(1) for state_cls, role, value_dim in ((KDAState, "draft0", 1), (GDNState, "main", 3)): @@ -809,23 +809,23 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas device="cuda", ) state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=True) - manager.set_att_state(role, state) - manager.finish_state_creation() + manager.register_state(role, state) + manager.finish_state_registration() states = manager.get_bounded_attention_states() - layouts = manager.get_cpu_transfer_layouts() + layouts = manager.get_cpu_page_layouts() # 两份快照分别含 10/22 字节数据，预留 16/24 字节；页 token 数不放大快照。 - assert [min(r.byte_start for r in layouts[state.resource_id].regions) for state in states] == [0, 16] + assert [min(r.byte_start for r in layouts[state.state_name].regions) for state in states] == [0, 16] assert manager.get_cpu_page_bytes() == 48 assert [(layout.byte_start, layout.byte_end) for layout in layouts.values()] == [(0, 16), (16, 40)] for i, state in enumerate(states): state.conv_state.fill_(i + 3) state.ssm_state.fill_(i + 7) source_slots = manager.bounded_att_big_page_allocator.alloc(1) - manager.save_bounded_att_pages( + manager.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=source_slots, + page_slot_ids=source_slots, ) pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) manager.write_req_to_cpu( @@ -854,7 +854,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas assert (state.ssm_state[:, 0] == i + 7).all() -@pytest.mark.parametrize("transfer", ["cpu", "pd_full_att", "pd_bounded_att_or_tail"]) +@pytest.mark.parametrize("transfer", ["cpu", "pd_token", "pd_runtime"]) def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation( monkeypatch: pytest.MonkeyPatch, transfer: str, @@ -862,15 +862,15 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation ): pd_args.pd_kv_page_size = 3 pd_args.cpu_cache_token_page_size = 3 - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) config = FullAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=3, dtype=torch.bfloat16, ) - manager.set_att_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) - manager.set_att_state("main", GDNState( + manager.register_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + manager.register_state("main", GDNState( ConvSSMStateConfig( logical_layer_ids=(0,), global_k_heads=1, @@ -883,29 +883,29 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation ), request_capacity=1, device="cpu", )) - manager.finish_state_creation() - getter = getattr(manager, f"get_{transfer}_transfer_layouts") + manager.finish_state_registration() + getter = getattr(manager, f"get_{transfer}_page_layouts") before = getter() assert getter() is before - state = manager.main_gdn_conv_bf16_ssm_bf16_state if transfer == "pd_bounded_att_or_tail" else manager.main_gqa_bf16_state - before_layout = before[state.resource_id] + state = manager.main_gdn_conv_bf16_ssm_bf16_state if transfer == "pd_runtime" else manager.main_gqa_bf16_state + before_layout = before[state.state_name] monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.initialize_storage( + manager.allocate_storage( token_capacity=8, page_size=4, big_page_tokens=8, device="cpu", ) after = getter() - after_layout = after[state.resource_id] + after_layout = after[state.state_name] assert after is not before and getter() is after assert after_layout is not before_layout full_state = manager.main_gqa_bf16_state assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() - if transfer == "pd_bounded_att_or_tail": + if transfer == "pd_runtime": assert after_layout.byte_end == 8 and page_bytes == 40 elif transfer == "cpu": layout_end = max(layout.byte_end for layout in after.values()) @@ -936,16 +936,16 @@ def get_cpu_transfer_regions(self): head_dim=2, dtype=torch.float32, ) - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") - prepared_full_state.initialize_storage(8) - manager.set_att_state("main", prepared_full_state) + prepared_full_state.allocate_storage(8) + manager.register_state("main", prepared_full_state) prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") - prepared_full_state.initialize_storage(8) - manager.set_att_state("draft0", prepared_full_state) + prepared_full_state.allocate_storage(8) + manager.register_state("draft0", prepared_full_state) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - manager.finish_state_creation() - layouts = manager.get_cpu_transfer_layouts() if transfer == "cpu" else manager.get_pd_full_att_transfer_layouts() + manager.finish_state_registration() + layouts = manager.get_cpu_page_layouts() if transfer == "cpu" else manager.get_pd_token_page_layouts() for index, state in enumerate(manager.get_full_attention_states()): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() @@ -954,7 +954,7 @@ def get_cpu_transfer_regions(self): # 1 层、4 tokens，K/V 各 1 head、每 head 2 个 FP32 元素。 state_bytes = 4 * 2 * 2 * torch.float32.itemsize k_offset = state_bytes // 2 if transfer == "cpu" else 0 - k_region, v_region = layouts[state.resource_id].regions + k_region, v_region = layouts[state.state_name].regions assert k_region.field == "k" and v_region.field == "v" assert k_region.byte_start == index * state_bytes + k_offset assert v_region.byte_start == index * state_bytes + state_bytes // 2 - k_offset @@ -985,14 +985,14 @@ def get_cpu_transfer_regions(self): layouts=layouts, ) else: - manager.write_to_pd_buffer( + manager.write_to_pd_page( pd_page_buffers=pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, dtype=torch.long, device="cuda"), ) - manager.read_from_pd_buffer( + manager.read_from_pd_page( pd_page_buffers=pages, pd_page_index=0, kind="kv", diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index f4fb55fed7..3a70b0fd58 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -12,8 +12,8 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( @@ -38,15 +38,15 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, ) for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() - manager = AttResourceManager(request_capacity=2) - manager.set_att_state(role, state) + manager = AttentionStateManager(request_capacity=2) + manager.register_state(role, state) manager.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() state = getattr(manager, f"{role}_{template.get_state_type()}_state") infer = InferStateInfo() - infer.att_resource_manager = manager + infer.att_state_manager = manager infer.is_draft_model = role == "draft0" - infer.att_resource_role = role + infer.att_state_role = role infer.mem_index = torch.tensor([1, 5], dtype=torch.int32, device="cuda") heads = 1 if kind == "packed_glm" else 2 * state.config.head_num dim = state.config.kv_dim if kind == "packed_glm" else state.config.head_dim diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 9f943a64df..6f28089874 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -14,9 +14,9 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.basemodel.attention.base_att import AttControl, BaseAttBackend from lightllm.common.basemodel.attention.fa3 import fp8 as fa3 from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -52,9 +52,9 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( scales=scales, q_scales=q_scales, )) - manager = AttResourceManager(request_capacity=2) - manager.set_att_state("main", states[0]) - manager.set_att_state("draft0", states[1]) + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", states[0]) + manager.register_state("draft0", states[1]) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) # Another encoding in the same role must not supply the FP8 calibration. plain = definition("gqa") @@ -64,14 +64,14 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( token_capacity=8, device="cuda", ) - manager.set_att_state(role, plain) - manager.finish_state_creation() + manager.register_state(role, plain) + manager.finish_state_registration() selected = getattr(manager, f"{role}_{template.get_state_type()}_state") k, v = selected.get_layer_view(8) infer = InferStateInfo() - infer.att_resource_manager = manager + infer.att_state_manager = manager infer.is_draft_model = role == "draft0" - infer.att_resource_role = role + infer.att_state_role = role infer.batch_size, infer.max_q_seq_len = 2, 1 infer.b_seq_len = torch.ones(2, dtype=torch.int32, device="cuda") backend = object.__new__(BaseAttBackend) @@ -111,6 +111,7 @@ def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.devic q=q, k=k, v=v, + att_control=AttControl(logical_layer_id=8), alloc_func=allocate, ) else: @@ -118,6 +119,7 @@ def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.devic q=q, k=k, v=v, + att_control=AttControl(logical_layer_id=8), alloc_func=allocate, ) if kind == "static_head": @@ -134,7 +136,7 @@ def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.devic @pytest.mark.parametrize("head_num", [1, 3]) @pytest.mark.parametrize("scale_mode", ["per_tensor", "per_head"]) def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_num: int, scale_mode: str): - from lightllm.common.att_resources.full_att import StaticFP8AttentionConfig, StaticFP8FullAttentionState + from lightllm.common.attention_state.full_att import StaticFP8AttentionConfig, StaticFP8FullAttentionState from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 scale_count = 2 if scale_mode == "per_tensor" else 2 * head_num @@ -158,7 +160,7 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu state=state, scales=scales, q_scales=torch.ones((1, head_num), dtype=torch.float32, device="cuda") if scale_mode == "per_head" else None, ) - state.initialize_storage(8) + state.allocate_storage(8) state.k_storage.zero_() state.v_storage.zero_() indices = torch.tensor([1, 6], dtype=torch.int32, device="cuda") diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index e28b973a78..8791e9c808 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -10,7 +10,7 @@ from unit_tests.common.att_resources.model_helpers import full_state_test_buffers from unit_tests.common.att_resources.test_state_identity import allocated_state from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout -from lightllm.common.att_resources.types import AxisSlice +from lightllm.common.attention_state.types import AxisMapping @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -107,15 +107,15 @@ def test_direct_cpu_pages_only_update_selected_head_slice_and_skip_replica_write regions = [] for region in layout.regions: slices = tuple( - AxisSlice( + AxisMapping( axis="head", - src_start=1, - dst_start=0, + local_start=1, + external_start=0, length=1, ) if item.axis == "head" else item - for item in region.axis_slices + for item in region.axis_mappings ) - regions.append(replace(region, axis_slices=slices)) + regions.append(replace(region, axis_mappings=slices)) layout = replace(layout, regions=tuple(regions)) pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) ids = torch.tensor([1, 5], dtype=torch.int64, device="cuda") @@ -275,7 +275,7 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_full_att_tokens(src, dst) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.copy_tokens_from_other_dp(src, dst, state) + state.copy_tokens_from_dp(src, dst, state) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.write_to_cpu( token_indices=cpu_ids, @@ -291,14 +291,14 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: layout=layout, ) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=cpu_ids, pd_page_buffers=pd_pages, pd_page_index=0, layout=layout, ) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pd_pages, token_indices=cpu_ids, pd_page_index=0, diff --git a/unit_tests/common/att_resources/test_full_att_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py index 29859af637..c07b8d8453 100644 --- a/unit_tests/common/att_resources/test_full_att_dp_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -10,8 +10,8 @@ import torch import torch.multiprocessing as mp -from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState -from lightllm.common.att_resources.full_att.dp_transfer import copy_gqa_tokens +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState +from lightllm.common.attention_state.full_att.dp_transfer import copy_gqa_tokens from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, rebuild_full_state from unit_tests.common.att_resources.test_capacity import definition @@ -83,7 +83,7 @@ def reject_intermediate(*args: object, **kwargs: object) -> None: for method in ("index_select", "index_copy_", "to", "tolist", "cpu", "contiguous"): patch.setattr(torch.Tensor, method, reject_intermediate) with torch.cuda.stream(stream): - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=src_ids, dst_token_indices=dst_ids, remote_dp_full_att_state=source, @@ -167,7 +167,7 @@ def _ipc_copy_worker(connection: Connection, payload: bytes) -> None: expected = {name: tensor.cpu().view(torch.uint8) for name, tensor in target_buffers.items()} for name, tensor in full_state_test_buffers(source).items(): expected[name][:, [8, 2, 5]] = tensor.cpu().view(torch.uint8)[:, [3, 1, 3]] - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=src_ids, dst_token_indices=torch.tensor([8, 2, 5], dtype=torch.int64, device="cuda:0"), remote_dp_full_att_state=source, diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py index 7cfeee833c..6567ec1211 100644 --- a/unit_tests/common/att_resources/test_full_att_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -10,8 +10,8 @@ import pytest import torch -from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState -from lightllm.common.att_resources.types import StateTransferLayout +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, full_state_test_fields, rebuild_full_state from unit_tests.common.att_resources.test_capacity import definition @@ -72,7 +72,7 @@ def test_full_att_pages_copy_all_fields_with_bounded_grid( shape=(region.shape[1], region.shape[0], *region.shape[2:]), ) for region in reversed(regions)) layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=page_tokens, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -127,13 +127,13 @@ def reject_gather(*args: object, **kwargs: object) -> None: layout=layout, ) else: - source.write_to_pd_buffer( + source.write_to_pd_page( token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, @@ -185,7 +185,7 @@ def test_all_full_att_schemes_transfer_through_peer_ipc( expected_fields[name].tensor.copy_(other=field.tensor[:, 1:2]) regions = state.get_pd_transfer_regions() layouts.append(StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -233,13 +233,13 @@ def _peer_transfer_worker( target_ids = torch.tensor(data=[5], dtype=torch.int64, device="cuda:0") for state, layout in zip(states, layouts): pages = torch.empty(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda:0") - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=source_ids, pd_page_buffers=pages, pd_page_index=1, layout=layout, ) - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, token_indices=target_ids, pd_page_index=1, @@ -276,7 +276,7 @@ def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( buffer.view(torch.uint8).random_(0, 256) regions = state.get_cpu_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=5, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -315,7 +315,7 @@ def reject_allocation(*args: object, **kwargs: object) -> None: cpu_page_buffers=pages[:, 16:-16], cpu_page_indices=page_indices, layout=layout, - cpu_page_readies=readies, + cpu_page_ready_flags=readies, ) torch.cuda.synchronize() # 整页比较同时覆盖 KV、量化 scale、packed indexer 字段及前后 padding。 diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 8b6e1ed742..787025b01e 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -12,14 +12,14 @@ import pytest import torch -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( FullAttentionState, GQAFullAttentionState, MLAFullAttentionState, FullAttentionConfig, ) -from lightllm.common.att_resources.types import ( - AxisSlice, StateTransferLayout, TransferRegion, +from lightllm.common.attention_state.types import ( + AxisMapping, StateTransferLayout, TransferRegion, ) @@ -30,7 +30,7 @@ def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: st head_dim=4, dtype=torch.float32, ), page_size=4, device=device) - prepared_full_state.initialize_storage(12) + prepared_full_state.allocate_storage(12) return prepared_full_state @@ -46,17 +46,17 @@ def layout_for(state: FullAttentionState): dtype=torch.float32, axes=("token", "layer", "head", "dim"), shape=(4, 2, 1, 4), - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="head", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=1, ), - AxisSlice( + AxisMapping( axis="dim", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=4, ), ), @@ -64,7 +64,7 @@ def layout_for(state: FullAttentionState): ) offset += 128 return StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=0, byte_end=offset + 16, @@ -119,21 +119,21 @@ def reject_host_list(tensor: torch.Tensor): ) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=ids, pd_page_buffers=pages, pd_page_index=1, layout=layout, ) with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, token_indices=ids, pd_page_index=1, layout=layout, ) return - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=ids, pd_page_buffers=pages, pd_page_index=1, @@ -142,7 +142,7 @@ def reject_host_list(tensor: torch.Tensor): target = full_state(cls, device) for storage_buffer in full_state_test_buffers(target).values(): storage_buffer.zero_() - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=pages, token_indices=torch.tensor([5, 6], dtype=torch.long, device=device), pd_page_index=1, @@ -210,7 +210,7 @@ def test_invalid_late_region_shape_cannot_partially_write_pd_page(): device="cuda", ) with pytest.raises(RuntimeError, match="shape"): - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=torch.tensor([0], dtype=torch.long, device="cuda"), pd_page_buffers=buffer.unsqueeze(0), pd_page_index=0, @@ -230,21 +230,21 @@ def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int) dtype=torch.float32, ) source = GQAFullAttentionState(config, page_size=4, device=device) - source.initialize_storage(12) + source.allocate_storage(12) target = GQAFullAttentionState(config, page_size=4, device=device) - target.initialize_storage(12) + target.allocate_storage(12) for storage_buffer in full_state_test_buffers(source).values(): storage_buffer.copy_(torch.arange(storage_buffer.numel(), device=device).reshape_as(storage_buffer)) for storage_buffer in full_state_test_buffers(target).values(): storage_buffer.fill_(-1) if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=torch.tensor([2, 0]), dst_token_indices=torch.tensor([6, 7]), remote_dp_full_att_state=source, ) return - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), dst_token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), remote_dp_full_att_state=source, @@ -257,11 +257,11 @@ def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int) single_layer_source = GQAFullAttentionState( replace(config, logical_layer_groups=config.logical_layer_groups[:1]), page_size=4, device=device ) - single_layer_source.initialize_storage(12) + single_layer_source.allocate_storage(12) for storage_buffer in full_state_test_buffers(single_layer_source).values(): storage_buffer.fill_(9) with pytest.raises(AssertionError, match="equivalent TP state configurations"): - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), dst_token_indices=torch.tensor([8, 9], dtype=torch.long, device=device), remote_dp_full_att_state=single_layer_source, diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index d9f4bdfa94..32a1849db8 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -16,7 +16,7 @@ import pytest import torch -from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState +from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState from lightllm.utils.envs_utils import get_env_start_args from unit_tests.common.att_resources.test_capacity import allocated_bytes from unit_tests.server.router.model_infer.test_cpu_resource_layout import CASES, make_runtime, runtime @@ -103,14 +103,14 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( model.tp_rank_, model.tp_world_size_ = 0, 1 model.max_total_token_num, model.max_req_num, model.mem_fraction = None, 1, 1.0 prepare_model_states(model) - model.att_resource_manager.finish_state_creation() - model.att_resource_manager.load_full_att_calibration(args=args) - model.att_resource_manager.initialize_resources( + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) - (state,) = model.att_resource_manager.get_full_attention_states("main") + (state,) = model.att_state_manager.get_full_attention_states("main") assert not hasattr(model, "full_att_state") assert state.capacity == 16 assert allocated_bytes(state) == storage_bytes if mode == "prefill": - assert model.att_resource_manager.get_pd_page_bytes() == reference.get_pd_page_bytes() + assert model.att_state_manager.get_pd_page_bytes() == reference.get_pd_page_bytes() diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py index 762208c425..ffc4607ee0 100644 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -8,8 +8,8 @@ import pytest import torch -from lightllm.common.att_resources.full_att import FullAttentionConfig, GQAFullAttentionState -from lightllm.common.att_resources.types import StateTransferLayout +from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout if TYPE_CHECKING: from easydict import EasyDict @@ -45,7 +45,7 @@ def test_gqa_pd_direct_copy_between_tp_layouts( device="cuda", ) for state in (source, target): - state.initialize_storage(token_capacity=8) + state.allocate_storage(token_capacity=8) for tensor in (source.k_storage, source.v_storage): tensor.view(torch.uint8).random_(0, 256) for tensor in (target.k_storage, target.v_storage): @@ -63,7 +63,7 @@ def test_gqa_pd_direct_copy_between_tp_layouts( shape=tuple(region.shape[region.axes.index(axis)] for axis in axes), )) layouts.append(StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=16, byte_end=max(region.byte_end for region in regions), @@ -91,13 +91,13 @@ def reject_gather(*args: object, **kwargs: object) -> None: patch.setattr(torch.Tensor, "index_copy_", reject_gather) patch.setattr(torch.Tensor, "tolist", reject_gather) with torch.cuda.stream(stream): - source.write_to_pd_buffer( + source.write_to_pd_page( token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=1, layout=layout, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=1, @@ -117,7 +117,7 @@ def reject_gather(*args: object, **kwargs: object) -> None: # 反向验证分片写出：rank 1 只修改全局 heads [2, 4)，rank 0 区域保持原始字节。 pages.fill_(91) - target.write_to_pd_buffer( + target.write_to_pd_page( token_indices=dst_ids, pd_page_buffers=page_view, pd_page_index=2, @@ -148,12 +148,12 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: page_size=4, device="cuda", ) - state.initialize_storage(token_capacity=8) + state.allocate_storage(token_capacity=8) state.k_storage.fill_(19) state.v_storage.fill_(19) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -164,13 +164,13 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: pages[1, region.byte_start:region.byte_end].view(region.dtype).fill_(value) snapshot = pages.clone() ids = torch.tensor(data=[5, 1], dtype=torch.int32, device="cuda") - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=ids, pd_page_buffers=pages, pd_page_index=1, layout=layout, ) - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, token_indices=ids, pd_page_index=1, @@ -203,7 +203,7 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt page_size=4, device="cuda:1", ) - state.initialize_storage(token_capacity=8) + state.allocate_storage(token_capacity=8) state.k_storage.zero_() state.v_storage.zero_() state.k_storage[:, 1].fill_(3) @@ -211,7 +211,7 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt torch.cuda.synchronize(device=1) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -250,13 +250,13 @@ def _peer_pd_worker(connection: Connection, payload: bytes, layout: StateTransfe stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - state.write_to_pd_buffer( + state.write_to_pd_page( token_indices=torch.tensor(data=[1], dtype=torch.int32, device="cuda:0"), pd_page_buffers=pages, pd_page_index=1, layout=layout, ) - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, token_indices=torch.tensor(data=[5], dtype=torch.int64, device="cuda:0"), pd_page_index=1, diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index 089f81d8c3..8c0441f70a 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -5,9 +5,9 @@ import pytest import torch -from lightllm.common.att_resources.full_att import FullAttentionConfig, GQAFullAttentionState -from lightllm.common.att_resources.full_att.gqa_transfer import transfer_gqa_pages -from lightllm.common.att_resources.types import StateTransferLayout +from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.full_att.gqa_transfer import transfer_gqa_pages +from lightllm.common.attention_state.types import StateTransferLayout if TYPE_CHECKING: from easydict import EasyDict @@ -26,15 +26,15 @@ def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_a ) source = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") target = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") - source.initialize_storage(token_capacity=128) - target.initialize_storage(token_capacity=128) + source.allocate_storage(token_capacity=128) + target.allocate_storage(token_capacity=128) for tensor in (source.k_storage, source.v_storage): tensor.view(torch.uint8).random_(0, 256) for tensor in (target.k_storage, target.v_storage): tensor.view(torch.uint8).fill_(173) regions = source.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=32, byte_start=0, byte_end=max(region.byte_end for region in regions), @@ -107,7 +107,7 @@ def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( pd_args: "EasyDict", monkeypatch: pytest.MonkeyPatch, ) -> None: - from lightllm.common.att_resources.full_att import gqa_transfer + from lightllm.common.attention_state.full_att import gqa_transfer pd_args.pd_kv_page_size = 2 config = FullAttentionConfig( @@ -118,8 +118,8 @@ def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( ) source = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") target = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") - source.initialize_storage(token_capacity=8) - target.initialize_storage(token_capacity=8) + source.allocate_storage(token_capacity=8) + target.allocate_storage(token_capacity=8) for name in ("k_storage", "v_storage"): storage = getattr(source, name) if dim_stride != 1: @@ -129,7 +129,7 @@ def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( getattr(target, name).view(torch.uint8).fill_(173) regions = source.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=2, byte_start=0, byte_end=max(region.byte_end for region in regions), diff --git a/unit_tests/common/att_resources/test_kpool_transfer.py b/unit_tests/common/att_resources/test_kpool_transfer.py index e58399c1e9..334d68f071 100644 --- a/unit_tests/common/att_resources/test_kpool_transfer.py +++ b/unit_tests/common/att_resources/test_kpool_transfer.py @@ -5,9 +5,9 @@ import pytest import torch -from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig -from lightllm.common.att_resources.att_tail import kpool_transfer -from lightllm.common.att_resources.types import AxisSlice, StateTransferLayout +from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.attention_state.att_tail import kpool_transfer +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -57,7 +57,7 @@ def test_kpool_single_launch_preserves_bits_and_other_slots( byte_end=byte_start + region.byte_end, ) layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=None, byte_start=byte_start, byte_end=region.byte_end, @@ -90,14 +90,14 @@ def reject_copy(*args: object, **kwargs: object) -> None: for method in ("copy_", "fill_", "zero_", "index_select", "index_copy_", "roll"): patch.setattr(torch.Tensor, method, reject_copy) with torch.cuda.stream(stream): - source.write_to_pd_buffer( + source.write_to_pd_page( req_idx=1, pd_page_buffers=pages, pd_page_index=1, state_token_len=11, layout=layout, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=pages, req_idx=2, pd_page_index=1, @@ -117,11 +117,11 @@ def reject_copy(*args: object, **kwargs: object) -> None: assert (pages[1, :byte_start] == 253).all() and (pages[1, layout.byte_end:] == 253).all() -@pytest.mark.parametrize("src_circular", [False, True]) -@pytest.mark.parametrize("dst_circular", [False, True]) +@pytest.mark.parametrize("local_circular", [False, True]) +@pytest.mark.parametrize("external_circular", [False, True]) def test_kpool_runtime_axis_overrides_are_bidirectional( - src_circular: bool, - dst_circular: bool, + local_circular: bool, + external_circular: bool, ) -> None: config = KPoolStateConfig( logical_layer_ids=(0, 2), @@ -139,20 +139,20 @@ def test_kpool_runtime_axis_overrides_are_bidirectional( original = state.tail_buffer[:, 1].clone() regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=None, byte_start=0, byte_end=regions[0].byte_end, regions=regions, ) pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - window = AxisSlice( + window = AxisMapping( axis="window", - src_start=5 if src_circular else 1, - dst_start=3 if dst_circular else 0, + local_start=5 if local_circular else 1, + external_start=3 if external_circular else 0, length=4, - src_circular=src_circular, - dst_circular=dst_circular, + local_circular=local_circular, + external_circular=external_circular, ) kpool_transfer.transfer_kpool_pages( tail=state.tail_buffer[:, 1], @@ -161,14 +161,14 @@ def test_kpool_runtime_axis_overrides_are_bidirectional( state_token_len=11, layout=layout, write=True, - axis_slice_overrides=(window,), + axis_mapping_overrides=(window,), grid_num=1, ) external = pages[0].view(config.dtype).view(2, 4, 256) for index in range(4): assert torch.equal( - external[:, (window.dst_start + index) % 4].view(torch.uint8), - original[:, (window.src_start + index) % 7].view(torch.uint8), + external[:, (window.external_start + index) % 4].view(torch.uint8), + original[:, (window.local_start + index) % 7].view(torch.uint8), ) state.tail_buffer[:, 1].fill_(37) kpool_transfer.transfer_kpool_pages( @@ -178,15 +178,15 @@ def test_kpool_runtime_axis_overrides_are_bidirectional( state_token_len=11, layout=layout, write=False, - axis_slice_overrides=(window,), + axis_mapping_overrides=(window,), grid_num=1, ) expected = torch.full_like(original, 37) for index in range(4): - slot = (window.src_start + index) % 7 + slot = (window.local_start + index) % 7 expected[:, slot] = original[:, slot] assert torch.equal(state.tail_buffer[:, 1].view(torch.uint8), expected.view(torch.uint8)) - assert layout.regions[0].axis_slices[0].src_start == 0 + assert layout.regions[0].axis_mappings[0].local_start == 0 def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: @@ -206,7 +206,7 @@ def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: ) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=None, byte_start=0, byte_end=regions[0].byte_end, @@ -214,7 +214,7 @@ def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: ) pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") pages.view(torch.bfloat16).fill_(13) - state.write_to_pd_buffer( + state.write_to_pd_page( req_idx=1, pd_page_buffers=pages, pd_page_index=0, @@ -222,7 +222,7 @@ def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: layout=layout, ) assert (pages.view(torch.bfloat16) == 13).all() - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, req_idx=1, pd_page_index=0, @@ -250,7 +250,7 @@ def test_kpool_window_positions_do_not_create_kernel_variants() -> None: ) regions = state.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=None, byte_start=0, byte_end=regions[0].byte_end, @@ -260,14 +260,14 @@ def test_kpool_window_positions_do_not_create_kernel_variants() -> None: kernel = kpool_transfer._transfer_kpool_pages cache_size = None for length in (*range(30), (1 << 33) + 7): - state.write_to_pd_buffer( + state.write_to_pd_page( req_idx=1, pd_page_buffers=pages, pd_page_index=0, state_token_len=length, layout=layout, ) - state.read_from_pd_buffer( + state.read_from_pd_page( pd_page_buffers=pages, req_idx=0, pd_page_index=0, @@ -305,7 +305,7 @@ def test_kpool_large_window_with_one_program(length: int) -> None: target.tail_buffer.fill_(37) regions = source.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=None, byte_start=0, byte_end=regions[0].byte_end, diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 116fbd05d8..c7d98e1437 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -7,8 +7,8 @@ import torch from torch.utils._pytree import tree_leaves -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.full_att import FullAttentionConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import FullAttentionConfig from lightllm.models.gemma4.model import Gemma4TpPartModel from unit_tests.common.att_resources.model_helpers import ( full_state_test_buffers, @@ -117,9 +117,9 @@ def test_shared_layers_transfer_each_physical_layer_once( buffer.zero_() managers = [] for state in (source, target): - manager = AttResourceManager(request_capacity=1) - manager.set_att_state(model_name="main", state=state) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state(model_role="main", state=state) + manager.finish_state_registration() managers.append(manager) src_ids = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") dst_ids = torch.tensor([1, 4, 7], dtype=torch.int32, device="cuda") @@ -129,29 +129,29 @@ def test_shared_layers_transfer_each_physical_layer_once( token_indices=src_ids, cpu_page_buffers=pages, cpu_page_indices=[1], - layouts=managers[0].get_cpu_transfer_layouts(), + layouts=managers[0].get_cpu_page_layouts(), ) managers[1].read_req_from_cpu( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=dst_ids, - layouts=managers[1].get_cpu_transfer_layouts(), + layouts=managers[1].get_cpu_page_layouts(), ) else: pages = torch.zeros((2, managers[0].get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - managers[0].write_to_pd_buffer( + managers[0].write_to_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=src_ids, - layouts=managers[0].get_pd_full_att_transfer_layouts(), + layouts=managers[0].get_pd_token_page_layouts(), ) - managers[1].read_from_pd_buffer( + managers[1].read_from_pd_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=dst_ids, - layouts=managers[1].get_pd_full_att_transfer_layouts(), + layouts=managers[1].get_pd_token_page_layouts(), ) torch.cuda.synchronize() source_fields = full_state_test_fields(source, source_buffers) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 94c089bf32..b3ea892073 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -6,9 +6,9 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.full_att import FullAttentionConfig, FullAttentionState, GQAFullAttentionState -from lightllm.common.att_resources.bounded_att import KDAState, GDNState +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import FullAttentionConfig, FullAttentionState, GQAFullAttentionState +from lightllm.common.attention_state.bounded_att import KDAState, GDNState from unit_tests.common.att_resources.test_bounded_state import make_state from unit_tests.common.att_resources.test_tail_state import make_tail from unit_tests.common.att_resources.test_capacity import definition @@ -36,19 +36,19 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: LlamaTpPartModel, "create_full_att_states", classmethod(lambda cls, **kwargs: [state]), ) - model.prepare_att_resources() + model.prepare_attention_states() - assert model.att_resource_manager.main_gqa_bf16_state is state + assert model.att_state_manager.main_gqa_bf16_state is state assert not hasattr(model.layers_infer[0], "full_att_state_name") assert not hasattr(model.layers_infer[0].indexer, "full_att_state_name") @pytest.mark.parametrize("role", ["main", "draft0"]) def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, monkeypatch: pytest.MonkeyPatch) -> None: - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) for model_role in ("main", "draft0"): for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): - manager.set_att_state(model_role, GQAFullAttentionState( + manager.register_state(model_role, GQAFullAttentionState( layout_config=FullAttentionConfig( logical_layer_groups=groups, global_head_num=1, @@ -61,8 +61,8 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo model.mtp_previous_draft_models = [] model.max_req_num = 0 model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) - model.main_model = SimpleNamespace(att_resource_manager=manager) - monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttResourceManager", lambda request_capacity: manager) + model.main_model = SimpleNamespace(att_state_manager=manager) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttentionStateManager", lambda request_capacity: manager) model.config, model.args = {}, SimpleNamespace(mtp_mode=None) model.tp_rank_, model.tp_world_size_ = 0, 1 model.layers_infer = [] @@ -71,132 +71,132 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo layer.layer_num_ = logical_layer_id model.layers_infer.append(layer) - model.prepare_att_resources() + model.prepare_attention_states() assert all(not hasattr(layer, "full_att_state_name") for layer in model.layers_infer) - manager.finish_state_creation() - state = manager.get_att_state(role, "gqa", 0) - assert state.resource_id == f"{role}_gqa_bf16_state" - assert manager.get_att_state(role, "gqa", 2) is state - assert manager.get_att_state(role, "gqa", 3).resource_id == f"{role}_gqa_fp16_state" + manager.finish_state_registration() + state = manager.get_state(role, "gqa", 0) + assert state.state_name == f"{role}_gqa_bf16_state" + assert manager.get_state(role, "gqa", 2) is state + assert manager.get_state(role, "gqa", 3).state_name == f"{role}_gqa_fp16_state" @pytest.mark.parametrize("finalized", [False, True]) def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> None: from dataclasses import replace - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) plain = definition("gqa") quantized = definition("int8") quantized = type(quantized)( layout_config=replace(quantized.config, logical_layer_groups=((10, 12),)), device="meta", ) - manager.set_att_state("main", plain) - manager.set_att_state("main", quantized) - manager.set_att_state("draft0", definition("gqa")) + manager.register_state("main", plain) + manager.register_state("main", quantized) + manager.register_state("draft0", definition("gqa")) logical_layer_id = plain.config.logical_layer_groups[0][0] if not finalized: - with pytest.raises(AssertionError, match="State creation must finish"): - manager.get_att_state("main", "gqa", logical_layer_id) - manager.finish_state_creation() - - assert manager.get_att_state("main", "gqa", logical_layer_id) is plain - assert manager.get_att_state("main", "gqa", 10) is quantized - assert manager.get_att_state("main", "int8", 10) is quantized - assert manager.get_att_state("main", "gqa", 12) is quantized - assert manager.get_att_state("draft0", "gqa", logical_layer_id) is not plain + with pytest.raises(AssertionError, match="State registration must finish"): + manager.get_state("main", "gqa", logical_layer_id) + manager.finish_state_registration() + + assert manager.get_state("main", "gqa", logical_layer_id) is plain + assert manager.get_state("main", "gqa", 10) is quantized + assert manager.get_state("main", "int8", 10) is quantized + assert manager.get_state("main", "gqa", 12) is quantized + assert manager.get_state("draft0", "gqa", logical_layer_id) is not plain with pytest.raises(AssertionError, match="expected exactly one"): - manager.get_att_state("main", "gqa", 99) + manager.get_state("main", "gqa", 99) def test_selection_requires_finished_creation_and_reuses_cache() -> None: - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) plain = definition("gqa") - manager.set_att_state("main", plain) + manager.register_state("main", plain) logical_layer_id = plain.config.logical_layer_groups[0][0] - with pytest.raises(AssertionError, match="State creation must finish"): - manager.get_att_state("main", "gqa", logical_layer_id) - manager.finish_state_creation() - assert manager.get_att_state("main", "gqa", logical_layer_id) is plain - hits = manager.get_att_state.cache_info().hits - assert manager.get_att_state("main", "gqa", logical_layer_id) is plain - assert manager.get_att_state.cache_info().hits == hits + 1 + with pytest.raises(AssertionError, match="State registration must finish"): + manager.get_state("main", "gqa", logical_layer_id) + manager.finish_state_registration() + assert manager.get_state("main", "gqa", logical_layer_id) is plain + hits = manager.get_state.cache_info().hits + assert manager.get_state("main", "gqa", logical_layer_id) is plain + assert manager.get_state.cache_info().hits == hits + 1 def test_selection_rejects_ambiguous_finished_members() -> None: - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) plain = definition("gqa") - manager.set_att_state("main", plain) - manager.set_att_state("main", definition("int8")) - manager.finish_state_creation() + manager.register_state("main", plain) + manager.register_state("main", definition("int8")) + manager.finish_state_registration() with pytest.raises(AssertionError, match="expected exactly one"): - manager.get_att_state("main", "gqa", plain.config.logical_layer_groups[0][0]) + manager.get_state("main", "gqa", plain.config.logical_layer_groups[0][0]) @pytest.mark.parametrize("state_class", [KDAState, GDNState]) def test_bounded_state_layer_ownership(state_class: type[KDAState] | type[GDNState]) -> None: state = make_state(cls=state_class) - manager = AttResourceManager(request_capacity=3) - manager.set_att_state("main", state) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", state) + manager.finish_state_registration() assert state.contains_logical_layer(1) and state.contains_logical_layer(3) assert not state.contains_logical_layer(0) - assert manager.get_att_state("main", state.get_state_type(), 3) is state + assert manager.get_state("main", state.get_state_type(), 3) is state def test_mla_selection_excludes_same_layer_tail() -> None: from dataclasses import replace - manager = AttResourceManager(request_capacity=3) + manager = AttentionStateManager(request_capacity=3) tail = make_tail() template = definition("mla") state = type(template)( layout_config=replace(template.config, logical_layer_groups=((3, 7),)), device="meta", ) - manager.set_att_state("main", state) - manager.set_att_state("main", tail) - manager.finish_state_creation() + manager.register_state("main", state) + manager.register_state("main", tail) + manager.finish_state_registration() assert tail.contains_logical_layer(3) and tail.contains_logical_layer(7) assert not tail.contains_logical_layer(0) - assert manager.get_att_state("main", "mla", 7) is state - assert manager.get_att_state("main", "kpool", 7) is tail - assert manager.get_att_state( - model_name="main", - att_type="", + assert manager.get_state("main", "mla", 7) is state + assert manager.get_state("main", "kpool", 7) is tail + assert manager.get_state( + model_role="main", + state_type_filter="", logical_layer_id=7, state_class=FullAttentionState, ) is state with pytest.raises(AssertionError, match="expected exactly one"): - manager.get_att_state("main", "", 7) + manager.get_state("main", "", 7) @pytest.mark.parametrize("role", ["draft1", "draft10"]) def test_query_matches_complete_model_role(role: str) -> None: - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) for model_name in ("draft1", "draft10"): - manager.set_att_state(model_name, definition("gqa")) - manager.finish_state_creation() - assert manager.get_att_state(role, "gqa", 0).resource_id == f"{role}_gqa_bf16_state" + manager.register_state(model_name, definition("gqa")) + manager.finish_state_registration() + assert manager.get_state(role, "gqa", 0).state_name == f"{role}_gqa_bf16_state" @pytest.mark.parametrize("algorithm", ["kda", "gdn", "kpool"]) def test_precision_variants_have_distinct_names_and_can_coexist(algorithm: str) -> None: from dataclasses import replace - manager = AttResourceManager(request_capacity=3) + manager = AttentionStateManager(request_capacity=3) template = make_tail() if algorithm == "kpool" else make_state(cls=KDAState if algorithm == "kda" else GDNState) states = [] for dtype in (torch.bfloat16, torch.float32): config = replace(template.config, **({"dtype": dtype} if algorithm == "kpool" else {"conv_dtype": dtype})) state = type(template)(config=config, request_capacity=3, device="cpu") - manager.set_att_state("main", state) + manager.register_state("main", state) states.append(state) - manager.finish_state_creation() + manager.finish_state_registration() logical_layer_id = template.config.logical_layer_ids[0] for state in states: - assert manager.get_att_state("main", state.get_state_type(), logical_layer_id) is state - assert getattr(manager, state.resource_id) is state - assert states[0].resource_id != states[1].resource_id + assert manager.get_state("main", state.get_state_type(), logical_layer_id) is state + assert getattr(manager, state.state_name) is state + assert states[0].state_name != states[1].state_name with pytest.raises(AssertionError, match="expected exactly one"): - manager.get_att_state("main", algorithm, logical_layer_id) + manager.get_state("main", algorithm, logical_layer_id) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 04aa3124bc..83753c87d7 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -10,11 +10,11 @@ import pytest import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator, StatePageAllocator -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.types import PageKind, StateTransferLayout +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.types import PageKind, StateTransferLayout from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager -from lightllm.common.att_resources.att_tail import KPoolAttentionTailState +from lightllm.common.attention_state.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state, layout_for from unit_tests.common.att_resources.test_bounded_state import page_layout from unit_tests.common.att_resources.test_tail_state import make_tail, tail_layout @@ -27,7 +27,7 @@ def composition(device: str | torch.device | int = "cpu", with_tail: bool = True return MainDraftTestResourceManager(device, with_tail=with_tail, tail_pool_size=tail_pool_size) -def layouts_for(manager: AttResourceManager, kind: str = "cpu"): +def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): layouts = [] if kind != "att_state": layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] @@ -39,7 +39,7 @@ def layouts_for(manager: AttResourceManager, kind: str = "cpu"): result = {} for layout in layouts: state_end = offset + layout.byte_end - result[layout.resource_id] = replace( + result[layout.state_name] = replace( layout, byte_start=offset + layout.byte_start, byte_end=state_end, @@ -52,7 +52,7 @@ def layouts_for(manager: AttResourceManager, kind: str = "cpu"): return result -def init_source(manager: AttResourceManager): +def init_source(manager: AttentionStateManager): for i, state in enumerate((manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.copy_( @@ -67,11 +67,11 @@ def init_source(manager: AttResourceManager): state.tail_buffer[:, 0].fill_(3) state.tail_buffer[:, 1].fill_(9) slots = manager.bounded_att_big_page_allocator.alloc(2) - manager.save_bounded_att_pages( + manager.save_bounded_att_runtime_to_pages( req_indices=[0, 1], state_token_lens=[4, 8], page_kind="big", - slot_ids=slots, + page_slot_ids=slots, ) if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: torch.cuda.synchronize() @@ -82,11 +82,11 @@ def init_source(manager: AttResourceManager): def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.MonkeyPatch): manager = composition(device="cuda") slots = init_source(manager) - manager.restore_bounded_att_pages( + manager.restore_bounded_att_runtime_from_pages( req_indices=[2], state_token_lens=[4], page_kind="big", - slot_ids=[slots[0]], + page_slot_ids=[slots[0]], ) for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): assert (state.conv_state[:, 2, :, :2] == 11 + i).all() @@ -95,13 +95,13 @@ def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.Mon def fail(*args: object, **kwargs: object): raise RuntimeError("simulated copy failure") - monkeypatch.setattr(manager.draft0_gdn_conv_fp32_ssm_fp32_state, "save_bounded_att_pages", fail) + monkeypatch.setattr(manager.draft0_gdn_conv_fp32_ssm_fp32_state, "save_bounded_att_runtime_to_pages", fail) with pytest.raises(RuntimeError): - manager.save_bounded_att_pages( + manager.save_bounded_att_runtime_to_pages( req_indices=[0, 1], state_token_lens=[4, 8], page_kind="big", - slot_ids=slots, + page_slot_ids=slots, ) assert manager.bounded_att_big_page_allocator.free_count == 1 manager.bounded_att_big_page_allocator.free(slots) @@ -109,26 +109,26 @@ def fail(*args: object, **kwargs: object): def test_disabled_domains(): - manager = AttResourceManager(request_capacity=3) - manager.set_att_state("main", full_state()) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", full_state()) manager.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() manager.reset_bounded_att_state([0]) manager.reset_attention_tail_state([0]) with pytest.raises(AssertionError, match="disabled"): - manager.save_bounded_att_pages( + manager.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[4], page_kind="big", - slot_ids=[0], + page_slot_ids=[0], ) - empty_manager = AttResourceManager(request_capacity=3) - empty_manager.finish_state_creation() + empty_manager = AttentionStateManager(request_capacity=3) + empty_manager.finish_state_registration() with pytest.raises(AssertionError, match="no full-attention"): empty_manager.copy_full_att_tokens(torch.tensor([0], dtype=torch.long), torch.tensor([1], dtype=torch.long)) -@pytest.mark.parametrize("operation", ["save_bounded_att_pages", "restore_bounded_att_pages"]) +@pytest.mark.parametrize("operation", ["save_bounded_att_runtime_to_pages", "restore_bounded_att_runtime_from_pages"]) def test_page_operations_forward_each_request_length(operation: str, monkeypatch: pytest.MonkeyPatch): manager = composition(with_tail=False) calls = [] @@ -137,9 +137,9 @@ def record_lengths( req_indices: list[int], state_token_lens: list[int], page_kind: PageKind, - slot_ids: list[int], + page_slot_ids: list[int], ) -> None: - calls.append((req_indices, state_token_lens, page_kind, slot_ids)) + calls.append((req_indices, state_token_lens, page_kind, page_slot_ids)) for state in manager.get_bounded_attention_states(): monkeypatch.setattr(state, operation, record_lengths) @@ -147,7 +147,7 @@ def record_lengths( req_indices=[0, 2], state_token_lens=[17, 29], page_kind="big", - slot_ids=[0, 1], + page_slot_ids=[0, 1], ) assert calls == [([0, 2], [17, 29], "big", [0, 1])] * len(manager.get_bounded_attention_states()) @@ -160,18 +160,18 @@ def test_cpu_restore_forwards_absolute_length_for_partial_prefix(monkeypatch: py for state in manager.get_full_attention_states(): monkeypatch.setattr(state, "read_from_cpu", lambda **kwargs: None) for state in manager.get_bounded_attention_states(): - monkeypatch.setattr(state, "read_page_state_from_cpu", lambda **kwargs: None) + monkeypatch.setattr(state, "read_state_pages_from_cpu", lambda **kwargs: None) restored = [] def record_restore( req_indices: list[int], state_token_lens: list[int], page_kind: PageKind, - slot_ids: list[int], + page_slot_ids: list[int], ) -> None: - restored.append((req_indices, state_token_lens, page_kind, slot_ids)) + restored.append((req_indices, state_token_lens, page_kind, page_slot_ids)) - monkeypatch.setattr(manager, "restore_bounded_att_pages", record_restore) + monkeypatch.setattr(manager, "restore_bounded_att_runtime_from_pages", record_restore) # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 manager.read_req_from_cpu( cpu_page_buffers=torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), @@ -185,7 +185,7 @@ def record_restore( assert restored == [([1], [70], "big", [2])] -@pytest.mark.parametrize("operation", ["write_to_pd_buffer", "read_from_pd_buffer"]) +@pytest.mark.parametrize("operation", ["write_to_pd_page", "read_from_pd_page"]) def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): manager = composition(with_tail=False) layouts = layouts_for(manager, kind="att_state") @@ -224,10 +224,10 @@ def test_reset_tail_clears_all_tail_states_and_preserves_linear(): ) first.tail_buffer.fill_(5) second.tail_buffer.fill_(7) - tail_manager = AttResourceManager(request_capacity=3) - tail_manager.set_att_state("main", first) - tail_manager.set_att_state("main", second) - tail_manager.finish_state_creation() + tail_manager = AttentionStateManager(request_capacity=3) + tail_manager.register_state("main", first) + tail_manager.register_state("main", second) + tail_manager.finish_state_registration() tail_manager.reset_attention_tail_state([0, 1]) assert not first.tail_buffer[:, :2].any() and not second.tail_buffer[:, :2].any() assert (first.tail_buffer[:, 2] == 5).all() and (second.tail_buffer[:, 2] == 7).all() @@ -275,7 +275,7 @@ def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | tor torch.cuda.synchronize() tail = getattr(target, f"main_{target.att_tail_type}_state") tail.tail_buffer.fill_(37) - assert tail.resource_id not in layouts + assert tail.state_name not in layouts target_slots = target.bounded_att_big_page_allocator.alloc(2) target.read_req_from_cpu( cpu_page_buffers=pages, @@ -321,11 +321,11 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i small_slots = [] if token_count % 4: small_slots = source.bounded_att_small_page_allocator.alloc(1) - source.save_bounded_att_pages( + source.save_bounded_att_runtime_to_pages( req_indices=[page_count - 1], state_token_lens=[token_count], page_kind="small", - slot_ids=small_slots, + page_slot_ids=small_slots, ) torch.cuda.synchronize() source.write_req_to_cpu( @@ -417,14 +417,14 @@ def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.d page = torch.empty(max(layout.byte_end for layout in layouts.values()), dtype=torch.uint8, device=device) src_args = {"token_indices": torch.tensor([0, 2], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 1, "state_token_len": 8} dst_args = {"token_indices": torch.tensor([5, 7], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 2, "state_token_len": 8} - source.write_to_pd_buffer( + source.write_to_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=layouts, **src_args, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, @@ -433,10 +433,10 @@ def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.d ) dp_tail = dp_target.get_attention_tail_states()[0] dp_tail.tail_buffer.fill_(19) - dp_target.copy_from_other_dp( + dp_target.copy_from_dp( remote_att_manager=source, - src_token_indexes=torch.tensor([0, 2], dtype=torch.int32, device=device), - dst_token_indexes=torch.tensor([5, 7], dtype=torch.long, device=device), + src_token_indices=torch.tensor([0, 2], dtype=torch.int32, device=device), + dst_token_indices=torch.tensor([5, 7], dtype=torch.long, device=device), src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=1, @@ -469,10 +469,10 @@ def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): tails = target.get_attention_tail_states() for state in tails: state.tail_buffer.fill_(19) - target.copy_from_other_dp( + target.copy_from_dp( remote_att_manager=source, - src_token_indexes=torch.empty(0, dtype=torch.long), - dst_token_indexes=torch.empty(0, dtype=torch.long), + src_token_indices=torch.empty(0, dtype=torch.long), + dst_token_indices=torch.empty(0, dtype=torch.long), src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=1, @@ -507,10 +507,10 @@ def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: boo for state in target.get_attention_tail_states(): state.tail_buffer.fill_(51) - target.copy_from_other_dp( + target.copy_from_dp( remote_att_manager=source, - src_token_indexes=torch.tensor([0, 2] if copy_tokens else [], dtype=torch.long, device=device), - dst_token_indexes=torch.tensor([5, 7] if copy_tokens else [], dtype=torch.long, device=device), + src_token_indices=torch.tensor([0, 2] if copy_tokens else [], dtype=torch.long, device=device), + dst_token_indices=torch.tensor([5, 7] if copy_tokens else [], dtype=torch.long, device=device), src_big_page_ids=[2, 0] if copy_pages else [], dst_big_page_ids=[0, 1] if copy_pages else [], src_req_idx=1, @@ -518,7 +518,7 @@ def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: boo ) for state in target.get_full_attention_states(): - src_state = getattr(source, state.resource_id) + src_state = getattr(source, state.state_name) values = read_kv_storage(state) if copy_tokens: assert torch.equal(values[:, [5, 7]], read_kv_storage(src_state)[:, [0, 2]]) @@ -526,7 +526,7 @@ def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: boo assert (values == -1).all() assert (values[:, :5] == -1).all() for state in target.get_bounded_attention_states(): - src_pool = getattr(source, state.resource_id).big_pages_pool + src_pool = getattr(source, state.state_name).big_pages_pool dst_pool = state.big_pages_pool if copy_pages: for src_page, dst_page in ((2, 0), (0, 1)): @@ -537,7 +537,7 @@ def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: boo assert (dst_pool.conv_state[2] == -1).all() and (dst_pool.ssm_state[2] == -2).all() assert (state.small_pages_pool.conv_state == 31).all() assert (state.small_pages_pool.ssm_state == 32).all() - source_state = getattr(source, state.resource_id) + source_state = getattr(source, state.state_name) assert torch.equal( state.conv_state[:, 2, :, :state.config.conv_width], source_state.conv_state[:, 1, :, :source_state.config.conv_width], @@ -562,11 +562,11 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo source = composition("cuda", with_tail=False) big_slots = init_source(source) small_slots = source.bounded_att_small_page_allocator.alloc(1) - source.save_bounded_att_pages( + source.save_bounded_att_runtime_to_pages( req_indices=[1], state_token_lens=[6], page_kind="small", - slot_ids=small_slots, + page_slot_ids=small_slots, ) torch.cuda.synchronize() layouts = layouts_for(source) @@ -600,7 +600,7 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo cpu_page_buffers=pages, cpu_page_indices=page_ids, layouts=layouts, - cpu_page_readies=readies, + cpu_page_ready_flags=readies, big_page_slot_ids=masked_big_slots, small_page_slot_ids=masked_small_slots, ) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index f96e96fca6..5439dd3bb2 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -12,7 +12,7 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.models.qwen3next.model import Qwen3NextTpPartModel @@ -96,10 +96,10 @@ def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encod layer_start += layer_count # 反向创建、注册，标定仍按 main 和 draft 的数值顺序绑定。 - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) for index in reversed(range(len(models))): model = models[index] - model.att_resource_manager = manager + model.att_state_manager = manager states = [] if not (model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att")): states = model.create_full_att_states( @@ -110,13 +110,13 @@ def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encod tp_world_size=model.tp_world_size_, ) for state in states: - assert state.resource_id is None + assert state.state_name is None assert state.scales is None and state.q_scales is None - manager.set_att_state(model.att_resource_role, state) - manager.finish_state_creation() + manager.register_state(model.att_state_role, state) + manager.finish_state_registration() manager.load_full_att_calibration(args=args) for index, model in enumerate(models): - states = manager.get_full_attention_states(model.att_resource_role) + states = manager.get_full_attention_states(model.att_state_role) state = states[0] if states else None expected_layers, (start, end) = definitions[index][-2:] if not expected_layers: @@ -134,7 +134,7 @@ def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, tp_world_size: int, ): from builtins import open as original_open - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state # 第 0 行属于普通精度的共享层；其余行分别属于 main、draft2、draft10 的 FP8 State。 scales = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) @@ -158,7 +158,7 @@ def count_open(file: str, *args: object, **kwargs: object): monkeypatch.setattr("builtins.open", count_open) for tp_rank in range(tp_world_size): - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) expected = [] # 特意打乱角色及同角色 State 的注册顺序。 for role, groups, row, encoding in ( @@ -177,12 +177,12 @@ def count_open(file: str, *args: object, **kwargs: object): tp_world_size=tp_world_size, ), kv_type=encoding, device="meta", ) - manager.set_att_state(role, state) + manager.register_state(role, state) if encoding != "None": expected.append((state, row)) - with pytest.raises(AssertionError, match="State creation"): + with pytest.raises(AssertionError, match="State registration"): manager.load_full_att_calibration(args=args) - manager.finish_state_creation() + manager.finish_state_registration() manager.load_full_att_calibration(args=args) heads = 4 // tp_world_size local_heads = slice(tp_rank * heads, (tp_rank + 1) * heads) @@ -226,7 +226,7 @@ def test_cpu_layout_does_not_require_calibration_file( @pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): - from lightllm.common.att_resources.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state state = create_gqa_state( config=FullAttentionConfig( @@ -236,16 +236,16 @@ def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): dtype=torch.bfloat16, ), kv_type=encoding, device="cpu", ) - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", state) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() with pytest.raises(AssertionError, match="calibration must be bound"): - state.initialize_storage(4) + state.allocate_storage(4) if encoding == "fp8kv_sph": - with pytest.raises(AssertionError, match="requires a calibration file"): + with pytest.raises(AssertionError, match="requires a KV calibration file"): manager.load_full_att_calibration(args=StartArgs()) else: manager.load_full_att_calibration(args=StartArgs()) assert torch.equal(state.scales, torch.ones((2, 2))) - state.initialize_storage(4) + state.allocate_storage(4) assert state.k_storage.shape[0] == 2 diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 8296c60bf9..7745ef9a32 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -3,7 +3,7 @@ from unit_tests.common.att_resources.model_helpers import read_kv_storage from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState from unit_tests.common.att_resources.model_helpers import owned_states import dataclasses @@ -98,7 +98,7 @@ def test_real_model_resource_hooks_and_shared_draft( observed = [] def check_runtime_before_profile(device: str | torch.device | int): - resources = model.att_resource_manager + resources = model.att_state_manager states = resources.get_full_attention_states() assert main_linear_state(resources).conv_state.is_cuda assert main_linear_state(resources).big_pages_pool is None @@ -130,8 +130,8 @@ def check_runtime_before_profile(device: str | torch.device | int): ) draft.max_seq_length, draft.batch_max_tokens = model.max_seq_length, model.batch_max_tokens prepare_model_states(draft, logical_layer_ids=(4,)) - assert draft.att_resource_manager is model.att_resource_manager - assert read_kv_storage(owned_states(draft.att_resource_manager.__dict__, FullAttentionState)[0]) is None + assert draft.att_state_manager is model.att_state_manager + assert read_kv_storage(owned_states(draft.att_state_manager.__dict__, FullAttentionState)[0]) is None drafts.append(draft) models = (model, *drafts) req_manager = ReqManager( @@ -141,13 +141,13 @@ def check_runtime_before_profile(device: str | torch.device | int): for item in models: item.bind_req_manager(req_manager=req_manager) assert all(item.req_manager is req_manager for item in models) - model.att_resource_manager.finish_state_creation() - model.att_resource_manager.load_full_att_calibration(args=args) - model.att_resource_manager.initialize_resources( + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction ) - model.req_manager.bind_hold_page(model.att_resource_manager.full_att_allocator.hold_token_indices) - resources = model.att_resource_manager + model.req_manager.bind_hold_page(model.att_state_manager.full_att_allocator.hold_token_indices) + resources = model.att_state_manager (full_state,) = resources.get_full_attention_states("main") assert type(model.req_manager) is ReqManager assert not hasattr(model.req_manager, "req_to_ssm_state") @@ -164,13 +164,13 @@ def check_runtime_before_profile(device: str | torch.device | int): assert not hasattr(model, "mtp_state_cursors") assert (main_linear_state(resources).mtp_state_index is not None) == bool(mtp) if mode == "prefill": - assert resources.get_pd_full_att_transfer_layouts() - assert resources.get_pd_bounded_att_or_tail_transfer_layouts() + assert resources.get_pd_token_page_layouts() + assert resources.get_pd_runtime_page_layouts() if mtp: draft = drafts[0] assert draft.req_manager is model.req_manager - assert draft.att_resource_manager is resources - assert draft.att_resource_manager.full_att_allocator is resources.full_att_allocator + assert draft.att_state_manager is resources + assert draft.att_state_manager.full_att_allocator is resources.full_att_allocator assert not hasattr(draft, "full_att_state") (draft_state,) = resources.get_full_attention_states("draft0") assert draft_state is not full_state diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 68de6067a5..57b521a440 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -13,13 +13,13 @@ import pytest import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator from unit_tests.common.att_resources.model_helpers import create_full_attention_roles -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, FullAttentionConfig, ) -from lightllm.common.att_resources.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig +from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -68,7 +68,7 @@ def make_resources( ), page_size=4, ) - full.initialize_storage(16) + full.allocate_storage(16) else: full = GQAFullAttentionState( FullAttentionConfig( @@ -81,7 +81,7 @@ def make_resources( ), page_size=4, ) - full.initialize_storage(16) + full.allocate_storage(16) main, draft = create_full_attention_roles(full, 4) resources = build_hybrid_resources( config=config, @@ -137,11 +137,11 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( for h in range(cfg.num_v_heads): state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_v_heads + h + 7) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=[slot], + page_slot_ids=[slot], ) torch.cuda.synchronize() expected.append((state.big_pages_pool.conv_state[slot].clone(), state.big_pages_pool.ssm_state[slot].clone())) @@ -152,7 +152,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=buffer, cpu_page_indices=[1], - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=[0], ) torch.cuda.synchronize() @@ -170,7 +170,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( device="cuda", ), state_token_len=8, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=[slot], req_idx=1, ) @@ -199,5 +199,5 @@ def test_cpu_layout_query_does_not_require_cpu_cache_enabled(monkeypatch: pytest rank=0, mtp=0, ) - assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_transfer_layouts().values()) + assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_page_layouts().values()) get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index d8c57ed244..067063166c 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -29,8 +29,8 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state if TYPE_CHECKING: @@ -123,20 +123,20 @@ def test_model_and_cpu_service_use_identical_storage_geometry( model.tp_rank_, model.tp_world_size_ = rank, args.tp model.max_total_token_num, model.max_req_num = 19, 1 prepare_model_states(model) - resources = model.att_resource_manager - resources.finish_state_creation() - resources.initialize_resources(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) + resources = model.att_state_manager + resources.finish_state_registration() + resources.initialize_storage(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) (state,) = resources.get_full_attention_states("main") assert state.config.head_num == local_heads and state.config.head_dim == dim storage_layers = 2 - (config.get("num_kv_shared_layers") or 0) assert state.k_storage.shape == state.v_storage.shape == (storage_layers, 20, local_heads, dim) if storage_layers == 1: assert state.config.logical_layer_groups == ((0, 1),) - assert resources.get_att_state("main", "gqa", 0) is resources.get_att_state("main", "gqa", 1) + assert resources.get_state("main", "gqa", 0) is resources.get_state("main", "gqa", 1) assert state.get_layer_view(logical_layer_id=0)[0].data_ptr() == state.get_layer_view(logical_layer_id=1)[0].data_ptr() assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) assert resources.full_att_allocator.size == 16 - layout = resources.get_cpu_transfer_layouts()[state.resource_id] + layout = resources.get_cpu_page_layouts()[state.state_name] assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) # 按物理层数计算传输页；共享消费者不重复增加 KV 或传输容量。 expected_pd_bytes = storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() @@ -153,10 +153,10 @@ def test_gpu_profile_counts_token_fields_hold_and_transfer_reserve( pd_args: EasyDict, ): state = make_quant_state(kind, "cpu") - reference = AttResourceManager(request_capacity=1) - reference.set_att_state("main", state) + reference = AttentionStateManager(request_capacity=1) + reference.register_state("main", state) reference.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) - reference.finish_state_creation() + reference.finish_state_registration() args = StartArgs( page_size=4, run_mode="prefill", @@ -177,9 +177,9 @@ def memory_info(device: str | torch.device | int): monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.att_resources.att_resource_manager.get_current_device_id", lambda: 3) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 3) monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) template = rebuild_full_state( template=state, @@ -187,10 +187,10 @@ def memory_info(device: str | torch.device | int): token_capacity=0, device="meta", ) - resources = AttResourceManager(request_capacity=1) - resources.set_att_state("main", rebuild_full_state(template, config=template.config)) - resources.finish_state_creation() - resources.initialize_resources( + resources = AttentionStateManager(request_capacity=1) + resources.register_state("main", rebuild_full_state(template, config=template.config)) + resources.finish_state_registration() + resources.initialize_storage( args=args, token_capacity=None, mem_fraction=0.75, @@ -205,7 +205,7 @@ def test_cpu_layout_uses_all_model_factory_states( args: StartArgs | EasyDict, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ): """同一模型返回多种方案和共享层时，CPU 页采用实际 State 列表。""" - from lightllm.common.att_resources.full_att import FullAttentionState, GQAFullAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState, GQAFullAttentionState class MixedStorageModel(LlamaTpPartModel): @classmethod @@ -250,7 +250,7 @@ def create_full_att_states( model.config, model.args = config, args model.tp_rank_, model.tp_world_size_, model.max_req_num = 0, 1, 1 prepare_model_states(model) - model.att_resource_manager.finish_state_creation() - runtime_layouts = model.att_resource_manager.get_cpu_transfer_layouts() + model.att_state_manager.finish_state_registration() + runtime_layouts = model.att_state_manager.get_cpu_page_layouts() assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_state", "main_gqa_fp16_state"} assert cpu_layouts == runtime_layouts diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index edc514cb6d..a50ff4fc71 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -8,21 +8,21 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.bounded_att import GDNState -from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.attention_state.allocators import StatePageAllocator from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens from unit_tests.common.att_resources.test_bounded_state import make_state def make_manager(device: str | torch.device | int): - manager = AttResourceManager(request_capacity=3) - manager.set_att_state("main", make_state(device=device)) - manager.set_att_state("main", make_state(GDNState, device=device)) - manager.set_att_state("draft0", make_state(GDNState, device=device)) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", make_state(device=device)) + manager.register_state("main", make_state(GDNState, device=device)) + manager.register_state("draft0", make_state(GDNState, device=device)) manager.bounded_att_big_page_allocator = StatePageAllocator(3) manager.bounded_att_small_page_allocator = StatePageAllocator(2) - manager.finish_state_creation() + manager.finish_state_registration() return manager @@ -55,8 +55,8 @@ def test_commit_selects_state_for_one_role(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_commit_without_recurrent_state_and_invalid_role(): - manager = AttResourceManager(request_capacity=3) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=3) + manager.finish_state_registration() manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), @@ -81,7 +81,7 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i manager = make_manager("cuda") backend = SimpleNamespace( model=SimpleNamespace( - att_resource_manager=manager, + att_state_manager=manager, req_manager=SimpleNamespace( req_sampling_params_manager=SimpleNamespace( req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index f450ed9ce9..d2a63e1ec6 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -30,8 +30,8 @@ if TYPE_CHECKING: from collections.abc import Sequence - from lightllm.common.att_resources.att_resource_manager import AttResourceManager - from lightllm.common.att_resources.full_att.packed_mla_state import PackedMLAFullAttentionState + from lightllm.common.attention_state.manager import AttentionStateManager + from lightllm.common.attention_state.full_att.packed_mla_state import PackedMLAFullAttentionState from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter from multiprocessing.connection import Connection @@ -129,7 +129,7 @@ def _task(peer: PDAgentMetadata, kind: str): def _snapshot( - resources: AttResourceManager, + resources: AttentionStateManager, full: PackedMLAFullAttentionState, indices: Sequence[int] | torch.Tensor, req_idx: int, @@ -222,7 +222,7 @@ def _worker(connection: Connection, backend: str, sender: bool): task_kind=kind, req_idx=1, state_token_len=task.end_kv_index, - write_to_pd_buffer=True, + write_to_pd_page=True, ) # The production sender publishes a page only after its copy event completes. copy_stream.synchronize() @@ -273,7 +273,7 @@ def _worker(connection: Connection, backend: str, sender: bool): task_kind=kind, req_idx=2, state_token_len=request.end_kv_index, - write_to_pd_buffer=False, + write_to_pd_page=False, ) copy_stream.synchronize() expected = _snapshot( diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index 0887dcaeff..67a53fb694 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -8,7 +8,7 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.server.core.objs.start_args_type import StartArgs @@ -17,8 +17,8 @@ def draft_model(mode: str): - resources = AttResourceManager(request_capacity=3) - resources.set_att_state("main", definition("gqa")) + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", definition("gqa")) draft = LlamaTpPartModel.__new__(LlamaTpPartModel) draft.is_mtp_draft_model = True draft.args = StartArgs(mtp_mode=mode, data_type="bfloat16") @@ -34,7 +34,7 @@ def draft_model(mode: str): main = LlamaTpPartModel.__new__(LlamaTpPartModel) main.args, main.config = draft.args, dict(draft.config, num_hidden_layers=2) main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(2)] - main.att_resource_manager = resources + main.att_state_manager = resources draft.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) draft.main_model, draft.mtp_previous_draft_models = main, [] return draft, resources @@ -44,11 +44,11 @@ def draft_model(mode: str): @pytest.mark.parametrize("prefill", [False, True]) def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.MonkeyPatch, mode: str, prefill: bool): draft, resources = draft_model(mode) - draft.prepare_att_resources() - resources.finish_state_creation() + draft.prepare_attention_states() + resources.finish_state_registration() draft._init_att_backend() draft._init_att_backend1() - assert draft.att_resource_manager is resources + assert draft.att_state_manager is resources assert not hasattr(draft, "full_att_state") assert not resources.get_full_attention_states("draft0") assert not any(name.startswith("draft0_") and name.endswith("_state") for name in vars(resources)) @@ -62,7 +62,7 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo model_input = _make_prefill_input([10, 11], req_idx=0, is_decode_req=False) if prefill else _make_decode_input() state = draft._create_inferstate(model_input) assert not hasattr(state, "full_att_state") - assert state.att_resource_manager is resources + assert state.att_state_manager is resources assert state.is_draft_model assert state.prefill_att_state is None and state.decode_att_state is None assert state.req_manager is draft.req_manager @@ -71,11 +71,11 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo @pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: str): draft, resources = draft_model(mode) - draft.prepare_att_resources() - assert draft.att_resource_manager is resources - resources.finish_state_creation() + draft.prepare_attention_states() + assert draft.att_state_manager is resources + resources.finish_state_registration() (state,) = resources.get_full_attention_states("draft0") assert not hasattr(draft, "full_att_state") - assert state.resource_id == "draft0_gqa_bf16_state" + assert state.state_name == "draft0_gqa_bf16_state" assert state.config.logical_layer_groups == ((2,),) assert read_kv_storage(state) is None diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index 17c3b44141..b279ea4ce1 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -7,8 +7,8 @@ import torch import torch.multiprocessing as mp -from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNPagesPool, KDAPagesPool -from lightllm.common.att_resources.bounded_att import utils +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNPagesPool, KDAPagesPool +from lightllm.common.attention_state.bounded_att import utils if TYPE_CHECKING: from easydict import EasyDict diff --git a/unit_tests/common/att_resources/test_paged_storage.py b/unit_tests/common/att_resources/test_paged_storage.py index fae8ecf18a..e146261f11 100644 --- a/unit_tests/common/att_resources/test_paged_storage.py +++ b/unit_tests/common/att_resources/test_paged_storage.py @@ -2,7 +2,7 @@ import pytest -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator def test_kv_cache_allocation_requires_complete_pages(): diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 2560d0eefa..21b5773a47 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -15,7 +15,7 @@ import pytest import torch -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, StaticFP8AttentionConfig, @@ -23,13 +23,13 @@ PackedMLAAttentionConfig, PackedMLAFullAttentionState, ) -from lightllm.common.att_resources.types import ( - AxisSlice, StateTransferLayout, TransferRegion, +from lightllm.common.attention_state.types import ( + AxisMapping, StateTransferLayout, TransferRegion, ) if TYPE_CHECKING: from collections.abc import Sequence - from lightllm.common.att_resources.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.full_att.base_state import FullAttentionState KINDS = ("int4", "int8", "fp8_tensor", "fp8_head", "packed", "packed_fp8") @@ -45,7 +45,7 @@ def make_quant_state(kind: str, device: str | torch.device | int): bits=int(kind[-1]), ) prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) - prepared_full_state.initialize_storage(8) + prepared_full_state.allocate_storage(8) return prepared_full_state if kind.startswith("fp8_"): per_head = kind == "fp8_head" @@ -67,7 +67,7 @@ def make_quant_state(kind: str, device: str | torch.device | int): state=prepared_full_state, scales=scales, q_scales=torch.ones((2, config.head_num), dtype=torch.float32) if per_head else None, ) - prepared_full_state.initialize_storage(8) + prepared_full_state.allocate_storage(8) return prepared_full_state fp8 = kind == "packed_fp8" config = PackedMLAAttentionConfig( @@ -78,7 +78,7 @@ def make_quant_state(kind: str, device: str | torch.device | int): fp8_nope_dim=512 if fp8 else 0, ) prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) - prepared_full_state.initialize_storage(8) + prepared_full_state.allocate_storage(8) return prepared_full_state @@ -96,11 +96,11 @@ def transfer_layout(state: FullAttentionState): dtype=field.tensor.dtype, axes=field.axes, shape=shape, - axis_slices=tuple( - AxisSlice( + axis_mappings=tuple( + AxisMapping( axis=axis, - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=dim, ) for axis, dim in zip(field.axes, shape) if axis not in ("layer", "token") @@ -109,7 +109,7 @@ def transfer_layout(state: FullAttentionState): ) offset += prod(shape) * itemsize return StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=4, byte_start=0, byte_end=(offset + 15) // 16 * 16, @@ -185,27 +185,27 @@ def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( page = torch.empty(layout.byte_end, dtype=torch.uint8, device=device) if device == "cpu": with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - source.write_to_pd_buffer( + source.write_to_pd_page( token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), pd_page_buffers=page.unsqueeze(0), pd_page_index=0, layout=layout, ) with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page.unsqueeze(0), token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, layout=layout, ) else: - source.write_to_pd_buffer( + source.write_to_pd_page( token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), pd_page_buffers=page.unsqueeze(0), pd_page_index=0, layout=layout, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page.unsqueeze(0), token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), pd_page_index=0, @@ -235,11 +235,11 @@ def reject_index_conversion(tensor: torch.Tensor, *args, **kwargs): patch.setattr(torch.Tensor, "to", reject_index_conversion) if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=src_indices, dst_token_indices=dst_indices, remote_dp_full_att_state=source, ) return - target.copy_tokens_from_other_dp( + target.copy_tokens_from_dp( src_token_indices=src_indices, dst_token_indices=dst_indices, remote_dp_full_att_state=source, @@ -309,6 +309,6 @@ def test_deferred_storage_preserves_real_fp8_calibration(kind: str): assert read_kv_storage(state) is None and not state.scales.is_meta assert torch.equal(state.scales, source.scales) state.device = torch.device("cpu") - state.initialize_storage(source.capacity) + state.allocate_storage(source.capacity) assert torch.equal(state.scales, source.scales) assert torch.equal(torch.from_numpy(state.cpu_scales), source.scales) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py index e88cb3131a..896a6073db 100644 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -2,7 +2,7 @@ import torch -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, ) @@ -21,7 +21,7 @@ def test_quantized_kv_buffers_include_the_full_hold_page(): ), page_size=4, device="cpu", ) - state.initialize_storage(32) + state.allocate_storage(32) allocations.extend(tuple(buffer.shape) for buffer in full_state_test_buffers(state).values()) assert allocations == [ diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 92bc2eeaed..013555fd3e 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -47,17 +47,17 @@ def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( main_linear_state(source).ssm_state.fill_(5) page = torch.empty(source.get_pd_page_bytes(), dtype=torch.uint8, device="cuda") for kind, getter in ( - ("kv", "get_pd_full_att_transfer_layouts"), - ("att_state", "get_pd_bounded_att_or_tail_transfer_layouts"), + ("kv", "get_pd_token_page_layouts"), + ("att_state", "get_pd_runtime_page_layouts"), ): - source.write_to_pd_buffer( + source.write_to_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=getattr(source, getter)(), **({"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 0, "state_token_len": 8}), ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, @@ -72,10 +72,10 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: with monkeypatch.context() as patch: patch.setattr(torch.Tensor, "tolist", reject_host_list) - target.copy_from_other_dp( + target.copy_from_dp( remote_att_manager=source, - src_token_indexes=src_indices, - dst_token_indexes=dst_indices, + src_token_indices=src_indices, + dst_token_indices=dst_indices, src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, @@ -95,11 +95,11 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: assert not tail[:, 2].any() # DP fetch resets the aligned prefix. slot = source.bounded_att_big_page_allocator.alloc(1)[0] - source.save_bounded_att_pages( + source.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=[slot], + page_slot_ids=[slot], ) torch.cuda.synchronize() cpu = torch.empty((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) @@ -107,7 +107,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=cpu, cpu_page_indices=[0], - layouts=source.get_cpu_transfer_layouts(), + layouts=source.get_cpu_page_layouts(), big_page_slot_ids=[slot], ) torch.cuda.synchronize() @@ -122,7 +122,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: device="cuda", ), state_token_len=8, - layouts=target.get_cpu_transfer_layouts(), + layouts=target.get_cpu_page_layouts(), big_page_slot_ids=[dst_slot], req_idx=1, ) diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 485822701c..de556140f8 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -16,8 +16,8 @@ import torch import torch.multiprocessing as mp -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.server.router.model_infer import shared_att_resources_utils from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from unit_tests.common.att_resources.test_model_resources import make_resources @@ -42,10 +42,10 @@ def test_cpu_pd_and_dp_model_layouts_cover_quantized_fields( src, dst = make_quant_state(kind, device), make_quant_state(kind, device) managers = [] for state in (src, dst): - manager = AttResourceManager(request_capacity=2) - manager.set_att_state("main", state) + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", state) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() managers.append(manager) src, dst = [manager.get_full_attention_states("main")[0] for manager in managers] for i, field in enumerate(full_state_test_fields(src, full_state_test_buffers(src)).values()): @@ -70,7 +70,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, - write_to_pd_buffer=True, + write_to_pd_page=True, ) else: copy_page( @@ -80,7 +80,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: dp_rank_in_node=0, tp_world_size=1, token_indices=source_indices, - write_to_pd_buffer=True, + write_to_pd_page=True, ) copy_page( node_att_managers=[managers[1]], @@ -89,7 +89,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: dp_rank_in_node=0, tp_world_size=1, token_indices=pd_indices, - write_to_pd_buffer=False, + write_to_pd_page=False, ) assert_token_fields_equal( source=src, @@ -105,14 +105,14 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: token_indices=source_indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], - layouts=managers[0].get_cpu_transfer_layouts(), + layouts=managers[0].get_cpu_page_layouts(), ) else: managers[0].write_req_to_cpu( token_indices=source_indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], - layouts=managers[0].get_cpu_transfer_layouts(), + layouts=managers[0].get_cpu_page_layouts(), ) if device == "cuda": torch.cuda.synchronize() @@ -120,7 +120,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: cpu_page_buffers=cpu_pages, cpu_page_indices=[0], token_indices=cpu_indices, - layouts=managers[1].get_cpu_transfer_layouts(), + layouts=managers[1].get_cpu_page_layouts(), ) assert_token_fields_equal( source=src, @@ -131,20 +131,20 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): - managers[1].copy_from_other_dp( + managers[1].copy_from_dp( remote_att_manager=managers[0], - src_token_indexes=source_indices, - dst_token_indexes=dp_indices, + src_token_indices=source_indices, + dst_token_indices=dp_indices, src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, dst_req_idx=0, ) return - managers[1].copy_from_other_dp( + managers[1].copy_from_dp( remote_att_manager=managers[0], - src_token_indexes=source_indices, - dst_token_indexes=dp_indices, + src_token_indices=source_indices, + dst_token_indices=dp_indices, src_big_page_ids=[], dst_big_page_ids=[], src_req_idx=0, @@ -224,8 +224,8 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( dst_req=req, src_dp_rank=1, src_req_idx=1, - src_token_indexes=torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=f"cuda:{source_device}"), - dst_token_indexes=torch.tensor([8, 9, 10, 11]), + src_token_indices=torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=f"cuda:{source_device}"), + dst_token_indices=torch.tensor([8, 9, 10, 11]), src_big_page_ids=src_big_page_ids, dst_big_page_ids=dst_big_page_ids, ) @@ -241,8 +241,8 @@ def barrier(group: str): monkeypatch.setattr(module.dist, "barrier", barrier) worker = module.DPKVSharedMoudle.__new__(module.DPKVSharedMoudle) worker.backend = SimpleNamespace( - model=SimpleNamespace(att_resource_manager=target), - dp_att_resource_managers={1: source}, + model=SimpleNamespace(att_state_manager=target), + dp_att_state_managers={1: source}, is_hybrid_att_model=True, is_master_in_dp=False, node_nccl_group="node", @@ -325,11 +325,11 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( pool.ssm_state.zero_() slot = resources.bounded_att_big_page_allocator.alloc(1) with torch.cuda.device(source_device): - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=slot, + page_slot_ids=slot, ) mapping = torch.zeros((3, 16), dtype=torch.int32, device=f"cuda:{source_device}") mapping[1, 0] = 7 @@ -420,7 +420,7 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe task_kind=kind, req_idx=1, state_token_len=8, - write_to_pd_buffer=True, + write_to_pd_page=True, ) copy_page( node_att_managers=[m for m, _ in target], @@ -432,7 +432,7 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe task_kind=kind, req_idx=2, state_token_len=8, - write_to_pd_buffer=False, + write_to_pd_page=False, ) for rank, (manager, full) in enumerate(target): for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): @@ -460,7 +460,7 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: torch.cuda.set_device(0) for kind in ipc_kinds: resources, mapping = shared_att_resources_utils.load_att_resources(0, kind=kind) - assert resources._states_created + assert resources._state_registration_finished assert resources.full_att_allocator is None assert resources.bounded_att_big_page_allocator is None assert resources.bounded_att_small_page_allocator is None @@ -496,5 +496,5 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: except BaseException: connection.send(traceback.format_exc()) finally: - AttResourceManager._get_states.cache_clear() + AttentionStateManager._get_states.cache_clear() connection.close() diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index d069551fc9..acc8e837ac 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -9,15 +9,15 @@ import pytest import torch -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( FullAttentionConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, StaticFP8FullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, create_gqa_state, create_mla_state, ) -from lightllm.common.att_resources.bounded_att import ( +from lightllm.common.attention_state.bounded_att import ( BoundedAttentionState, ConvSSMStateConfig, create_gdn_state, create_kda_state, ) -from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state +from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state @pytest.mark.parametrize( @@ -54,7 +54,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( group_size=16, ) assert state.get_state_type() == expected_type - assert state.resource_id is None + assert state.state_name is None assert state.config.logical_layer_groups == ((3,), (7,)) assert state.config.global_head_num == 8 and state.config.head_num == 2 assert state.config.tp_rank == 2 and state.config.tp_world_size == 4 @@ -67,7 +67,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( if q_scales is not None: assert torch.equal(state.q_scales, q_scales) state.device = torch.device("cpu") - state.initialize_storage(8) + state.allocate_storage(8) assert state.k_storage.shape == state.v_storage.shape == (2, 12, 2, width) assert state.k_storage.untyped_storage().data_ptr() != state.v_storage.untyped_storage().data_ptr() assert state.k_storage.dtype == state.v_storage.dtype == storage_dtype @@ -112,7 +112,7 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( assert state.config.head_dim == row_width and config.head_dim == kv_dim assert state.capacity == 0 and read_kv_storage(state) is None state.device = torch.device("cpu") - state.initialize_storage(8) + state.allocate_storage(8) assert read_kv_storage(state).shape == (2, 12, 1, row_width) if packed: assert state.config.kv_dim == kv_dim @@ -137,7 +137,7 @@ def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Calla mtp_step=2, device="cpu", ) - assert state.config is config and state.resource_id is None + assert state.config is config and state.state_name is None assert state.conv_state.shape == (2, 3, 48, 5) assert state.ssm_state.shape == (2, 9, 2, 8, 8) assert state.mtp_state_index.shape == (3,) @@ -157,7 +157,7 @@ def test_kpool_creation_preserves_pool_and_mtp_dimensions(): mtp_step=2, device="cpu", ) - assert state.config is config and state.resource_id is None + assert state.config is config and state.state_name is None assert state.tail_buffer.shape == (2, 3, 6, 256) assert state.get_state_type() == "kpool_c4_bf16" diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index e84d4e89d6..2af879e03f 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -14,9 +14,9 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.att_resources.full_att import StaticFP8FullAttentionState +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState from unit_tests.common.att_resources.model_helpers import create_full_attention_roles from unit_tests.common.att_resources.test_capacity import definition @@ -41,7 +41,7 @@ def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_arg assert read_kv_storage(state) is None original_config = state.config state.device = torch.device("cpu") - state.initialize_storage(8) + state.allocate_storage(8) assert state.config is original_config assert state.capacity == 8 and read_kv_storage(state).shape[1] == 8 + state.page_size assert not read_kv_storage(state).is_meta @@ -52,14 +52,14 @@ def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_arg assert actual_bytes == token_bytes * (8 + state.page_size) -def test_numbered_drafts_keep_resource_id_order_and_unique_members(): - manager = AttResourceManager(request_capacity=1) +def test_numbered_drafts_keep_state_name_order_and_unique_members(): + manager = AttentionStateManager(request_capacity=1) for role in ("draft35", "draft10", "draft1", "main", "draft0", "draft2"): - manager.set_att_state(role, definition("gqa")) + manager.register_state(role, definition("gqa")) with pytest.raises(AssertionError, match="already bound"): - manager.set_att_state("draft1", definition("gqa")) - manager.finish_state_creation() - assert [state.resource_id.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ + manager.register_state("draft1", definition("gqa")) + manager.finish_state_registration() + assert [state.state_name.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ "draft0", "draft10", "draft1", "draft2", "draft35", "main", ] assert manager.get_full_attention_states("draft1") == (manager.draft1_gqa_bf16_state,) @@ -68,25 +68,25 @@ def test_numbered_drafts_keep_resource_id_order_and_unique_members(): def test_custom_state_type_controls_registered_resource_name(): from unit_tests.common.att_resources.test_full_state import full_state - from lightllm.common.att_resources.full_att import GQAFullAttentionState + from lightllm.common.attention_state.full_att import GQAFullAttentionState class CustomGQAState(GQAFullAttentionState): def get_state_type(self) -> str: return "gqa_custom_fp32" state = full_state(CustomGQAState) - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("draft0", state) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("draft0", state) + manager.finish_state_registration() assert manager.draft0_gqa_custom_fp32_state is state - assert state.resource_id == "draft0_gqa_custom_fp32_state" + assert state.state_name == "draft0_gqa_custom_fp32_state" assert manager.get_full_attention_states("draft0") == (state,) assert not hasattr(state, "state_type") @pytest.mark.parametrize("kind", ["full_meta", "full", "bounded", "tail"]) @pytest.mark.parametrize("model_name", ["main", "draft1"]) -def test_set_att_state_preserves_object_and_storage(kind: str, model_name: str): +def test_register_state_preserves_object_and_storage(kind: str, model_name: str): from unit_tests.common.att_resources.test_bounded_state import make_state from unit_tests.common.att_resources.test_tail_state import make_tail @@ -104,12 +104,12 @@ def test_set_att_state_preserves_object_and_storage(kind: str, model_name: str): state = make_tail() storage = state.tail_buffer config = state.config - manager = AttResourceManager(request_capacity=3) - manager.set_att_state(model_name, state) + manager = AttentionStateManager(request_capacity=3) + manager.register_state(model_name, state) name = f"{model_name}_{state.get_state_type()}_state" assert getattr(manager, name) is state - assert state.resource_id == name - assert state.config is config and not hasattr(config, "resource_id") + assert state.state_name == name + assert state.config is config and not hasattr(config, "state_name") if kind.startswith("full"): assert all(actual is original for actual, original in zip(full_state_test_buffers(state).values(), storage)) elif kind == "bounded": @@ -118,11 +118,11 @@ def test_set_att_state_preserves_object_and_storage(kind: str, model_name: str): assert state.tail_buffer is storage other_role = "draft0" if model_name == "main" else "main" with pytest.raises(AssertionError, match="multiple model roles"): - manager.set_att_state(other_role, state) - assert state.resource_id == name and getattr(manager, name) is state - manager.finish_state_creation() - with pytest.raises(AssertionError, match="creation is already finished"): - manager.set_att_state(other_role, definition("gqa")) + manager.register_state(other_role, state) + assert state.state_name == name and getattr(manager, name) is state + manager.finish_state_registration() + with pytest.raises(AssertionError, match="registration is already finished"): + manager.register_state(other_role, definition("gqa")) def test_model_resource_role_numbers_the_36th_draft_and_rejects_the_37th(): @@ -131,17 +131,17 @@ def test_model_resource_role_numbers_the_36th_draft_and_rejects_the_37th(): draft = LlamaTpPartModel.__new__(LlamaTpPartModel) draft.is_mtp_draft_model = True draft.mtp_previous_draft_models = [object()] * 35 - assert draft.att_resource_role == "draft35" + assert draft.att_state_role == "draft35" draft.mtp_previous_draft_models.append(object()) with pytest.raises(AssertionError, match="at most 36"): - _ = draft.att_resource_role + _ = draft.att_state_role @pytest.mark.parametrize("role", ["draft", "draft36", "draft-1", "draft01"]) -def test_set_att_state_rejects_invalid_draft_numbers(role: str): - manager = AttResourceManager(request_capacity=1) +def test_register_state_rejects_invalid_draft_numbers(role: str): + manager = AttentionStateManager(request_capacity=1) with pytest.raises(AssertionError, match="unknown model role"): - manager.set_att_state(role, definition("gqa")) + manager.register_state(role, definition("gqa")) def allocated_state(kind: str, device: str = "cpu"): @@ -172,58 +172,58 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp if buffer is not None: assert buffer.untyped_storage().data_ptr() != original.untyped_storage().data_ptr() - source = AttResourceManager(request_capacity=2) - source.set_att_state("main", main) - source.set_att_state("draft0", draft) + source = AttentionStateManager(request_capacity=2) + source.register_state("main", main) + source.register_state("draft0", draft) source.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - source.finish_state_creation() - assert getattr(source, main.resource_id) is main - assert getattr(source, draft.resource_id) is draft + source.finish_state_registration() + assert getattr(source, main.state_name) is main + assert getattr(source, draft.state_name) is draft source.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) target_physical = allocated_state(kind, device="cuda") target_main, target_draft = create_full_attention_roles(target_physical, 1) - target = AttResourceManager(request_capacity=2) - target.set_att_state("main", target_main) - target.set_att_state("draft0", target_draft) + target = AttentionStateManager(request_capacity=2) + target.register_state("main", target_main) + target.register_state("draft0", target_draft) target.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - target.finish_state_creation() + target.finish_state_registration() page = torch.zeros((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) source.write_req_to_cpu( token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), cpu_page_buffers=page, cpu_page_indices=[0], - layouts=source.get_cpu_transfer_layouts(), + layouts=source.get_cpu_page_layouts(), ) target.read_req_from_cpu( cpu_page_buffers=page, cpu_page_indices=[0], token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), - layouts=source.get_cpu_transfer_layouts(), + layouts=source.get_cpu_page_layouts(), ) for src, dst in ((main, target_main), (draft, target_draft)): fields = {name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items()} for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): if "token" in field.axes: assert torch.equal(field.tensor[:, :2], fields[name][:, 6:8]) - assert set(source.get_cpu_transfer_layouts()) == {main.resource_id, draft.resource_id} + assert set(source.get_cpu_page_layouts()) == {main.state_name, draft.state_name} @pytest.mark.parametrize("overlap", [False, True]) def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_layers(overlap: bool): main = allocated_state("gqa") draft = allocated_state("gqa") - manager = AttResourceManager(request_capacity=2) - manager.set_att_state("main", main) - manager.set_att_state("draft0", draft) + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", main) + manager.register_state("draft0", draft) manager.full_att_allocator = FullAttKVAllocator(main.capacity, main.page_size, pin_memory=False) template = allocated_state("int8") quant = rebuild_full_state(template, config=replace(template.config, logical_layer_groups=((4,), (6,)))) - manager.set_att_state("main", quant) + manager.register_state("main", quant) if overlap: overlapping = allocated_state("static_head") - manager.set_att_state("main", overlapping) - manager.finish_state_creation() + manager.register_state("main", overlapping) + manager.finish_state_registration() assert manager.main_gqa_bf16_state is main assert manager.draft0_gqa_bf16_state is draft assert manager.main_gqa_int8_group8_bf16_scale_state is quant @@ -232,7 +232,7 @@ def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_la if overlap: assert manager.main_gqa_fp8_static_per_head_state is overlapping expected.append(overlapping) - assert manager.get_full_attention_states("main") == tuple(sorted(expected, key=lambda state: state.resource_id)) + assert manager.get_full_attention_states("main") == tuple(sorted(expected, key=lambda state: state.state_name)) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -246,15 +246,15 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): field.tensor.fill_(5) field.tensor[:, 4:6].zero_() - manager = AttResourceManager(request_capacity=2) + manager = AttentionStateManager(request_capacity=2) for state in (full, quant): - manager.set_att_state("main", state) - manager.finish_state_creation() + manager.register_state("main", state) + manager.finish_state_registration() manager.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): assert (field.tensor[:, 4:6] == 5).all() - layouts = manager.get_cpu_transfer_layouts() + layouts = manager.get_cpu_page_layouts() page = torch.zeros((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) manager.write_req_to_cpu( token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 166d93a64d..28c04506ff 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -13,9 +13,9 @@ import pytest from easydict import EasyDict -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.att_resources.full_att import FullAttentionState +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.server.core.objs.start_args_type import StartArgs from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager from unit_tests.common.att_resources.test_capacity import definition @@ -24,36 +24,36 @@ def test_created_state_queries_keep_the_same_objects_through_storage_initialization(monkeypatch: pytest.MonkeyPatch): args = StartArgs(page_size=4) definitions = [] - manager = AttResourceManager(request_capacity=1) + manager = AttentionStateManager(request_capacity=1) for role, kind in (("main", "gqa"), ("draft0", "int8")): - manager.set_att_state(role, definition(kind)) + manager.register_state(role, definition(kind)) definitions.append(next( state for state in vars(manager).values() - if isinstance(state, FullAttentionState) and state.resource_id.startswith(role + "_") + if isinstance(state, FullAttentionState) and state.state_name.startswith(role + "_") )) for query in ( manager.get_full_attention_states, manager.get_bounded_attention_states, manager.get_attention_tail_states, ): - with pytest.raises(AssertionError, match="State creation"): + with pytest.raises(AssertionError, match="State registration"): query() assert manager._get_states.cache_info().currsize == 0 cached = [] - initialize_storage = AttResourceManager.initialize_storage + allocate_storage = AttentionStateManager.allocate_storage - def check_created_before_allocation(manager: AttResourceManager, *args: object, **kwargs: object): + def check_created_before_allocation(manager: AttentionStateManager, *args: object, **kwargs: object): states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) assert manager.full_att_allocator is None cached.append(states) - initialize_storage(manager, *args, **kwargs) + allocate_storage(manager, *args, **kwargs) - monkeypatch.setattr(AttResourceManager, "initialize_storage", check_created_before_allocation) + monkeypatch.setattr(AttentionStateManager, "allocate_storage", check_created_before_allocation) monkeypatch.setattr( - "lightllm.common.att_resources.att_resource_manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) - manager.finish_state_creation() - manager.initialize_resources( + manager.finish_state_registration() + manager.initialize_storage( args=args, token_capacity=8, mem_fraction=1.0, @@ -90,10 +90,10 @@ def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): def test_finished_state_creation_allows_unallocated_storage(): - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", definition("gqa")) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", definition("gqa")) manager.full_att_allocator = FullAttKVAllocator(0, 4, pin_memory=False) - manager.finish_state_creation() + manager.finish_state_registration() assert read_kv_storage(manager.get_full_attention_states()[0]) is None @@ -109,7 +109,7 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): assert getattr(source, name) is not None assert getattr(exported, name) is None for state, (big_pool, small_pool) in zip(source_states, owner_pools): - copied = getattr(exported, state.resource_id) + copied = getattr(exported, state.state_name) assert copied is not state assert copied.conv_state is state.conv_state assert copied.ssm_state is state.ssm_state @@ -121,10 +121,10 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): assert big_pool.size == source.bounded_att_big_page_allocator.size assert small_pool.size == source.bounded_att_small_page_allocator.size assert exported.get_bounded_attention_states() == tuple( - getattr(exported, state.resource_id) for state in source_states + getattr(exported, state.state_name) for state in source_states ) for state in (*source.get_full_attention_states(), *source.get_attention_tail_states()): - assert getattr(exported, state.resource_id) is state + assert getattr(exported, state.state_name) is state @@ -134,7 +134,7 @@ def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share source = MainDraftTestResourceManager() pd_copy = source.create_pd_ipc_copy() for state in source.get_bounded_attention_states(): - exported = getattr(pd_copy, state.resource_id) + exported = getattr(pd_copy, state.state_name) assert exported.big_pages_pool is None and exported.small_pages_pool is None assert state.big_pages_pool is not None and state.small_pages_pool is not None if not share_pools: @@ -146,7 +146,7 @@ def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share assert getattr(dp_copy, name) is None assert getattr(source, name) is not None for state in source.get_bounded_attention_states(): - exported = getattr(dp_copy, state.resource_id) + exported = getattr(dp_copy, state.state_name) assert exported.big_pages_pool is state.big_pages_pool assert exported.small_pages_pool is state.small_pages_pool assert exported.conv_state is state.conv_state @@ -158,7 +158,7 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): source.get_bounded_attention_states() source.get_attention_tail_states() imported = pickle.loads(pickle.dumps(source.create_pd_ipc_copy())) - assert imported._states_created + assert imported._state_registration_finished assert imported.full_att_allocator is None assert imported.bounded_att_big_page_allocator is None assert imported.bounded_att_small_page_allocator is None @@ -167,7 +167,7 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): imported.get_full_attention_states, imported.get_bounded_attention_states, imported.get_attention_tail_states, ): for state in query(): - name = state.resource_id + name = state.state_name assert state is getattr(imported, name) and state is not getattr(source, name) assert imported._get_states.cache_info().misses == misses + 3 @@ -182,6 +182,6 @@ def test_clearing_query_cache_releases_manager_and_states(): del manager gc.collect() assert manager_ref() is not None and state_ref() is not None - AttResourceManager._get_states.cache_clear() + AttentionStateManager._get_states.cache_clear() gc.collect() assert manager_ref() is None and state_ref() is None diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py index 426f79b781..04d2d92ba3 100644 --- a/unit_tests/common/att_resources/test_tail_state.py +++ b/unit_tests/common/att_resources/test_tail_state.py @@ -3,9 +3,9 @@ import pytest import torch -from lightllm.common.att_resources.att_tail import KPoolAttentionTailState, KPoolStateConfig -from lightllm.common.att_resources.types import ( - AxisSlice, StateTransferLayout, TransferRegion, +from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.attention_state.types import ( + AxisMapping, StateTransferLayout, TransferRegion, ) @@ -27,7 +27,7 @@ def tail_layout(state: KPoolAttentionTailState): shape = (2, state.config.pool_size, 4) nbytes = shape[0] * shape[1] * shape[2] * 4 return StateTransferLayout( - resource_id=state.resource_id, + state_name=state.state_name, tokens_per_page=None, byte_start=0, byte_end=nbytes + 16, @@ -39,18 +39,18 @@ def tail_layout(state: KPoolAttentionTailState): dtype=torch.float32, axes=("layer", "window", "channel"), shape=shape, - axis_slices=( - AxisSlice( + axis_mappings=( + AxisMapping( axis="window", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=shape[1], - src_circular=True, + local_circular=True, ), - AxisSlice( + AxisMapping( axis="channel", - src_start=0, - dst_start=0, + local_start=0, + external_start=0, length=4, ), ), @@ -93,14 +93,14 @@ def test_pd_logical_window_restores_different_mtp_rings( dtype=torch.uint8, device=device, ) - source.write_to_pd_buffer( + source.write_to_pd_page( req_idx=1, pd_page_buffers=pages, pd_page_index=1, state_token_len=length, layout=layout, ) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=pages, req_idx=2, pd_page_index=1, diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index 33e7cb8eec..8919a81d1d 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -15,9 +15,9 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.attention_state.full_att import ( FullAttentionConfig, GQAFullAttentionState, MLAFullAttentionState, @@ -26,15 +26,15 @@ if TYPE_CHECKING: from collections.abc import Sequence from easydict import EasyDict - from lightllm.common.att_resources.bounded_att.base_state import BoundedAttentionState - from lightllm.common.att_resources.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + from lightllm.common.attention_state.full_att.base_state import FullAttentionState def cpu_layout(state: FullAttentionState | BoundedAttentionState): - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("main", state) - manager.finish_state_creation() - return manager.get_cpu_transfer_layouts()[state.resource_id] + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + return manager.get_cpu_page_layouts()[state.state_name] @pytest.mark.parametrize( @@ -65,7 +65,7 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( replicated=replicated, ) reference = state_cls(config, page_size=4, device="cuda") - reference.initialize_storage(4) + reference.allocate_storage(4) for storage_buffer in full_state_test_buffers(reference).values(): storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape(storage_buffer.shape)) fields = {name: field.tensor for name, field in full_state_test_fields(reference, full_state_test_buffers(reference)).items()} @@ -77,7 +77,7 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( assert "head_num" not in vars(restored_config) for rank, (start, end) in enumerate(source_ranges): source = state_cls(replace(config, tp_rank=rank, tp_world_size=len(source_ranges)), page_size=4, device="cuda") - source.initialize_storage(4) + source.allocate_storage(4) layout = cpu_layout(source) assert layout.byte_end == reference_layout.byte_end for name, field in full_state_test_fields(source, full_state_test_buffers(source)).items(): @@ -91,7 +91,7 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( for rank, (start, end) in enumerate(target_ranges): target = state_cls(replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda") - target.initialize_storage(4) + target.allocate_storage(4) layout = cpu_layout(target) assert layout.byte_end == reference_layout.byte_end target.read_from_cpu( @@ -175,10 +175,10 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( ) else: other = rebuild_full_state(source, config=other_config) - manager = AttResourceManager(request_capacity=1) - manager.set_att_state("draft0", source) - manager.set_att_state("draft1", other) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("draft0", source) + manager.register_state("draft1", other) + manager.finish_state_registration() query = manager.get_bounded_attention_states if bounded else manager.get_full_attention_states assert query("draft0") == (source,) assert query("draft1") == (other,) @@ -192,14 +192,14 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( assert page_bytes == page_tokens * 192 assert manager.get_pd_page_bytes() == page_bytes getter = ( - manager.get_pd_bounded_att_or_tail_transfer_layouts if bounded else manager.get_pd_full_att_transfer_layouts + manager.get_pd_runtime_page_layouts if bounded else manager.get_pd_token_page_layouts ) layouts = getter() assert getter() is layouts - assert set(layouts) == {source.resource_id, other.resource_id} + assert set(layouts) == {source.state_name, other.state_name} assert all(layout.tokens_per_page == (None if bounded else page_tokens) for layout in layouts.values()) unused_getter = ( - manager.get_pd_full_att_transfer_layouts if bounded else manager.get_pd_bounded_att_or_tail_transfer_layouts + manager.get_pd_token_page_layouts if bounded else manager.get_pd_runtime_page_layouts ) assert unused_getter() == {} assert manager.get_pd_page_bytes() == page_bytes diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 241e6cdc4f..356cb9a112 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -6,8 +6,8 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, GDNState +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 @@ -357,10 +357,10 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st infer_state.req_manager = SimpleNamespace( HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, ) - infer_state.att_resource_role = role + infer_state.att_state_role = role if req_to_mtp_state_index is not None: request_capacity = req_to_mtp_state_index.numel() - resources = AttResourceManager(request_capacity=request_capacity) + resources = AttentionStateManager(request_capacity=request_capacity) config = ConvSSMStateConfig( logical_layer_ids=(3,), global_k_heads=1, global_v_heads=1, key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, @@ -368,9 +368,9 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st for model_role in ("main", "draft0"): linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") linear.mtp_state_index = req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) - resources.set_att_state(model_role, linear) - resources.finish_state_creation() - infer_state.att_resource_manager = resources + resources.register_state(model_role, linear) + resources.finish_state_registration() + infer_state.att_state_manager = resources state = gdn.LinearAttDecodeAttState( backend=SimpleNamespace( model=model, diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index 381e18e202..c814494d6e 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -7,7 +7,7 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager +from lightllm.common.attention_state import AttentionStateManager from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 from lightllm.models.llama.infer_struct import LlamaInferStateInfo @@ -38,12 +38,12 @@ def test_fp8_backend_selects_calibration_from_logical_layer( state.scales[0].fill_(2) state.scales[1].fill_(5) expected_scale = 2 if storage_layer_index == 0 else 5 - manager = AttResourceManager(request_capacity=1) - manager.set_att_state(role, state) - manager.finish_state_creation() + manager = AttentionStateManager(request_capacity=1) + manager.register_state(role, state) + manager.finish_state_registration() infer = LlamaInferStateInfo() - infer.att_resource_manager = manager - infer.att_resource_role = role + infer.att_state_manager = manager + infer.att_state_role = role infer.b_seq_len = torch.tensor([1], dtype=torch.int32) infer.max_q_seq_len = 1 calls = [] diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 3ae663ccae..09a5e10448 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -52,17 +52,17 @@ def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypat mem_indexes = backend._alloc_req_kv_mem(req, alloc_token_num=4) assert mem_indexes.tolist() == [0, 1, 2, 3] assert req.hold_kv_len == 4 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4] assert context.req_manager.req_to_token_indexs[0, :4].tolist() == [0, 1, 2, 3] req.cur_kv_len = 3 assert backend._alloc_req_kv_mem(req, alloc_token_num=0) is None - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4] req.cur_kv_len = 4 backend._alloc_req_kv_mem(req, alloc_token_num=4) assert req.hold_kv_len == 8 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [4, 4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4, 4] assert context.req_manager.req_to_token_indexs[0, :8].tolist() == list(range(8)) @@ -96,7 +96,7 @@ def __getitem__(self, key): target = _CopyTarget() context = SimpleNamespace( - att_resource_manager=SimpleNamespace(full_att_allocator=_FakeAllocator()), + att_state_manager=SimpleNamespace(full_att_allocator=_FakeAllocator()), req_manager=SimpleNamespace( req_to_token_indexs=_ReqToTokenIndexes(target), ), @@ -189,7 +189,7 @@ def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( context.req_manager.req_to_token_indexs[req.req_idx, :hold_kv_len] = torch.arange( start_index, start_index + hold_kv_len, dtype=torch.int32 ) - context.att_resource_manager.full_att_allocator.next_index = 2 * hold_kv_len + context.att_state_manager.full_att_allocator.next_index = 2 * hold_kv_len backend._filter_not_ready_reqs = lambda req_ids: reqs copy_modes = [] original_alloc_req_kv_mem = backend._alloc_req_kv_mem @@ -209,7 +209,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert prefill_reqs == [reqs[0]] assert decode_reqs == [] assert prefill_tokens == 3 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == expected_alloc_sizes + assert context.att_state_manager.full_att_allocator.alloc_sizes == expected_alloc_sizes assert reqs[0].hold_kv_len == 8 assert not reqs[0].wait_pause assert reqs[1].hold_kv_len == hold_kv_len @@ -245,7 +245,7 @@ def test_decode_scheduler_uses_non_blocking_request_table_copy(monkeypatch): req.get_cur_total_len = lambda: 5 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.att_resource_manager.full_att_allocator.next_index = 4 + context.att_state_manager.full_att_allocator.next_index = 4 backend._filter_not_ready_reqs = lambda req_ids: [req] copy_modes = [] @@ -267,7 +267,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert decode_reqs == [req] assert prefill_tokens == 0 assert req.hold_kv_len == 16 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [12] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [12] assert copy_modes == [True] @@ -307,7 +307,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.get_cur_total_len = lambda: req.cur_kv_len + 1 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :page_size] = torch.arange(page_size, dtype=torch.int32) - context.att_resource_manager.full_att_allocator.next_index = page_size + context.att_state_manager.full_att_allocator.next_index = page_size backend._filter_not_ready_reqs = lambda req_ids: [req] for _ in range(expected_alloc_token_num): @@ -316,7 +316,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.cur_kv_len += 1 assert req.hold_kv_len == page_size + expected_alloc_token_num - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [expected_alloc_token_num] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [expected_alloc_token_num] def test_decode_reserves_mtp_headroom(monkeypatch): @@ -332,7 +332,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): req.args = context.args context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.att_resource_manager.full_att_allocator.next_index = 4 + context.att_state_manager.full_att_allocator.next_index = 4 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == 6 assert alloc_token_num == 8 @@ -343,7 +343,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): assert run_reqs == [req, req, req] assert model_input.b_seq_len.tolist() == [4, 5, 6] assert req.hold_kv_len == 12 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [8] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [8] assert context.req_manager.req_to_token_indexs[0, :12].tolist() == list(range(12)) @@ -360,7 +360,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): req.get_radix_cache_shared_len = lambda: 0 req.args = context.args context.req_manager.req_to_token_indexs[0, :3] = torch.arange(3, dtype=torch.int32) - context.att_resource_manager.full_att_allocator.next_index = 3 + context.att_state_manager.full_att_allocator.next_index = 3 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == alloc_token_num == 6 backend._alloc_req_kv_mem(req, alloc_token_num) @@ -368,7 +368,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): model_input, _ = generic_pre_process.prepare_decode_inputs([req]) assert req.hold_kv_len == 9 - assert context.att_resource_manager.full_att_allocator.alloc_sizes == [6] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [6] assert context.req_manager.req_to_token_indexs[0, :9].tolist() == list(range(9)) @@ -516,7 +516,7 @@ def _make_context(monkeypatch): context = SimpleNamespace( args=SimpleNamespace(page_size=4), req_manager=req_manager, - att_resource_manager=SimpleNamespace(full_att_allocator=mem_manager), + att_state_manager=SimpleNamespace(full_att_allocator=mem_manager), radix_cache=None, ) monkeypatch.setattr(generic_pre_process, "g_infer_context", context) diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index f9293fd1df..c07fd123f0 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -12,8 +12,8 @@ @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" - from lightllm.common.att_resources import AttResourceManager - from lightllm.common.att_resources.allocators import full_att_kv_allocator + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.allocators import full_att_kv_allocator from lightllm.utils.envs_utils import get_unique_server_name monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") @@ -29,7 +29,7 @@ def track_counter(name: str): monkeypatch.setattr(full_att_kv_allocator, "SharedInt", track_counter) yield - AttResourceManager._get_states.cache_clear() + AttentionStateManager._get_states.cache_clear() names = set() for counter in counters: if counter.shm.name not in names: diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 0875d3f9b5..080e4ef99d 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -5,9 +5,9 @@ import pytest import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, PackedMLAAttentionConfig, ) @@ -96,7 +96,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ), page_size=args.page_size, ) - full.initialize_storage(16) + full.allocate_storage(16) req = ReqManager(3, 16) from unit_tests.common.att_resources.model_helpers import create_full_attention_roles @@ -143,22 +143,22 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind=kind, - slot_ids=[slot], + page_slot_ids=[slot], ) torch.cuda.synchronize() linear.conv_state[:, 0].zero_() linear.ssm_state[:, 0].zero_() tail.tail_buffer.normal_() unchanged = tail.tail_buffer[:, [0, 1, 3]].clone() - resources.restore_bounded_att_pages( + resources.restore_bounded_att_runtime_from_pages( req_indices=[2], state_token_lens=[8], page_kind=kind, - slot_ids=[slot], + page_slot_ids=[slot], ) resources.reset_attention_tail_state([2]) torch.cuda.synchronize() diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 05d509feaa..d1a0eaf29f 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -588,11 +588,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): get_layer_view=lambda logical_layer_id: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), ) infer = SimpleNamespace( - att_resource_role="main", - att_resource_manager=SimpleNamespace( - get_att_state=lambda att_type, **kwargs: ( + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda state_type_filter, **kwargs: ( SimpleNamespace(get_layer_view=lambda logical_layer_id: tail) - if "kpool" in att_type else full_att_state + if "kpool" in state_type_filter else full_att_state ), ), is_draft_model=False, diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index ddaf9fc602..657289a1f5 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -170,9 +170,9 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) state = InferStateInfo() state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - state.att_resource_role = "main" - state.att_resource_manager = SimpleNamespace( - get_att_state=lambda **kwargs: SimpleNamespace( + state.att_state_role = "main" + state.att_state_manager = SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace( get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv) ) ) @@ -417,26 +417,26 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): assert draft.req_manager is main.req_manager drafts.append(draft) - assert all(not hasattr(model, "att_resource_manager") for model in (main, *drafts)) + assert all(not hasattr(model, "att_state_manager") for model in (main, *drafts)) assert all(not hasattr(model, "linear_config") for model in (main, *drafts)) from contextlib import nullcontext for model in (main, *drafts): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) - main.prepare_att_resources() - assert main.att_resource_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.is_cuda - assert main.att_resource_manager.main_kpool_c4_bf16_state.tail_buffer.is_cuda + main.prepare_attention_states() + assert main.att_state_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.is_cuda + assert main.att_state_manager.main_kpool_c4_bf16_state.tail_buffer.is_cuda for step, draft in enumerate(drafts): - draft.prepare_att_resources() - assert draft.att_resource_manager is main.att_resource_manager + draft.prepare_attention_states() + assert draft.att_state_manager is main.att_state_manager draft._init_att_backend1() - state = getattr(draft.att_resource_manager, f"draft{step}_mla_packed_bf16_state") - assert state is not main.att_resource_manager.main_mla_packed_bf16_state + state = getattr(draft.att_state_manager, f"draft{step}_mla_packed_bf16_state") + assert state is not main.att_state_manager.main_mla_packed_bf16_state assert not hasattr(draft, "full_att_state") assert state.storage is None assert state.config.logical_layer_groups == ((4 + step,),) - assert state.resource_id == f"draft{step}_mla_packed_bf16_state" + assert state.state_name == f"draft{step}_mla_packed_bf16_state" assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None - assert getattr(draft.att_resource_manager, f"draft{step}_kpool_c4_bf16_state").tail_buffer.is_cuda + assert getattr(draft.att_state_manager, f"draft{step}_kpool_c4_bf16_state").tail_buffer.is_cuda # NoPE state initialization must not require a model-owned RoPE cache. for model in (main, *drafts): diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index d2466bc6d1..fe8e149f3c 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -2,8 +2,8 @@ import pytest import torch -from lightllm.common.att_resources import AttResourceManager -from lightllm.common.att_resources.bounded_att import ConvSSMStateConfig, KDAState +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -194,7 +194,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): uses_dynamic_spec_verify_layout=lambda: dynamic, split_qkv=lambda x: x.split(hidden, -1), ) - resources = AttResourceManager(request_capacity=4) + resources = AttentionStateManager(request_capacity=4) config = ConvSSMStateConfig( logical_layer_ids=(0,), global_k_heads=heads, global_v_heads=heads, key_dim=dim, value_dim=dim, conv_width=3, conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, @@ -204,16 +204,16 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) if model_role == role: linear.conv_state, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) - resources.set_att_state(model_role, linear) - resources.finish_state_creation() + resources.register_state(model_role, linear) + resources.finish_state_registration() manager = SimpleNamespace(HOLD_REQUEST_ID=3) infer = InferStateInfo() infer.batch_size = tokens infer.b_req_idx = torch.tensor(request_rows, device="cuda", dtype=torch.int32) infer.b_mtp_index = torch.tensor(offsets, device="cuda", dtype=torch.int32) infer.req_manager = manager - infer.att_resource_manager = resources - infer.att_resource_role = role + infer.att_state_manager = resources + infer.att_state_role = role weight = SimpleNamespace( get_merged_kda_conv_weight=lambda: conv_weight, linear_A_log=SimpleNamespace(weight=a), diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index c36d9e610d..9415496fef 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -7,10 +7,10 @@ import pytest import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator from unit_tests.common.att_resources.model_helpers import create_full_attention_roles -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, PackedMLAAttentionConfig, ) @@ -80,7 +80,7 @@ def create(tp, target_mtp_step=None): ), page_size=args.page_size, ) - full.initialize_storage(32) + full.allocate_storage(32) main, draft = create_full_attention_roles(full, config.all_layer_num) resources = build_hybrid_resources( config=config, @@ -191,7 +191,7 @@ def test_pd_kv_and_runtime_state_roundtrip( source_buffer = allocate_transfer_buffer(source_resources[0], 1, device="cuda") dest_buffer = allocate_transfer_buffer(dest_resources[0], 1, device="cuda") assert source_buffer.shape == dest_buffer.shape - assert source_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() == dest_resources[0].get_pd_bounded_att_or_tail_transfer_layouts().keys() + assert source_resources[0].get_pd_runtime_page_layouts().keys() == dest_resources[0].get_pd_runtime_page_layouts().keys() for kind in ("kv", "att_state"): copy_page( node_att_managers=source_resources, @@ -203,7 +203,7 @@ def test_pd_kv_and_runtime_state_roundtrip( task_kind=kind, req_idx=src_req, state_token_len=length, - write_to_pd_buffer=True, + write_to_pd_page=True, ) dest_buffer.copy_(source_buffer) copy_page( @@ -216,7 +216,7 @@ def test_pd_kv_and_runtime_state_roundtrip( task_kind=kind, req_idx=dst_req, state_token_len=length, - write_to_pd_buffer=False, + write_to_pd_page=False, ) torch.cuda.synchronize() @@ -284,9 +284,9 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_full_att_transfer_layouts().values()) + full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_token_page_layouts().values()) runtime_bytes = max( - layout.byte_end for layout in mem.resources.get_pd_bounded_att_or_tail_transfer_layouts().values() + layout.byte_end for layout in mem.resources.get_pd_runtime_page_layouts().values() ) assert runtime_bytes > full_bytes assert mem.resources.get_pd_page_bytes() == runtime_bytes @@ -297,8 +297,8 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, @pytest.mark.parametrize("prefix", [1, 3, 4, 5, 11, 12]) def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, prefix): """真实 prefill 导出后，在不同大小的 ring 上验证、拒绝回退并再次验证。""" - from lightllm.common.att_resources.att_tail import KPoolStateConfig, create_kpool_state - from lightllm.common.att_resources.types import StateTransferLayout + from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.common.attention_state.types import StateTransferLayout torch.manual_seed(53) config = KPoolStateConfig( @@ -361,14 +361,14 @@ def prefill(values, destination, tail): regions = source.get_pd_transfer_regions() assert regions == target.get_pd_transfer_regions() layout = StateTransferLayout( - resource_id=source.resource_id, + state_name=source.state_name, tokens_per_page=None, byte_start=0, byte_end=regions[-1].byte_end, regions=regions, ) pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - source.write_to_pd_buffer( + source.write_to_pd_page( req_idx=0, pd_page_buffers=pages, pd_page_index=0, @@ -376,7 +376,7 @@ def prefill(values, destination, tail): layout=layout, ) target.tail_buffer.fill_(17) - target.read_from_pd_buffer( + target.read_from_pd_page( pd_page_buffers=pages, req_idx=0, pd_page_index=0, diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py index 124fe208b4..8f8c191b2c 100644 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -65,7 +65,7 @@ def test_parallel_block_decode_commits_target_hiddens_directly(model_class): model = model_class.__new__(model_class) model._cos_cached = torch.arange(24).view(6, 4) model._sin_cached = model._cos_cached + 100 - model.att_resource_manager = object() + model.att_state_manager = object() model.pre_post_weight = object() target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) @@ -105,7 +105,7 @@ def context_forward(self, hidden, infer_state, layer_weight): assert output.logits.shape == (2, 1) infer_state = observed_states[0] assert not hasattr(infer_state, "full_att_state") - assert infer_state.att_resource_manager is model.att_resource_manager + assert infer_state.att_state_manager is model.att_state_manager assert infer_state.is_draft_model assert infer_state.mem_index is mem_indexes assert torch.equal(infer_state.position_cos, model._cos_cached[[2, 4]]) diff --git a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py index 0928d0bf29..fd75b23e95 100644 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py @@ -79,9 +79,11 @@ def allocate(req, count): _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(1024).view(1, -1)), - att_resource_manager=SimpleNamespace( - get_cpu_transfer_layouts=lambda: {}, - read_req_from_cpu=lambda buffer, pages, tokens, **kwargs: copied_pages.extend(pages), + att_state_manager=SimpleNamespace( + get_cpu_page_layouts=lambda: {}, + read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: copied_pages.extend( + cpu_page_indices + ), ), ), ) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index 6d42fdd052..04e03ee770 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -247,10 +247,10 @@ def test_case9(): def test_flush_cache_requires_released_references_and_returns_allocator_pages(): - from lightllm.common.att_resources.allocators import FullAttKVAllocator + from lightllm.common.attention_state.allocators import FullAttKVAllocator allocator = FullAttKVAllocator(100, pin_memory=False) - tree = RadixCache(100, 0, att_resource_manager=SimpleNamespace(full_att_allocator=allocator)) + tree = RadixCache(100, 0, att_state_manager=SimpleNamespace(full_att_allocator=allocator)) values = allocator.alloc(5) key = torch.arange(5) tree.insert(key[:3], values[:3]) @@ -273,24 +273,24 @@ def test_hybrid_flush_releases_shared_kv_and_state_slots(monkeypatch): from unit_tests.common.att_resources.test_manager_helpers import composition monkeypatch.setattr( - "lightllm.common.att_resources.bounded_att.utils.get_env_start_args", + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), ) resources = composition(device="cuda", with_tail=False) allocator = resources.full_att_allocator big = resources.bounded_att_big_page_allocator.alloc(1)[0] small = resources.bounded_att_small_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=[big], + page_slot_ids=[big], ) - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[1], state_token_lens=[12], page_kind="small", - slot_ids=[small], + page_slot_ids=[small], ) values = allocator.alloc(12) tree = HybridAttPagedRadixCache(12, 106, 4, 2, resources) @@ -346,7 +346,7 @@ def test_page_size_must_match_allocator(): mem_manager = SimpleNamespace(page_size=8) with pytest.raises(ValueError, match="must match allocator page_size 8"): - RadixCache(100, 100, att_resource_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) + RadixCache(100, 100, att_state_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) @pytest.mark.parametrize("update_refs", [False, True]) @@ -463,7 +463,7 @@ def test_hybrid_match_refs_only_valid_prefix_after_tail_state_eviction(): rank_in_node=101, hash_page_size=4, big_page_num=2, - att_resource_manager=hybrid_resources(freed_small=freed_state_ids), + att_state_manager=hybrid_resources(freed_small=freed_state_ids), ) key = torch.arange(12, dtype=torch.int64) values = torch.arange(100, 112, dtype=torch.int64) @@ -503,7 +503,7 @@ def hybrid_tree(): rank_in_node=102, hash_page_size=4, big_page_num=2, - att_resource_manager=hybrid_resources(freed_big=freed_big_states), + att_state_manager=hybrid_resources(freed_big=freed_big_states), ) return tree, freed_big_states @@ -569,7 +569,7 @@ def test_hybrid_match_falls_back_along_matched_branch(monkeypatch, small_state_v rank_in_node=104, hash_page_size=64, big_page_num=4, - att_resource_manager=hybrid_resources(freed_big=freed_big_states), + att_state_manager=hybrid_resources(freed_big=freed_big_states), ) # The older request leaves a small-page state at 384, before the image. old_values = torch.arange(1000, 1384) @@ -676,7 +676,7 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) context.radix_cache = tree - context.att_resource_manager = tree.att_resource_manager + context.att_state_manager = tree.att_state_manager context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(100, 132).view(1, -1)) req = InferReq.__new__(InferReq) req.req_idx = req.cur_output_len = 0 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py index 393d589e97..2a90c7f943 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -8,7 +8,7 @@ import torch from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 -from lightllm.common.att_resources.allocators import StatePageAllocator +from lightllm.common.attention_state.allocators import StatePageAllocator from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans @@ -89,8 +89,8 @@ def test_build_reserves_only_new_big_pages_before_kv( task, = tasks assert task.dst_req is req and task.src_req_idx == 1 assert task.src_dp_rank == 1 - assert task.src_token_indexes.tolist() == list(range(40, 32 + src_kv_len)) - assert len(task.dst_token_indexes) == src_kv_len - 8 + assert task.src_token_indices.tolist() == list(range(40, 32 + src_kv_len)) + assert len(task.dst_token_indices) == src_kv_len - 8 assert allocations == [((src_kv_len + 3) // 4) * 4 - 8] assert req.cur_kv_len == 8 and req.hybrid_len_to_big_page_id == {8: 77} if src_kv_len == 26: @@ -128,15 +128,15 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( barriers = [] manager = SimpleNamespace( bounded_att_big_page_allocator=SimpleNamespace(alloc=allocate_pages), - copy_from_other_dp=copy_resources, + copy_from_dp=copy_resources, ) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = SimpleNamespace( is_hybrid_att_model=hybrid, is_master_in_dp=True, node_nccl_group="node", - model=SimpleNamespace(att_resource_manager=manager), - dp_att_resource_managers={}, + model=SimpleNamespace(att_state_manager=manager), + dp_att_state_managers={}, dp_req_to_token_indices={}, _alloc_req_kv_mem=allocate_tokens, ) @@ -198,7 +198,7 @@ def allocate(cur_req: SimpleNamespace, count: int) -> torch.Tensor: dp_world_size=2, rank_in_dp=1, node_nccl_group=None, - model=SimpleNamespace(att_resource_manager=SimpleNamespace(bounded_att_big_page_allocator=allocator)), + model=SimpleNamespace(att_state_manager=SimpleNamespace(bounded_att_big_page_allocator=allocator)), dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, _alloc_req_kv_mem=allocate, ) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index dd7f7acd2a..5f71ea7dad 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -137,7 +137,7 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): context = InferenceContext() context.is_hybrid_att_model = True context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) - context.att_resource_manager = SimpleNamespace( + context.att_state_manager = SimpleNamespace( bounded_att_small_page_allocator=SimpleNamespace(free=freed_small_pages.extend), bounded_att_big_page_allocator=SimpleNamespace(free=freed_big_pages.extend), ) @@ -182,7 +182,7 @@ def alloc_req_kv_mem(req, alloc_token_num): return mem_indexes resources = SimpleNamespace( - get_cpu_transfer_layouts=lambda: {}, + get_cpu_page_layouts=lambda: {}, read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded_indexes.extend( token_indices.tolist() ), @@ -192,7 +192,7 @@ def alloc_req_kv_mem(req, alloc_token_num): is_hybrid_att_model=False, radix_cache=None, model=SimpleNamespace( - att_resource_manager=resources, + att_state_manager=resources, req_manager=SimpleNamespace(req_to_token_indexs=table), ), _alloc_req_kv_mem=alloc_req_kv_mem, diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py index 97f1a7e0fb..a3c2bb5795 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -144,8 +144,8 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat assert req.hold_kv_len == 12 assert len(tasks) == 1 assert tasks[0].src_dp_rank == 1 - assert tasks[0].dst_token_indexes.tolist() == list(range(4, 10)) - assert len(tasks[0].src_token_indexes) == 6 + assert tasks[0].dst_token_indices.tolist() == list(range(4, 10)) + assert len(tasks[0].src_token_indices) == 6 @pytest.mark.parametrize( @@ -192,7 +192,7 @@ def load_storage(rank, kind): node_world_size=node_size, rank_in_node=rank_in_node, model=SimpleNamespace( - att_resource_manager=local_manager, + att_state_manager=local_manager, req_manager=SimpleNamespace(req_to_token_indexs=tables[rank_in_node]), ), ) @@ -205,10 +205,10 @@ def load_storage(rank, kind): remote_ranks = [rank for rank in expected_ranks if rank != rank_in_node] assert imported_ranks == registered_ranks == remote_ranks - assert list(backend.dp_att_resource_managers) == list(range(len(expected_ranks))) + assert list(backend.dp_att_state_managers) == list(range(len(expected_ranks))) assert list(backend.dp_req_to_token_indices) == list(range(len(expected_ranks))) for dp_rank, rank in enumerate(expected_ranks): - assert backend.dp_att_resource_managers[dp_rank] is managers[rank] + assert backend.dp_att_state_managers[dp_rank] is managers[rank] assert backend.dp_req_to_token_indices[dp_rank] is tables[rank] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py index e73d8b0a2a..01addc80f9 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -94,7 +94,7 @@ def capture_copy(**kwargs: object) -> None: device == "cpu" and not pinned and non_blocking and stream == module.copy_cuda_stream for device, pinned, non_blocking, stream in uploads ) - assert all(kwargs["write_to_pd_buffer"] is prefill for kwargs, _ in copies) + assert all(kwargs["write_to_pd_page"] is prefill for kwargs, _ in copies) assert all(indices is None or indices.is_cuda and indices.dtype == torch.long for _, indices in copies) diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index ed5fe0585b..b04d9ddc54 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -184,13 +184,13 @@ def make_runtime( draft.tp_world_size_, draft.tp_rank_ = tp, rank prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) drafts.append(draft) - model.att_resource_manager.finish_state_creation() - model.att_resource_manager.load_full_att_calibration(args=args) - model.att_resource_manager.initialize_resources( + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction ) assert not hasattr(model, "mem_manager") - managers.append(model.att_resource_manager) + managers.append(model.att_state_manager) return meta, managers @@ -222,8 +222,8 @@ def test_cpu_service_describes_main_and_draft_as_distinct_resources( layer for group in draft.config.logical_layer_groups for layer in group ) assert main.capacity == draft.capacity == resources.full_att_allocator.size - assert {main.resource_id, draft.resource_id} <= resources.get_cpu_transfer_layouts().keys() - layout_end = max(layout.byte_end for layout in resources.get_cpu_transfer_layouts().values()) + assert {main.state_name, draft.state_name} <= resources.get_cpu_page_layouts().keys() + layout_end = max(layout.byte_end for layout in resources.get_cpu_page_layouts().values()) assert (layout_end + 15) // 16 * 16 == meta.page_bytes assert resources.get_cpu_page_bytes() == meta.page_bytes @@ -251,7 +251,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( # No byte-size formulas duplicated in the test: check the service's actual # meta-device descriptors against every real model's layout. for resources in managers: - layout = next(iter(resources.get_cpu_transfer_layouts().values())) + layout = next(iter(resources.get_cpu_page_layouts().values())) assert 0 <= layout.byte_start < layout.byte_end <= meta.page_bytes assert meta.page_bytes == resources.get_cpu_page_bytes() meta.page_num = 3 @@ -269,11 +269,11 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( linear.conv_state.fill_(3) linear.ssm_state.fill_(5) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.save_bounded_att_pages( + resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], page_kind="big", - slot_ids=[slot], + page_slot_ids=[slot], ) kwargs = dict(big_page_slot_ids=[slot]) torch.cuda.synchronize() @@ -288,7 +288,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=buffer, cpu_page_indices=[2], - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), **kwargs, ) torch.cuda.synchronize() @@ -313,7 +313,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( device="cuda", ), state_token_len=8, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), **kwargs, ) if resources.bounded_att_big_page_allocator is not None: @@ -357,7 +357,7 @@ def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( write_to_cpu = resources.write_req_to_cpu def capture_write(*args: object, **kwargs: object) -> None: - write_batches.append((kwargs["cpu_page_indices"], kwargs["cpu_page_readies"])) + write_batches.append((kwargs["cpu_page_indices"], kwargs["cpu_page_ready_flags"])) write_to_cpu(*args, **kwargs) monkeypatch.setattr(resources, "write_req_to_cpu", capture_write) @@ -399,7 +399,7 @@ def capture_write(*args: object, **kwargs: object) -> None: is_hybrid_att_model=hybrid, dp_world_size=1, model=SimpleNamespace( - att_resource_manager=resources, + att_state_manager=resources, req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(16, device="cuda").view(1, -1)), ), ) @@ -441,7 +441,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: cpu_page_indices=[pages[1]], token_indices=torch.arange(8, dtype=torch.long, device="cuda"), state_token_len=8, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), **kwargs, ) for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): @@ -481,7 +481,7 @@ def test_disk_preserves_resource_pages( token_indices=torch.arange(16, dtype=torch.long, device="cuda"), cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=pages, - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), ) torch.cuda.synchronize() expected = client.cpu_kv_cache_tensor.clone() @@ -501,6 +501,6 @@ def test_disk_preserves_resource_pages( cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=pages, token_indices=torch.arange(16, dtype=torch.long, device="cuda"), - layouts=resources.get_cpu_transfer_layouts(), + layouts=resources.get_cpu_page_layouts(), ) assert any((tmp_path / "disk").glob("cache_file*")) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 6cd5713c57..0fcff8d09f 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -6,9 +6,9 @@ import pytest import torch -from lightllm.common.att_resources.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.att_resources.full_att import ( +from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, FullAttentionConfig, ) @@ -71,7 +71,7 @@ def linear_cache(monkeypatch, tmp_path, request): ), page_size=args.page_size, ) - full.initialize_storage(384) + full.allocate_storage(384) requests = ReqManager(2, 512) resources = build_hybrid_resources( config=config, @@ -82,7 +82,7 @@ def linear_cache(monkeypatch, tmp_path, request): big_page_capacity=5, small_page_capacity=2, ) - monkeypatch.setattr(g_infer_context, "att_resource_manager", resources) + monkeypatch.setattr(g_infer_context, "att_state_manager", resources) monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(g_infer_context, "req_manager", requests) monkeypatch.setattr( @@ -172,7 +172,7 @@ def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, - model=SimpleNamespace(req_manager=requests, att_resource_manager=resources), + model=SimpleNamespace(req_manager=requests, att_state_manager=resources), ) monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) task = module.write_req_to_cpu(req, torch.cuda.Stream()) @@ -201,7 +201,7 @@ def allocate(req, count): module.backend = SimpleNamespace( is_master_in_dp=True, _alloc_req_kv_mem=allocate, - model=SimpleNamespace(req_manager=requests, att_resource_manager=resources), + model=SimpleNamespace(req_manager=requests, att_state_manager=resources), is_hybrid_att_model=True, ) module.init_sync_group = None @@ -228,7 +228,7 @@ def allocate(req, count): rank_in_node=103, hash_page_size=64, big_page_num=2, - att_resource_manager=resources, + att_state_manager=resources, ) cache_len = max(req.hybrid_len_to_big_page_id) tree.insert( @@ -305,7 +305,7 @@ def write_to_cpu( is_master_in_dp=True, is_hybrid_att_model=True, dp_world_size=1, - model=SimpleNamespace(att_resource_manager=resources, req_manager=requests), + model=SimpleNamespace(att_state_manager=resources, req_manager=requests), ) module.cpu_cache_client = SimpleNamespace( lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 904cad0f17..d75d10e776 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -39,7 +39,7 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): from unit_tests.common.att_resources.test_manager_helpers import composition, init_source monkeypatch.setattr( - "lightllm.common.att_resources.bounded_att.utils.get_env_start_args", + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), ) resources = composition(device="cuda") @@ -47,18 +47,18 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): tail = resources.main_kpool_c4_fp32_state resources.reset_bounded_att_state([1]) assert (tail.tail_buffer[:, 1] == 9).all() - resources.restore_bounded_att_pages( + resources.restore_bounded_att_runtime_from_pages( req_indices=[1], state_token_lens=[8], page_kind="big", - slot_ids=[slots[-1]], + page_slot_ids=[slots[-1]], ) for state in resources.get_bounded_attention_states(): state.mtp_state_index.fill_(2) monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(g_infer_context, "radix_cache", None) - monkeypatch.setattr(g_infer_context, "att_resource_manager", resources) + monkeypatch.setattr(g_infer_context, "att_state_manager", resources) req = InferReq.__new__(InferReq) req.args = SimpleNamespace( linear_att_hash_page_size=4, @@ -136,10 +136,10 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( context.args = req.args context.is_hybrid_att_model = True context.radix_cache = SimpleNamespace(free_one_small_page_buffer=lambda: None) - context.att_resource_manager = SimpleNamespace( + context.att_state_manager = SimpleNamespace( bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: [next(slots)]), bounded_att_small_page_allocator=SimpleNamespace(alloc=lambda count: [7]), - save_bounded_att_pages=lambda req_indices, state_token_lens, page_kind, slot_ids: ( + save_bounded_att_runtime_to_pages=lambda req_indices, state_token_lens, page_kind, page_slot_ids: ( big_snapshots if page_kind == "big" else small_snapshots ).extend(state_token_lens), ) @@ -333,8 +333,8 @@ def allocate(req, count): _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), - att_resource_manager=SimpleNamespace( - get_cpu_transfer_layouts=lambda: {}, + att_state_manager=SimpleNamespace( + get_cpu_page_layouts=lambda: {}, bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded.append( req.hold_kv_len diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py index 3ffa28da55..a263c9bed9 100644 --- a/unit_tests/server/router/test_kv_capacity.py +++ b/unit_tests/server/router/test_kv_capacity.py @@ -4,8 +4,8 @@ import pytest -from lightllm.common.att_resources.allocators import FullAttKVAllocator -from lightllm.common.att_resources.allocators import full_att_kv_allocator +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import full_att_kv_allocator pytestmark = pytest.mark.usefixtures("kv_allocator_environment") diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 56433d1b3b..2cde8ec635 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -39,7 +39,7 @@ def make_executor(page_size: int = 4): req_manager = SimpleNamespace( req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), ) - model = SimpleNamespace(att_resource_manager=SimpleNamespace(full_att_allocator=allocator), req_manager=req_manager) + model = SimpleNamespace(att_state_manager=SimpleNamespace(full_att_allocator=allocator), req_manager=req_manager) args = SimpleNamespace(mtp_mode=None, mtp_step=0) executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) return executor, allocator, req_manager From fdf649d405ae9a9282ee338ec6ecec3f165e1e56 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 07:39:55 +0000 Subject: [PATCH 25/56] refactor: share attention transfer view preparation helpers --- docs/cache_manager_design.md | 8 +- .../att_tail/kpool_transfer.py | 29 +-- .../bounded_att/copy_kernel.py | 90 ++------- .../attention_state/full_att/dp_transfer.py | 31 +-- .../attention_state/full_att/gqa_transfer.py | 31 +-- .../full_att/group_quantized_transfer.py | 31 +-- .../attention_state/full_att/mla_transfer.py | 31 +-- .../full_att/packed_mla_transfer.py | 31 +-- .../full_att/static_fp8_transfer.py | 31 +-- lightllm/common/attention_state/utils.py | 140 ++++++++++++- .../att_resources/test_bounded_copy_kernel.py | 3 +- .../att_resources/test_full_att_cpu_kernel.py | 81 ++++++++ .../test_full_att_dp_transfer.py | 36 ++++ .../att_resources/test_transfer_views.py | 189 ++++++++++++++++++ 14 files changed, 523 insertions(+), 239 deletions(-) create mode 100644 unit_tests/common/att_resources/test_transfer_views.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 0f90a09e0d..87abcaa68a 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -302,7 +302,7 @@ lightllm/common/ │ └── basemodel_utils.py # 配置归一化、KV 选项与 full/linear State 定义构建 ├── attention_state/ │ ├── __init__.py -│ ├── utils.py # State 角色与精度命名 +│ ├── utils.py # State 命名、传输视图尾维合并及搬运 dtype 选择 │ ├── manager.py # AttentionStateManager：State 装配、容量与存储初始化、缓存和传输接口 │ ├── types.py # TensorField / StateTransferLayout 等 │ ├── allocators/ # 管理器直接持有，同一分配域共享实例 @@ -362,7 +362,7 @@ lightllm/utils/ `ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `state_name`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 -KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个私有流程函数顺序执行：`_view_min_dimensions` 合并连续尾维，`_view_max_dtype` 整理最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧不连续就停止，保留前面的维度及单元素轴，再按行宽、stride、storage_offset 及整批地址对齐条件创建最大 8 字节的整数 view，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 +KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个阶段顺序执行：公共 `utils.merge_contiguous_tail_dims` 合并连续尾维，公共 `utils.view_max_copy_dtype` 为整批选择最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。公共函数使用 PyTorch `_collapse_view_helper(..., must_be_valid=None)` 判断两端均可合并的尾部区间，随后调用 `view` 创建整批视图；只读取布局元数据，不复制数据。PyTorch 2.11 的 helper 对带 padding 的单元素轴可能过于宽松，最终创建视图时仅捕获 `view` 的 `RuntimeError` 并退回较小范围，以防错误跨越存储间隙；不使用可能复制数据的 `reshape`。`start_dim` 限制允许参与合并的最前轴，例如 `[layer, window, head, dim]` 使用 `start_dim=2` 保留环形 window，只合并 head/dim；`start_dim=-2` 按两侧末尾两维对齐，允许前缀维数、token 容量不同。默认直接减少维数；`keepdim=True` 将被合并的前面各轴保留为 1，使元素集中到末轴，供固定维数的 token/page kernel 使用。调用前通过缓存的 `_check_collapse_view_version` 断言，只允许已验证的 PyTorch 2.11.x，`end` 包含在合并区间内。该函数属于 PyTorch 内部 API，历史版本的参数和区间语义不同；升级 PyTorch 时须先核对签名、合并区间并通过 padding/环形边界回归，再调整允许的版本范围。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧无法合并就停止，保留前面的维度；尾部单元素轴由 PyTorch 按布局处理，再按行宽选择候选类型，通过 `view` 校验布局并尝试创建最大 8 字节的整数视图，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 默认参数在 H200 上固定 grid=24 联合调优，覆盖 Conv BF16 / SSM FP32、三种规模和 CUDA↔CUDA / CUDA↔pinned CPU / pinned CPU↔pinned CPU，共 24 个场景。第一轮比较 BLOCK_ELEMENTS=512/1024/2048/4096 与 warps=1/2/4/8，并保留 256/8/1 对照；第二轮为入选组合比较 num_stages=1/2/3/4，两轮共测试 29 个不同组合。每个候选均校验数据，预先创建指针表，以随机顺序进行 CUDA Graph 计时并取中位数。按场景等权的速度几何平均选择 1024/8/1；相比旧默认 grid=16、BLOCK_ELEMENTS=256、warps=8、stages=1，第二轮综合速度约为 1.58 倍。该组合并非每项最快，小块场景最多增加约 1.35 μs（Conv CPU→GPU 约 2.42→3.77 μs）。`grid_num`、`num_warps`、`num_stages` 可由调用方覆盖；流水深度用于 `tl.range` 的循环属性。 @@ -382,7 +382,7 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，其他外层 stride 和 Storage 偏移交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP/同卡 token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 @@ -419,7 +419,7 @@ manager.register_state("main", state) # 全部 State 准备完成后，由 manager 统一规划容量并分配存储。 ``` -根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 +根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名及传输视图的尾维合并、整批最大搬运 dtype 选择，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttentionStateManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 diff --git a/lightllm/common/attention_state/att_tail/kpool_transfer.py b/lightllm/common/attention_state/att_tail/kpool_transfer.py index 490774cc53..235246f0d0 100644 --- a/lightllm/common/attention_state/att_tail/kpool_transfer.py +++ b/lightllm/common/attention_state/att_tail/kpool_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import AxisMapping, StateTransferLayout +from ..utils import view_max_copy_dtype # ======================================================================== @@ -80,33 +81,7 @@ def transfer_kpool_pages( # 阶段三：只扩大连续 channel 的搬运单位，layer/window 维度和环形语义保持不变。 # 例如 BF16 的 [layers, ring_size, 256] 按 INT64 搬运时变为 [layers, ring_size, 64]。 - local_tail = tail - if local_tail.stride(-1) == external_tail.stride(-1) == 1: - element_bytes = local_tail.element_size() - for unit_bytes, unit_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if unit_bytes <= element_bytes: - break - elements_per_unit = unit_bytes // element_bytes - # 每个 channel 行都必须能完整分组，不能跨 window 拼接搬运单元。 - if channel_count % elements_per_unit != 0: - continue - # Tensor 起点和 storage_offset 都须满足宽类型 view 的对齐要求。 - if ( - local_tail.data_ptr() % unit_bytes != 0 - or external_tail.data_ptr() % unit_bytes != 0 - or local_tail.storage_offset() % elements_per_unit != 0 - or external_tail.storage_offset() % elements_per_unit != 0 - ): - continue - # layer/window 的步长必须能换算为新类型的整数元素数。 - if ( - any(stride % elements_per_unit != 0 for stride in local_tail.stride()[:-1]) - or any(stride % elements_per_unit != 0 for stride in external_tail.stride()[:-1]) - ): - continue - local_tail = local_tail.view(unit_dtype) - external_tail = external_tail.view(unit_dtype) - break + local_tail, external_tail = view_max_copy_dtype(copies=[(tail, external_tail)])[0] # 阶段四：正反方向都只处理逻辑窗口，窗口外的 MTP 槽位不参与读写。 # grid_num 是发射上限，大窗口由 kernel 的跳步循环处理。 diff --git a/lightllm/common/attention_state/bounded_att/copy_kernel.py b/lightllm/common/attention_state/bounded_att/copy_kernel.py index 580c526f5a..b4ec296fda 100644 --- a/lightllm/common/attention_state/bounded_att/copy_kernel.py +++ b/lightllm/common/attention_state/bounded_att/copy_kernel.py @@ -4,6 +4,8 @@ import triton import triton.language as tl +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype + # ======================================================================== # 一、聚合入口：只上传指针表，shape/stride 作为常量 @@ -31,10 +33,19 @@ def copy_bounded_att_tensors( if not copies: return assert grid_num > 0, "copy grid must be positive" + # 每批使用相同 shape/dtype 和同侧 stride；只核对首尾，保留两侧的布局差异。 + source, destination = copies[0] + last_source, last_destination = copies[-1] + assert source.shape == destination.shape == last_source.shape == last_destination.shape, "copy shapes must match" + assert source.dtype == destination.dtype == last_source.dtype == last_destination.dtype, "copy dtypes must match" + assert 1 <= source.ndim <= 6, "copy tensors must have one to six dimensions" + assert source.numel() > 0, "copy tensors must be nonempty" + assert source.stride() == last_source.stride(), "source strides must match across the batch" + assert destination.stride() == last_destination.stride(), "destination strides must match across the batch" # 阶段一：只合并两端共同连续的尾部维度，保留前面的布局。 - copies = _view_min_dimensions(copies=copies) + copies = merge_contiguous_tail_dims(copies=copies) # 阶段二：按整批地址和行布局选择最宽的搬运单元，不转换数值。 - copies = _view_max_dtype(copies=copies) + copies = view_max_copy_dtype(copies=copies) # 阶段三：上传两张指针表，按最终视图维数发射聚合拷贝算子。 _launch_copy_kernel( copies=copies, @@ -45,83 +56,10 @@ def copy_bounded_att_tensors( # ======================================================================== -# 二、视图整理与算子发射：降维 → 扩大搬运单元 → kernel +# 二、聚合算子发射：上传指针与共享布局 # ======================================================================== -def _view_min_dimensions( - copies: list[tuple[torch.Tensor, torch.Tensor]], -) -> list[tuple[torch.Tensor, torch.Tensor]]: - """按首对布局合并连续尾维，遇到任一侧不连续就停止。""" - # 每批复制使用同一种 shape；由调用方分别收集 Conv、SSM 或同一传输区域。 - source, destination = copies[0] - last_source, last_destination = copies[-1] - # 整批布局由调用方保证一致，只检查首尾两对 Tensor。 - assert source.shape == destination.shape == last_source.shape == last_destination.shape, "copy shapes must match" - assert source.dtype == destination.dtype == last_source.dtype == last_destination.dtype, "copy dtypes must match" - ndim = source.ndim - assert 1 <= ndim <= 6, "copy tensors must have one to six dimensions" - assert source.numel() > 0, "copy tensors must be nonempty" - shape = tuple(source.shape) - # 同一侧所有视图共用 stride；只核对首尾，源和目标之间允许不同布局。 - src_stride = source.stride() - dst_stride = destination.stride() - assert src_stride == last_source.stride(), "source strides must match across the batch" - assert dst_stride == last_destination.stride(), "destination strides must match across the batch" - # 从最后一维向前扩展连续区间；任一侧有间隙就停止，不再处理前面的轴。 - # 例如 (2, 3, 4, 5) 的末轴带 padding 时保持原形状，不额外合并前面的 (2, 3, 4)。 - first_tail_axis = ndim - 1 - tail_elements = shape[-1] - while first_tail_axis > 0: - outer_axis = first_tail_axis - 1 - if ( - src_stride[outer_axis] != src_stride[first_tail_axis] * shape[first_tail_axis] - or dst_stride[outer_axis] != dst_stride[first_tail_axis] * shape[first_tail_axis] - ): - break - tail_elements *= shape[outer_axis] - first_tail_axis = outer_axis - if first_tail_axis == ndim - 1: - return copies - transfer_shape = shape[:first_tail_axis] + (tail_elements,) - return [(source.view(transfer_shape), destination.view(transfer_shape)) for source, destination in copies] - - -def _view_max_dtype( - copies: list[tuple[torch.Tensor, torch.Tensor]], -) -> list[tuple[torch.Tensor, torch.Tensor]]: - """将整批 view 为最大可用整数宽度，保留原始位模式和存储间隙。""" - source, destination = copies[0] - if source.stride(-1) != 1 or destination.stride(-1) != 1: - return copies - - # 每个切片的起点可能不同；地址和 storage_offset 都须支持相同的宽类型 view。 - # 累计低位即可检查整批对齐，不能仅检查首尾，否则会漏掉中间错位的视图。 - pointer_alignment_bits = 0 - offset_alignment_bits = 0 - for src_tensor, dst_tensor in copies: - pointer_alignment_bits |= src_tensor.data_ptr() | dst_tensor.data_ptr() - offset_alignment_bits |= src_tensor.storage_offset() | dst_tensor.storage_offset() - - # 最内维连续时，将相邻元素作为一个更宽的整数搬运；优先 8 字节，其次 4/2 字节。 - # 行宽、外层 stride、storage_offset 及整批地址都须对齐，避免跨行 padding 或错位读取。 - element_bytes = source.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - pointer_alignment_bits % transfer_bytes != 0 - or source.shape[-1] % elements_per_unit != 0 - or offset_alignment_bits % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in source.stride()[:-1]) - or any(stride % elements_per_unit for stride in destination.stride()[:-1]) - ): - continue - return [(src.view(transfer_dtype), dst.view(transfer_dtype)) for src, dst in copies] - return copies - - def _launch_copy_kernel( copies: list[tuple[torch.Tensor, torch.Tensor]], grid_num: int, diff --git a/lightllm/common/attention_state/full_att/dp_transfer.py b/lightllm/common/attention_state/full_att/dp_transfer.py index c9fc88d752..8ded336af5 100644 --- a/lightllm/common/attention_state/full_att/dp_transfer.py +++ b/lightllm/common/attention_state/full_att/dp_transfer.py @@ -4,6 +4,8 @@ import triton import triton.language as tl +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype + def copy_gqa_tokens( src_k_storage: torch.Tensor, @@ -73,28 +75,17 @@ def _copy_token_storages( if token_count == 0: return - # 阶段一：只构造整数 view。字段仍为 [layer, token_slot, head, width]； - # 各字段可有不同 dtype/宽度，两侧 token 容量和实际 stride 也可以不同。 - # 连续末轴优先按 8/4/2 字节搬运，保留 FP8、NaN、动态 scale 等原始位模式。 + # 阶段一：先合并两侧连续的 head/width，保留各自 layer/token 容量和 stride。 + # head 轴保留为 1，沿用四维 token kernel；不连续时保持原布局。 + # 再按 8/4/2 字节搬运，保留 FP8、NaN、动态 scale 等原始位模式。 sources, destinations = [], [] for source, destination in zip(src_storages, dst_storages): - if source.stride(-1) == destination.stride(-1) == 1: - element_bytes = source.element_size() - for unit_bytes, unit_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if unit_bytes <= element_bytes: - break - ratio = unit_bytes // element_bytes - if ( - source.shape[-1] % ratio - or source.data_ptr() % unit_bytes or destination.data_ptr() % unit_bytes - or source.storage_offset() % ratio or destination.storage_offset() % ratio - or any(stride % ratio for stride in source.stride()[:-1]) - or any(stride % ratio for stride in destination.stride()[:-1]) - ): - continue - source = source.view(unit_dtype) - destination = destination.view(unit_dtype) - break + source, destination = merge_contiguous_tail_dims( + copies=[(source, destination)], + start_dim=-2, + keepdim=True, + )[0] + source, destination = view_max_copy_dtype(copies=[(source, destination)])[0] sources.append(source) destinations.append(destination) diff --git a/lightllm/common/attention_state/full_att/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa_transfer.py index 1782fb2672..4c6d1bb8a8 100644 --- a/lightllm/common/attention_state/full_att/gqa_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import StateTransferLayout +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== @@ -115,28 +116,14 @@ def transfer_gqa_pages( # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: - # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 - # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 - # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 - if local.stride(-1) == external.stride(-1) == 1: - element_bytes = local.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - local.shape[-1] % elements_per_unit != 0 - or local.data_ptr() % transfer_bytes != 0 - or external.data_ptr() % transfer_bytes != 0 - or local.storage_offset() % elements_per_unit != 0 - or external.storage_offset() % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in local.stride()[:-1]) - or any(stride % elements_per_unit for stride in external.stride()[:-1]) - ): - continue - local = local.view(transfer_dtype) - external = external.view(transfer_dtype) - break + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_gqa_pages[(grid_num,)]( LOCAL=local, diff --git a/lightllm/common/attention_state/full_att/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/group_quantized_transfer.py index 371913e758..5fec71c382 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_transfer.py +++ b/lightllm/common/attention_state/full_att/group_quantized_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import StateTransferLayout +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== @@ -125,28 +126,14 @@ def transfer_group_quantized_pages( # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 及动态 scale 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: - # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 - # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 - # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 - if local.stride(-1) == external.stride(-1) == 1: - element_bytes = local.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - local.shape[-1] % elements_per_unit != 0 - or local.data_ptr() % transfer_bytes != 0 - or external.data_ptr() % transfer_bytes != 0 - or local.storage_offset() % elements_per_unit != 0 - or external.storage_offset() % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in local.stride()[:-1]) - or any(stride % elements_per_unit for stride in external.stride()[:-1]) - ): - continue - local = local.view(transfer_dtype) - external = external.view(transfer_dtype) - break + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_group_quantized_pages[(grid_num,)]( LOCAL=local, diff --git a/lightllm/common/attention_state/full_att/mla_transfer.py b/lightllm/common/attention_state/full_att/mla_transfer.py index 2855077ff3..2172f65c69 100644 --- a/lightllm/common/attention_state/full_att/mla_transfer.py +++ b/lightllm/common/attention_state/full_att/mla_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import StateTransferLayout +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== @@ -110,28 +111,14 @@ def transfer_mla_pages( # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 for local, external in copy_views: - # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 - # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 - # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 - if local.stride(-1) == external.stride(-1) == 1: - element_bytes = local.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - local.shape[-1] % elements_per_unit != 0 - or local.data_ptr() % transfer_bytes != 0 - or external.data_ptr() % transfer_bytes != 0 - or local.storage_offset() % elements_per_unit != 0 - or external.storage_offset() % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in local.stride()[:-1]) - or any(stride % elements_per_unit for stride in external.stride()[:-1]) - ): - continue - local = local.view(transfer_dtype) - external = external.view(transfer_dtype) - break + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_mla_pages[(grid_num,)]( LOCAL=local, diff --git a/lightllm/common/attention_state/full_att/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla_transfer.py index c3d43d6b93..22e39de04e 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import StateTransferLayout +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== @@ -134,28 +135,14 @@ def transfer_packed_mla_pages( # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 KV、indexer 及各量化字段 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: - # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 - # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 - # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 - if local.stride(-1) == external.stride(-1) == 1: - element_bytes = local.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - local.shape[-1] % elements_per_unit != 0 - or local.data_ptr() % transfer_bytes != 0 - or external.data_ptr() % transfer_bytes != 0 - or local.storage_offset() % elements_per_unit != 0 - or external.storage_offset() % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in local.stride()[:-1]) - or any(stride % elements_per_unit for stride in external.stride()[:-1]) - ): - continue - local = local.view(transfer_dtype) - external = external.view(transfer_dtype) - break + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_packed_mla_pages[(grid_num,)]( LOCAL=local, diff --git a/lightllm/common/attention_state/full_att/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/static_fp8_transfer.py index 13ec1cba1c..eefbb15e41 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_transfer.py +++ b/lightllm/common/attention_state/full_att/static_fp8_transfer.py @@ -5,6 +5,7 @@ import triton.language as tl from ..types import StateTransferLayout +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== @@ -117,28 +118,14 @@ def transfer_static_fp8_pages( # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 FP8 K/V 的逐字段 launch 融合， # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 for local, external in copy_views: - # 连续字段按最大 8 字节整数 view 搬运，减少索引计算和访存指令；不做数值转换。 - # token/head/page 轴各有寻址语义，保持维数，只改变最内轴的搬运单位。 - # 行宽、两端起点及外层 stride 都须对齐；packed 字段或 TP 切片不满足时保留原宽度。 - if local.stride(-1) == external.stride(-1) == 1: - element_bytes = local.element_size() - for transfer_bytes, transfer_dtype in ((8, torch.int64), (4, torch.int32), (2, torch.int16)): - if transfer_bytes <= element_bytes: - break - elements_per_unit = transfer_bytes // element_bytes - if ( - local.shape[-1] % elements_per_unit != 0 - or local.data_ptr() % transfer_bytes != 0 - or external.data_ptr() % transfer_bytes != 0 - or local.storage_offset() % elements_per_unit != 0 - or external.storage_offset() % elements_per_unit != 0 - or any(stride % elements_per_unit for stride in local.stride()[:-1]) - or any(stride % elements_per_unit for stride in external.stride()[:-1]) - ): - continue - local = local.view(transfer_dtype) - external = external.view(transfer_dtype) - break + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] _copy_static_fp8_pages[(grid_num,)]( LOCAL=local, diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py index 669b40256a..d607cfab3b 100644 --- a/lightllm/common/attention_state/utils.py +++ b/lightllm/common/attention_state/utils.py @@ -1,8 +1,11 @@ -"""注意力资源的 State 命名工具。""" +"""注意力 State 的命名和传输视图工具。""" from __future__ import annotations +from functools import lru_cache + import torch +from torch._prims import _collapse_view_helper MAX_DRAFT_MODELS = 36 MODEL_ROLES = ("main", *(f"draft{i}" for i in range(MAX_DRAFT_MODELS))) @@ -14,3 +17,138 @@ def dtype_name(dtype: torch.dtype): return {torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32"}[dtype] + + +# ======================================================================== +# 二、传输视图：合并两端共同支持的尾部维度 +# ======================================================================== + + +def merge_contiguous_tail_dims( + copies: list[tuple[torch.Tensor, torch.Tensor]], + start_dim: int = 0, + keepdim: bool = False, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """合并批次两侧共同连续的尾部维度，只创建 view，不复制数据或改变 dtype。 + + 调用方保证同一侧的 shape/stride 相同，两侧待合并的尾部 shape 相同； + 前缀维度的数量和容量可以不同，例如 GPU token 存储与外部页集合。 + 根据首对 Tensor 确定共同的合并范围，再为整批分别创建两侧形状的视图。 + 从末轴向前尝试，遇到任一侧无法合并就停止，不单独整理前面的维度。 + + start_dim 是允许参与合并的最前轴，之前的轴保持独立。例如 + [layer, window, head, dim] 传入 start_dim=2，只合并 head/dim，保留环形 window。 + 支持负数，例如 start_dim=-2 对齐两侧最后两维，不要求两侧前缀维数一致。 + keepdim=True 将被合并的前面各轴保留为 1，元素数集中到末轴，供固定维数的 + token/page kernel 使用；默认 False 直接减少维数，供 bounded 聚合 kernel 使用。 + 长度为 1 的轴是否可以合并由 PyTorch 根据布局判断。 + 通过版本断言限制为已验证的 PyTorch 2.11.x，end 包含在合并区间内。 + helper 对带 padding 的单元素轴可能过于宽松，因此最终以 view 是否成功为准； + 只捕获创建视图的 RuntimeError，退回较小的合并范围，不使用可能复制数据的 reshape。 + """ + _check_collapse_view_version(torch_version=torch.__version__) + source, destination = copies[0] + assert -source.ndim <= start_dim < source.ndim, "invalid tail merge boundary" + if start_dim < 0: + start_dim += source.ndim + tail_dims = source.ndim - start_dim + assert destination.ndim >= tail_dims and source.shape[start_dim:] == destination.shape[-tail_dims:], ( + "copy tail shapes must match" + ) + transfer_shapes = [] + last_dim = source.ndim - 1 + for first_dim in range(last_dim - 1, start_dim - 1, -1): + # 按距末轴的距离对齐区间；两侧独立保留 layer/token/page 等前缀。 + merge_dims = source.ndim - first_dim + destination_first_dim = destination.ndim - merge_dims + source_shape, _ = _collapse_view_helper( + a=source, + start=first_dim, + end=last_dim, + must_be_valid=None, + ) + if source_shape is None: + break + destination_shape, _ = _collapse_view_helper( + a=destination, + start=destination_first_dim, + end=destination.ndim - 1, + must_be_valid=None, + ) + if destination_shape is None: + break + if keepdim: + singleton_axes = (1,) * (merge_dims - 1) + source_shape = source.shape[:first_dim] + singleton_axes + (source_shape[-1],) + destination_shape = destination.shape[:destination_first_dim] + singleton_axes + (destination_shape[-1],) + if source_shape != source.shape or destination_shape != destination.shape: + transfer_shapes.append((source_shape, destination_shape)) + + # 优先采用合并范围最大的形状。普通布局只创建一次整批视图； + # helper 误判时逐步缩小合并范围，保留跨轴的存储间隙。 + for source_shape, destination_shape in reversed(transfer_shapes): + try: + return [(src.view(source_shape), dst.view(destination_shape)) for src, dst in copies] + except RuntimeError: + continue + return copies + + +# ======================================================================== +# 三、传输视图：为整批选择最宽的搬运 dtype +# ======================================================================== + + +def view_max_copy_dtype( + copies: list[tuple[torch.Tensor, torch.Tensor]], +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """将整批视图按最大可用整数宽度搬运，不转换数值、不复制数据。 + + 调用方保证批次非空、dtype 相同，同一侧的 shape/stride 相同，两侧的末轴长度相同。 + 两侧的外层 shape/stride 可以不同。依次尝试 8/4/2 字节宽度，整批统一选择一种 dtype； + 不能扩大时返回原列表。只改变末轴的元素数，保留其余轴及存储间隙。 + 本函数不合并维度；有需要时先调用 merge_contiguous_tail_dims。 + 先用末轴字节数缩小候选范围，再由 view 校验外层 stride 和 Storage 偏移； + 只捕获创建视图的 RuntimeError，任一视图失败就整批退回较窄的 dtype。 + """ + source, destination = copies[0] + element_bytes = source.element_size() + if element_bytes >= 8 or source.stride(-1) != 1 or destination.stride(-1) != 1: + return copies + row_bytes = source.shape[-1] * element_bytes + # 按行宽选择候选类型，例如 12 字节只尝试 INT32、INT16。 + if row_bytes % 8 == 0: + copy_dtypes = ((8, torch.int64), (4, torch.int32), (2, torch.int16)) + elif row_bytes % 4 == 0: + copy_dtypes = ((4, torch.int32), (2, torch.int16)) + elif row_bytes % 2 == 0: + copy_dtypes = ((2, torch.int16),) + else: + return copies + + # 只尝试扩宽；布局是否兼容直接由 view 判断，任一项失败就整批回退。 + for transfer_bytes, transfer_dtype in copy_dtypes: + if transfer_bytes <= element_bytes: + break + try: + return [(src.view(transfer_dtype), dst.view(transfer_dtype)) for src, dst in copies] + except RuntimeError: + continue + return copies + + +# ======================================================================== +# 四、内部依赖约束 +# ======================================================================== + + +@lru_cache(maxsize=1) +def _check_collapse_view_version(torch_version: str) -> None: + """仅允许已验证的 PyTorch 版本使用内部维度合并接口，正常进程只校验一次。""" + # _collapse_view_helper 不是稳定公共 API：历史版本的 end 开闭区间及 + # must_be_valid 参数不同，单元素轴的处理也曾变化。升级后须先核对签名、 + # 区间语义并通过 padding/环形边界回归，再修改此版本范围。 + assert torch_version.startswith("2.11."), ( + f"tail dimension merging supports PyTorch 2.11.x, got {torch_version}; " + "revalidate _collapse_view_helper before updating the supported version" + ) diff --git a/unit_tests/common/att_resources/test_bounded_copy_kernel.py b/unit_tests/common/att_resources/test_bounded_copy_kernel.py index 0984df7a3d..243fb1fa00 100644 --- a/unit_tests/common/att_resources/test_bounded_copy_kernel.py +++ b/unit_tests/common/att_resources/test_bounded_copy_kernel.py @@ -217,7 +217,8 @@ def launch(**kwargs): shape = (1, 2, 1, 3, 1, 7) target_shape = (1, 2, 1, 3, 1, 9) width = 5 - expected_launch = (6, (1, 2, 1, 3, 1, 5), tl.uint32) + # 尾部单元素轴可并入 width；前面的 padding 边界仍然保留。 + expected_launch = (5, (1, 2, 1, 3, 5), tl.uint32) elif layout == "wide_singleton_prefix": # dtype 扩宽只改变末轴长度，前后的单元素轴均保持不变。 shape = (1, 3, 4) diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index 8791e9c808..c6f46ed9fe 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -11,6 +11,87 @@ from unit_tests.common.att_resources.test_state_identity import allocated_state from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout from lightllm.common.attention_state.types import AxisMapping +from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_narrow_head_fields_merge_before_page_transfer(page_device: str, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.full_att import gqa_transfer + + config = FullAttentionConfig( + logical_layer_groups=((0,), (1,)), + global_head_num=4, + head_dim=1, + dtype=torch.bfloat16, + ) + source = GQAFullAttentionState(layout_config=config, page_size=4, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=4, device="cuda") + source.allocate_storage(token_capacity=8) + target.allocate_storage(token_capacity=12) + src_ids = torch.tensor([5, 1, 7], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([9, 0, 6], dtype=torch.int64, device="cuda") + expected = [] + for src_tensor, dst_tensor in ((source.k_storage, target.k_storage), (source.v_storage, target.v_storage)): + src_tensor.view(torch.uint8).random_(0, 256) + dst_tensor.view(torch.uint8).fill_(173) + snapshot = dst_tensor.clone() + snapshot[:, dst_ids] = src_tensor[:, src_ids] + expected.append(snapshot) + + layout = transfer_layout(source) + pages = torch.full( + size=(2, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, 16:-16] + launches = [] + kernel = gqa_transfer._copy_gqa_pages + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((kwargs["HEADS"], kwargs["WIDTH"], kwargs["ELEMENT_BYTES"])) + return kernel[grid](**kwargs) + return launch + + monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) + if page_device == "cpu": + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=[1], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=[1], + token_indices=dst_ids, + layout=layout, + ) + else: + source.write_to_pd_page( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=1, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=1, + layout=layout, + ) + torch.cuda.synchronize() + # 单个 BF16 head 只有 2 字节；合并 4 个 head 后才能按 INT64 搬运。 + assert launches == [(1, 1, 8)] * 4 + for actual, reference in zip((target.k_storage, target.v_storage), expected): + assert torch.equal(actual.view(torch.uint8), reference.view(torch.uint8)) + assert (pages[0] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_full_att_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py index c07b8d8453..2a63449def 100644 --- a/unit_tests/common/att_resources/test_full_att_dp_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -19,6 +19,42 @@ SCHEMES = ("gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_dp_narrow_head_fields_merge_before_copy(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.full_att import dp_transfer + + source = torch.empty(size=(2, 17, 4, 1), dtype=torch.bfloat16, device="cuda") + target = torch.empty(size=(2, 23, 4, 1), dtype=torch.bfloat16, device="cuda") + source.view(torch.uint8).random_(0, 256) + target.view(torch.uint8).fill_(173) + src_ids = torch.tensor([12, 0, 12], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([19, 1, 6], dtype=torch.int64, device="cuda") + expected = target.clone() + expected[:, dst_ids] = source[:, src_ids] + kernel = dp_transfer._copy_indexed_token_storages + launches = [] + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((kwargs["SHAPES"], kwargs["ELEMENT_BYTES"])) + return kernel[grid](**kwargs) + return launch + + monkeypatch.setattr(dp_transfer, "_copy_indexed_token_storages", CaptureLaunch()) + copy_gqa_tokens( + src_k_storage=source, + src_v_storage=source, + dst_k_storage=target, + dst_v_storage=torch.empty_like(target), + src_token_indices=src_ids, + dst_token_indices=dst_ids, + ) + torch.cuda.synchronize() + assert launches == [(((2, 1, 1), (2, 1, 1)), (8, 8))] + assert torch.equal(target.view(torch.uint8), expected.view(torch.uint8)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kind", SCHEMES) @pytest.mark.parametrize("source_device", [0, 1]) diff --git a/unit_tests/common/att_resources/test_transfer_views.py b/unit_tests/common/att_resources/test_transfer_views.py new file mode 100644 index 0000000000..0c208de515 --- /dev/null +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -0,0 +1,189 @@ +"""传输视图整理必须保留存储别名、位模式、未参与区域和独立寻址轴。""" + +import pytest +import torch + +from lightllm.common.attention_state.utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +@pytest.mark.parametrize( + "layout, expected_shape, start_dim", + [ + ("dense", (48,), 0), + ("source_gap", (2, 3, 8), 0), + ("destination_gap", (2, 3, 8), 0), + ("strided_tail", (48,), 0), + ("singleton_gap", (2, 8), 0), + ("protected_window", (2, 5, 24), 2), + ], +) +def test_merge_tail_preserves_copy_regions( + layout: str, expected_shape: tuple[int, ...], start_dim: int, +) -> None: + source_shape = (2, 5, 3, 8) if layout == "protected_window" else (2, 3, 8) + destination_shape = source_shape + if layout == "source_gap": + source_shape = (2, 3, 10) + elif layout == "destination_gap": + destination_shape = (2, 3, 10) + elif layout == "strided_tail": + source_shape = (2, 3, 16) + elif layout == "singleton_gap": + source_shape, destination_shape = (2, 1, 10), (2, 1, 12) + + source_storage = torch.arange(torch.Size(source_shape).numel(), dtype=torch.float32).view(source_shape) + destination_storage = torch.full(size=destination_shape, fill_value=-1, dtype=torch.float32) + source = source_storage[..., ::2] if layout == "strided_tail" else source_storage[..., :8] + destination = destination_storage[..., :8] + # 直接按原始视图执行参考复制，保留目标 padding，随后比较整个底层存储。 + expected = destination_storage.clone() + expected.as_strided( + size=destination.shape, + stride=destination.stride(), + storage_offset=destination.storage_offset(), + ).copy_(source) + + copies = [(source, destination)] + merged = merge_contiguous_tail_dims(copies=copies, start_dim=start_dim) + source_view, destination_view = merged[0] + assert source_view.shape == destination_view.shape == expected_shape + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + if expected_shape == source.shape: + assert merged is copies + destination_view.copy_(source_view) + assert torch.equal(destination_storage, expected) + + +def test_merge_request_slices_preserves_batch_offsets() -> None: + source = torch.arange(2 * 4 * 3 * 8, dtype=torch.float32).view(2, 4, 3, 8) + destination = torch.full(size=(4, 2, 3, 8), fill_value=-1, dtype=torch.float32) + expected = destination.clone() + copies = [] + for req_idx in (0, 2): + copies.append((source[:, req_idx], destination[req_idx])) + expected[req_idx].copy_(source[:, req_idx]) + + for source_view, destination_view in merge_contiguous_tail_dims(copies=copies): + assert source_view.shape == destination_view.shape == (2, 24) + destination_view.copy_(source_view) + assert torch.equal(destination, expected) + + +@pytest.mark.parametrize("destination_prefix", [(2, 11), (3, 2, 5)]) +@pytest.mark.parametrize("keepdim", [False, True]) +def test_merge_tail_preserves_different_prefix_layouts(destination_prefix: tuple[int, ...], keepdim: bool) -> None: + source = torch.arange(2 * 7 * 4 * 2, dtype=torch.float32).view(2, 7, 4, 2) + destination_shape = destination_prefix + (4, 2) + destination = torch.arange(torch.Size(destination_shape).numel(), dtype=torch.float32).view(destination_shape) + source_view, destination_view = merge_contiguous_tail_dims( + copies=[(source, destination)], + start_dim=-2, + keepdim=keepdim, + )[0] + tail_shape = (1, 8) if keepdim else (8,) + assert source_view.shape == (2, 7) + tail_shape + assert destination_view.shape == destination_prefix + tail_shape + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + assert torch.equal(source_view.flatten(), source.flatten()) + assert torch.equal(destination_view.flatten(), destination.flatten()) + + +@pytest.mark.parametrize("gap_side", ["source", "destination"]) +def test_merge_tail_keeps_head_padding(gap_side: str) -> None: + source = torch.empty(size=(2, 7, 4, 3 if gap_side == "source" else 2))[..., :2] + destination = torch.empty(size=(3, 2, 5, 4, 3 if gap_side == "destination" else 2))[..., :2] + copies = [(source, destination)] + assert merge_contiguous_tail_dims(copies=copies, start_dim=-2, keepdim=True) is copies + + +@pytest.mark.parametrize("torch_version", ["2.9.0", "2.10.0", "2.12.0", "3.0.0"]) +def test_merge_tail_rejects_unverified_torch_version(torch_version: str, monkeypatch: pytest.MonkeyPatch) -> None: + source = torch.empty(size=(2, 4)) + destination = torch.empty_like(source) + monkeypatch.setattr(torch, "__version__", torch_version) + with pytest.raises(AssertionError, match="supports PyTorch 2.11.x"): + merge_contiguous_tail_dims(copies=[(source, destination)]) + + +@pytest.mark.parametrize("torch_version", ["2.11.0+cu130", "2.11.1"]) +def test_merge_tail_accepts_supported_torch_version(torch_version: str, monkeypatch: pytest.MonkeyPatch) -> None: + source = torch.empty(size=(2, 4)) + destination = torch.empty_like(source) + monkeypatch.setattr(torch, "__version__", torch_version) + source_view, destination_view = merge_contiguous_tail_dims(copies=[(source, destination)])[0] + assert source_view.shape == destination_view.shape == (8,) + + +@pytest.mark.parametrize( + "dtype, width, src_offset, dst_offset, expected_dtype", + [ + (torch.bfloat16, 8, 0, 0, torch.int64), + (torch.bfloat16, 6, 0, 0, torch.int32), + (torch.int8, 2, 0, 0, torch.int16), + (torch.float32, 3, 0, 0, torch.float32), + (torch.int64, 8, 0, 0, torch.int64), + (torch.bfloat16, 8, 1, 0, torch.bfloat16), + (torch.bfloat16, 8, 0, 1, torch.bfloat16), + (torch.bfloat16, 8, 2, 0, torch.int32), + ], +) +def test_max_copy_dtype_preserves_data_and_aliases( + dtype: torch.dtype, + width: int, + src_offset: int, + dst_offset: int, + expected_dtype: torch.dtype, +) -> None: + source = torch.arange(32).to(dtype=dtype) + destination = torch.full(size=(32,), fill_value=-1, dtype=dtype) + expected = destination.clone() + source_view = source[src_offset:src_offset + width] + destination_view = destination[dst_offset:dst_offset + width] + expected[dst_offset:dst_offset + width].copy_(source_view) + + copies = [(source_view, destination_view)] + converted = view_max_copy_dtype(copies=copies) + copy_source, copy_destination = converted[0] + assert copy_source.dtype == copy_destination.dtype == expected_dtype + assert copy_source.data_ptr() == source_view.data_ptr() + assert copy_destination.data_ptr() == destination_view.data_ptr() + assert copy_source.ndim == source_view.ndim + if expected_dtype == dtype: + assert converted is copies + copy_destination.copy_(copy_source) + assert torch.equal(destination, expected) + + +@pytest.mark.parametrize( + "src_padding, dst_padding, expected_dtype", + [(1, 0, torch.bfloat16), (0, 2, torch.int32), (0, 0, torch.int64)], +) +def test_max_copy_dtype_preserves_row_padding( + src_padding: int, dst_padding: int, expected_dtype: torch.dtype, +) -> None: + source = torch.arange(2 * (8 + src_padding)).to(dtype=torch.bfloat16).view(2, 8 + src_padding) + destination = torch.full(size=(2, 8 + dst_padding), fill_value=-1, dtype=torch.bfloat16) + expected = destination.clone() + expected[:, :8].copy_(source[:, :8]) + + source_view, destination_view = view_max_copy_dtype(copies=[(source[:, :8], destination[:, :8])])[0] + assert source_view.dtype == destination_view.dtype == expected_dtype + assert source_view.shape[0] == destination_view.shape[0] == 2 + destination_view.copy_(source_view) + assert torch.equal(destination, expected) + + +def test_max_copy_dtype_falls_back_for_middle_batch_offset() -> None: + source = torch.arange(64).to(dtype=torch.bfloat16) + destination = torch.full(size=(64,), fill_value=-1, dtype=torch.bfloat16) + copies = [(source[start:start + 8], destination[end:end + 8]) for start, end in ((0, 0), (9, 16), (24, 32))] + assert view_max_copy_dtype(copies=copies) is copies + + +def test_max_copy_dtype_keeps_strided_last_axis() -> None: + source = torch.arange(32, dtype=torch.float32)[::2] + destination = torch.empty(size=(16,), dtype=torch.float32) + copies = [(source, destination)] + assert view_max_copy_dtype(copies=copies) is copies From 1c090b95442e3c6508f7990c269d583306937bad Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 07:49:45 +0000 Subject: [PATCH 26/56] perf: prefilter incompatible attention copy dtype views --- lightllm/common/attention_state/utils.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py index d607cfab3b..ccecd7d2b6 100644 --- a/lightllm/common/attention_state/utils.py +++ b/lightllm/common/attention_state/utils.py @@ -108,7 +108,8 @@ def view_max_copy_dtype( 两侧的外层 shape/stride 可以不同。依次尝试 8/4/2 字节宽度，整批统一选择一种 dtype； 不能扩大时返回原列表。只改变末轴的元素数，保留其余轴及存储间隙。 本函数不合并维度；有需要时先调用 merge_contiguous_tail_dims。 - 先用末轴字节数缩小候选范围，再由 view 校验外层 stride 和 Storage 偏移； + 先用末轴字节数、首对 Tensor 的外层 stride 和 Storage 偏移过滤候选； + 同一侧的 stride 由整批共用，其余切片的 Storage 偏移继续由 view 校验。 只捕获创建视图的 RuntimeError，任一视图失败就整批退回较窄的 dtype。 """ source, destination = copies[0] @@ -126,10 +127,21 @@ def view_max_copy_dtype( else: return copies - # 只尝试扩宽；布局是否兼容直接由 view 判断，任一项失败就整批回退。 + # 首对布局即可排除必然失败的候选，避免 MTP Conv 等带间隙切片反复抛异常。 + # 这些值以原 dtype 的元素数计量，扩宽后必须仍能表示为整数 stride/offset。 + strides_and_offsets = ( + *source.stride()[:-1], + *destination.stride()[:-1], + source.storage_offset(), + destination.storage_offset(), + ) for transfer_bytes, transfer_dtype in copy_dtypes: if transfer_bytes <= element_bytes: break + elements_per_unit = transfer_bytes // element_bytes + if any(value % elements_per_unit for value in strides_and_offsets): + continue + # 批次中其他切片的起点可能不同；保留 view 的整批校验和较窄类型回退。 try: return [(src.view(transfer_dtype), dst.view(transfer_dtype)) for src, dst in copies] except RuntimeError: From 3bec22ba06f3bb7bd2390a743a8ca0b86fcdf0d4 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 08:03:41 +0000 Subject: [PATCH 27/56] perf: simplify attention tail dimension merging with stride checks --- docs/cache_manager_design.md | 4 +- lightllm/common/attention_state/utils.py | 112 +++++++----------- .../att_resources/test_transfer_views.py | 46 ++++--- 3 files changed, 77 insertions(+), 85 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 87abcaa68a..4345f1fdf6 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -362,7 +362,7 @@ lightllm/utils/ `ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `state_name`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 -KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个阶段顺序执行：公共 `utils.merge_contiguous_tail_dims` 合并连续尾维，公共 `utils.view_max_copy_dtype` 为整批选择最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。公共函数使用 PyTorch `_collapse_view_helper(..., must_be_valid=None)` 判断两端均可合并的尾部区间，随后调用 `view` 创建整批视图；只读取布局元数据，不复制数据。PyTorch 2.11 的 helper 对带 padding 的单元素轴可能过于宽松，最终创建视图时仅捕获 `view` 的 `RuntimeError` 并退回较小范围，以防错误跨越存储间隙；不使用可能复制数据的 `reshape`。`start_dim` 限制允许参与合并的最前轴，例如 `[layer, window, head, dim]` 使用 `start_dim=2` 保留环形 window，只合并 head/dim；`start_dim=-2` 按两侧末尾两维对齐，允许前缀维数、token 容量不同。默认直接减少维数；`keepdim=True` 将被合并的前面各轴保留为 1，使元素集中到末轴，供固定维数的 token/page kernel 使用。调用前通过缓存的 `_check_collapse_view_version` 断言，只允许已验证的 PyTorch 2.11.x，`end` 包含在合并区间内。该函数属于 PyTorch 内部 API，历史版本的参数和区间语义不同；升级 PyTorch 时须先核对签名、合并区间并通过 padding/环形边界回归，再调整允许的版本范围。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧无法合并就停止，保留前面的维度；尾部单元素轴由 PyTorch 按布局处理，再按行宽选择候选类型，通过 `view` 校验布局并尝试创建最大 8 字节的整数视图，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 +KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个阶段顺序执行：公共 `utils.merge_contiguous_tail_dims` 合并连续尾维，公共 `utils.view_max_copy_dtype` 为整批选择最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。公共函数只读取首对 Tensor 的 shape/stride，从末轴向前扫描一次，判断两端共同连续的尾部区间，再调用 `view` 创建整批视图。普通轴须满足外层 stride 等于已合并尾部元素数乘尾部元素 stride；长度为 1 的轴不约束 stride，内侧仅有一个元素时改用当前非单元素轴的 stride。这样可正确保留 padding，也支持末轴步长不为 1 的视图。不复制数据，不依赖 PyTorch 内部 API，不以异常作为正常回退路径，也不使用可能复制数据的 `reshape`。`start_dim` 限制允许参与合并的最前轴，例如 `[layer, window, head, dim]` 使用 `start_dim=2` 保留环形 window，只合并 head/dim；`start_dim=-2` 按两侧末尾两维对齐，允许前缀维数、token 容量不同。默认直接减少维数；`keepdim=True` 将被合并的前面各轴保留为 1，使元素集中到末轴，供固定维数的 token/page kernel 使用。实现仅依赖公开的 Tensor shape、stride 和 view 接口，不再需要针对 PyTorch 内部接口的版本断言。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧无法合并就停止，保留前面的维度；尾部单元素轴按上述规则处理，再按行宽选择候选类型，通过首对布局过滤及整批 `view` 校验尝试创建最大 8 字节的整数视图，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 默认参数在 H200 上固定 grid=24 联合调优，覆盖 Conv BF16 / SSM FP32、三种规模和 CUDA↔CUDA / CUDA↔pinned CPU / pinned CPU↔pinned CPU，共 24 个场景。第一轮比较 BLOCK_ELEMENTS=512/1024/2048/4096 与 warps=1/2/4/8，并保留 256/8/1 对照；第二轮为入选组合比较 num_stages=1/2/3/4，两轮共测试 29 个不同组合。每个候选均校验数据，预先创建指针表，以随机顺序进行 CUDA Graph 计时并取中位数。按场景等权的速度几何平均选择 1024/8/1；相比旧默认 grid=16、BLOCK_ELEMENTS=256、warps=8、stages=1，第二轮综合速度约为 1.58 倍。该组合并非每项最快，小块场景最多增加约 1.35 μs（Conv CPU→GPU 约 2.42→3.77 μs）。`grid_num`、`num_warps`、`num_stages` 可由调用方覆盖；流水深度用于 `tl.range` 的循环属性。 @@ -382,7 +382,7 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，其他外层 stride 和 Storage 偏移交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP/同卡 token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py index ccecd7d2b6..a1181d1d24 100644 --- a/lightllm/common/attention_state/utils.py +++ b/lightllm/common/attention_state/utils.py @@ -2,10 +2,7 @@ from __future__ import annotations -from functools import lru_cache - import torch -from torch._prims import _collapse_view_helper MAX_DRAFT_MODELS = 36 MODEL_ROLES = ("main", *(f"draft{i}" for i in range(MAX_DRAFT_MODELS))) @@ -31,7 +28,7 @@ def merge_contiguous_tail_dims( ) -> list[tuple[torch.Tensor, torch.Tensor]]: """合并批次两侧共同连续的尾部维度，只创建 view，不复制数据或改变 dtype。 - 调用方保证同一侧的 shape/stride 相同，两侧待合并的尾部 shape 相同； + 调用方保证批次及 Tensor 非空，同一侧的 shape/stride 相同，两侧待合并的尾部 shape 相同； 前缀维度的数量和容量可以不同，例如 GPU token 存储与外部页集合。 根据首对 Tensor 确定共同的合并范围，再为整批分别创建两侧形状的视图。 从末轴向前尝试，遇到任一侧无法合并就停止，不单独整理前面的维度。 @@ -41,57 +38,55 @@ def merge_contiguous_tail_dims( 支持负数，例如 start_dim=-2 对齐两侧最后两维，不要求两侧前缀维数一致。 keepdim=True 将被合并的前面各轴保留为 1，元素数集中到末轴，供固定维数的 token/page kernel 使用；默认 False 直接减少维数，供 bounded 聚合 kernel 使用。 - 长度为 1 的轴是否可以合并由 PyTorch 根据布局判断。 - 通过版本断言限制为已验证的 PyTorch 2.11.x，end 包含在合并区间内。 - helper 对带 padding 的单元素轴可能过于宽松，因此最终以 view 是否成功为准； - 只捕获创建视图的 RuntimeError，退回较小的合并范围，不使用可能复制数据的 reshape。 + 直接按 shape/stride 判断连续性，长度为 1 的轴不约束 stride。 + 只依赖公开 Tensor 接口，确定形状后一次性创建 view，不使用可能复制数据的 reshape。 """ - _check_collapse_view_version(torch_version=torch.__version__) source, destination = copies[0] - assert -source.ndim <= start_dim < source.ndim, "invalid tail merge boundary" + source_shape, destination_shape = source.shape, destination.shape + source_strides, destination_strides = source.stride(), destination.stride() + ndim = len(source_shape) + assert -ndim <= start_dim < ndim, "invalid tail merge boundary" if start_dim < 0: - start_dim += source.ndim - tail_dims = source.ndim - start_dim - assert destination.ndim >= tail_dims and source.shape[start_dim:] == destination.shape[-tail_dims:], ( + start_dim += ndim + tail_dims = ndim - start_dim + assert len(destination_shape) >= tail_dims and source_shape[start_dim:] == destination_shape[-tail_dims:], ( "copy tail shapes must match" ) - transfer_shapes = [] - last_dim = source.ndim - 1 - for first_dim in range(last_dim - 1, start_dim - 1, -1): - # 按距末轴的距离对齐区间；两侧独立保留 layer/token/page 等前缀。 - merge_dims = source.ndim - first_dim - destination_first_dim = destination.ndim - merge_dims - source_shape, _ = _collapse_view_helper( - a=source, - start=first_dim, - end=last_dim, - must_be_valid=None, - ) - if source_shape is None: - break - destination_shape, _ = _collapse_view_helper( - a=destination, - start=destination_first_dim, - end=destination.ndim - 1, - must_be_valid=None, - ) - if destination_shape is None: - break - if keepdim: - singleton_axes = (1,) * (merge_dims - 1) - source_shape = source.shape[:first_dim] + singleton_axes + (source_shape[-1],) - destination_shape = destination.shape[:destination_first_dim] + singleton_axes + (destination_shape[-1],) - if source_shape != source.shape or destination_shape != destination.shape: - transfer_shapes.append((source_shape, destination_shape)) - - # 优先采用合并范围最大的形状。普通布局只创建一次整批视图； - # helper 误判时逐步缩小合并范围，保留跨轴的存储间隙。 - for source_shape, destination_shape in reversed(transfer_shapes): - try: - return [(src.view(source_shape), dst.view(destination_shape)) for src, dst in copies] - except RuntimeError: - continue - return copies + + # 阶段一：从末轴向前检查；两侧的前缀容量和轴数独立，按距末轴的距离对齐。 + first_tail_dim = ndim - 1 + tail_elements = source_shape[-1] + source_tail_stride, destination_tail_stride = source_strides[-1], destination_strides[-1] + prefix_dim_difference = len(destination_shape) - ndim + for source_axis in range(ndim - 2, start_dim - 1, -1): + destination_axis = source_axis + prefix_dim_difference + axis_size = source_shape[source_axis] + if axis_size != 1: + if tail_elements == 1: + # 内侧只有一个元素，其 stride 不影响寻址；以当前非单元素轴为准。 + source_tail_stride = source_strides[source_axis] + destination_tail_stride = destination_strides[destination_axis] + elif ( + source_strides[source_axis] != tail_elements * source_tail_stride + or destination_strides[destination_axis] != tail_elements * destination_tail_stride + ): + break + # 当前轴长度为 1 时忽略其 stride，但保留已累计的连续区间和元素步长。 + tail_elements *= axis_size + first_tail_dim = source_axis + + if first_tail_dim == ndim - 1: + return copies + + # 阶段二：整批复用普通 tuple 形状，减少 view 的参数解析开销。 + # keepdim 保留固定维数 kernel 所需的单元素轴。 + merge_dims = ndim - first_tail_dim + tail_shape = (1,) * (merge_dims - 1) + (tail_elements,) if keepdim else (tail_elements,) + source_transfer_shape = tuple(source_shape[:-merge_dims]) + tail_shape + destination_transfer_shape = tuple(destination_shape[:-merge_dims]) + tail_shape + if source_transfer_shape == source_shape and destination_transfer_shape == destination_shape: + return copies + return [(src.view(source_transfer_shape), dst.view(destination_transfer_shape)) for src, dst in copies] # ======================================================================== @@ -147,20 +142,3 @@ def view_max_copy_dtype( except RuntimeError: continue return copies - - -# ======================================================================== -# 四、内部依赖约束 -# ======================================================================== - - -@lru_cache(maxsize=1) -def _check_collapse_view_version(torch_version: str) -> None: - """仅允许已验证的 PyTorch 版本使用内部维度合并接口，正常进程只校验一次。""" - # _collapse_view_helper 不是稳定公共 API：历史版本的 end 开闭区间及 - # must_be_valid 参数不同，单元素轴的处理也曾变化。升级后须先核对签名、 - # 区间语义并通过 padding/环形边界回归，再修改此版本范围。 - assert torch_version.startswith("2.11."), ( - f"tail dimension merging supports PyTorch 2.11.x, got {torch_version}; " - "revalidate _collapse_view_helper before updating the supported version" - ) diff --git a/unit_tests/common/att_resources/test_transfer_views.py b/unit_tests/common/att_resources/test_transfer_views.py index 0c208de515..0952c13503 100644 --- a/unit_tests/common/att_resources/test_transfer_views.py +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -98,22 +98,36 @@ def test_merge_tail_keeps_head_padding(gap_side: str) -> None: assert merge_contiguous_tail_dims(copies=copies, start_dim=-2, keepdim=True) is copies -@pytest.mark.parametrize("torch_version", ["2.9.0", "2.10.0", "2.12.0", "3.0.0"]) -def test_merge_tail_rejects_unverified_torch_version(torch_version: str, monkeypatch: pytest.MonkeyPatch) -> None: - source = torch.empty(size=(2, 4)) - destination = torch.empty_like(source) - monkeypatch.setattr(torch, "__version__", torch_version) - with pytest.raises(AssertionError, match="supports PyTorch 2.11.x"): - merge_contiguous_tail_dims(copies=[(source, destination)]) - - -@pytest.mark.parametrize("torch_version", ["2.11.0+cu130", "2.11.1"]) -def test_merge_tail_accepts_supported_torch_version(torch_version: str, monkeypatch: pytest.MonkeyPatch) -> None: - source = torch.empty(size=(2, 4)) - destination = torch.empty_like(source) - monkeypatch.setattr(torch, "__version__", torch_version) - source_view, destination_view = merge_contiguous_tail_dims(copies=[(source, destination)])[0] - assert source_view.shape == destination_view.shape == (8,) +@pytest.mark.parametrize( + "shape, source_strides, destination_strides, flat_shape, kept_shape", + [ + ((2, 1, 8), (10, 255, 1), (12, 7, 1), (2, 8), (2, 1, 8)), + ((2, 3, 1), (30, 2, 71), (11, 1, 59), (2, 3), (2, 1, 3)), + ((2, 1, 3, 1), (30, 500, 2, 99), (16, 1, 1, 32), (2, 3), (2, 1, 1, 3)), + ((2, 1, 1), (7, 100, 300), (11, 90, 20), (2,), (1, 1, 2)), + ], +) +@pytest.mark.parametrize("keepdim", [False, True]) +def test_merge_tail_singleton_axes_preserve_addressing( + shape: tuple[int, ...], + source_strides: tuple[int, ...], + destination_strides: tuple[int, ...], + flat_shape: tuple[int, ...], + kept_shape: tuple[int, ...], + keepdim: bool, +) -> None: + source_size = 1 + sum((size - 1) * stride for size, stride in zip(shape, source_strides)) + destination_size = 1 + sum((size - 1) * stride for size, stride in zip(shape, destination_strides)) + source = torch.arange(source_size).as_strided(size=shape, stride=source_strides) + destination = torch.arange(destination_size).as_strided(size=shape, stride=destination_strides) + source_view, destination_view = merge_contiguous_tail_dims( + copies=[(source, destination)], keepdim=keepdim, + )[0] + assert source_view.shape == destination_view.shape == (kept_shape if keepdim else flat_shape) + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + assert torch.equal(source_view.flatten(), source.flatten()) + assert torch.equal(destination_view.flatten(), destination.flatten()) @pytest.mark.parametrize( From 4e3c215d9ee3237883c7fff054228df34be52660 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 14:49:34 +0000 Subject: [PATCH 28/56] refactor: integrate DeepSeek V4 with atomic attention states Separate compressed history, sliding-window KV, and window compressor states. Migrate V4, MTP, and DSpark to the shared attention manager and allocator lifecycle. Add independent group64 E8M0 packed MLA storage for GLM, preserve Loop execution-layer cache isolation, and simplify state configuration, transfer naming, and layout caching. Validation: 1528 tests passed; 6 existing FlashMLA/vLLM compatibility cases deselected. Update cache manager documentation and regression coverage. --- docs/cache_manager_design.md | 162 +++++- .../attention_state/bounded_att/__init__.py | 19 +- .../attention_state/bounded_att/base_state.py | 9 + .../create_sliding_window_state.py | 19 + .../create_window_compressor_state.py | 19 + .../bounded_att/fp8_e8m0_swa_state.py | 365 +++++++++++++ .../bounded_att/fp8_e8m0_swa_transfer.py | 97 ++++ .../attention_state/bounded_att/gdn_state.py | 5 +- .../attention_state/bounded_att/kda_state.py | 5 +- .../bounded_att/window_compressor_state.py | 377 +++++++++++++ .../bounded_att/window_compressor_transfer.py | 44 ++ .../bounded_att/window_pages_pool.py | 27 + .../attention_state/full_att/__init__.py | 14 + .../attention_state/full_att/base_state.py | 27 +- .../full_att/compressed_att_state.py | 425 +++++++++++++++ .../full_att/create_compressed_att_state.py | 14 + .../full_att/create_mla_state.py | 25 +- .../attention_state/full_att/dp_transfer.py | 7 +- .../full_att/fp8_e8m0_packed_mla_state.py | 299 ++++++++++ .../full_att/fp8_e8m0_packed_mla_write.py} | 52 +- .../attention_state/full_att/gqa_state.py | 5 +- .../attention_state/full_att/gqa_transfer.py | 8 +- .../full_att/group_quantized_state.py | 5 +- .../full_att/group_quantized_transfer.py | 8 +- .../attention_state/full_att/mla_state.py | 142 ++--- .../full_att/packed_mla_state.py | 9 +- .../full_att/packed_mla_transfer.py | 8 +- .../{mla_transfer.py => packed_transfer.py} | 20 +- .../full_att/static_fp8_state.py | 5 + .../full_att/static_fp8_transfer.py | 8 +- lightllm/common/attention_state/manager.py | 27 +- lightllm/common/attention_state/utils.py | 23 +- .../attention/nsa/dsv4_fp8_flashmla_sparse.py | 91 ++-- .../deepseek4_mem_manager.py | 175 ------ .../fp8_glm5_next_mem_manager.py | 19 - lightllm/common/req_manager/deepseek4.py | 88 --- .../common/state_cache_manager/deepseek4.py | 146 ----- lightllm/models/deepseek_v4/config.py | 85 +++ lightllm/models/deepseek_v4/infer_struct.py | 72 ++- .../deepseek_v4/layer_infer/compressor.py | 28 +- .../models/deepseek_v4/layer_infer/indexer.py | 15 +- .../layer_infer/transformer_layer_infer.py | 10 +- lightllm/models/deepseek_v4/model.py | 137 ++++- .../triton_kernel/build_dspark_swa_indices.py | 13 +- .../triton_kernel/build_swa_indices.py | 10 +- .../deepseek_v4/triton_kernel/cache_io.py | 222 -------- .../deepseek_v4/triton_kernel/swa_cache.py | 139 ++--- .../models/deepseek_v4_dspark/infer_struct.py | 7 +- .../layer_infer/transformer_layer_infer.py | 5 +- .../pre_and_post_layer_weight.py | 4 +- .../layer_weights/transformer_layer_weight.py | 2 +- lightllm/models/deepseek_v4_dspark/model.py | 50 +- .../pre_and_post_layer_weight.py | 2 +- lightllm/models/deepseek_v4_mtp/model.py | 25 +- lightllm/models/glm5_next/model.py | 7 +- lightllm/models/glm5_next_mtp/model.py | 1 + lightllm/models/qwen3_5_mtp/model.py | 2 +- lightllm/models/qwen3next/config.py | 5 + lightllm/models/qwen3next/model.py | 3 +- lightllm/models/sensenova_loop/model.py | 31 +- unit_tests/common/att_resources/conftest.py | 10 + .../common/att_resources/test_capacity.py | 7 +- .../att_resources/test_dsv4_integration.py | 280 ++++++++++ .../common/att_resources/test_dsv4_state.py | 509 ++++++++++++++++++ .../test_fp8_e8m0_packed_mla_state.py | 244 +++++++++ .../att_resources/test_full_att_transfer.py | 2 +- .../att_resources/test_manager_helpers.py | 8 + .../att_resources/test_state_identity.py | 5 +- .../att_resources/test_transfer_views.py | 23 + .../basemodel/attention/test_dsv4_flashmla.py | 18 +- unit_tests/common/dsv4_test_utils.py | 473 ++++++++++++++++ .../common/test_deepseek4_paged_cache.py | 253 +++++---- unit_tests/models/deepseek_v4/test_model.py | 4 +- .../deepseek_v4/test_vision_integration.py | 9 +- unit_tests/models/glm5_next/test_fp8_kv.py | 91 +++- unit_tests/models/sensenova_loop/test_loop.py | 51 +- unit_tests/server/test_pd_start_args.py | 39 +- 77 files changed, 4341 insertions(+), 1358 deletions(-) create mode 100644 lightllm/common/attention_state/bounded_att/create_sliding_window_state.py create mode 100644 lightllm/common/attention_state/bounded_att/create_window_compressor_state.py create mode 100644 lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py create mode 100644 lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py create mode 100644 lightllm/common/attention_state/bounded_att/window_compressor_state.py create mode 100644 lightllm/common/attention_state/bounded_att/window_compressor_transfer.py create mode 100644 lightllm/common/attention_state/bounded_att/window_pages_pool.py create mode 100644 lightllm/common/attention_state/full_att/compressed_att_state.py create mode 100644 lightllm/common/attention_state/full_att/create_compressed_att_state.py create mode 100644 lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py rename lightllm/{models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py => common/attention_state/full_att/fp8_e8m0_packed_mla_write.py} (64%) rename lightllm/common/attention_state/full_att/{mla_transfer.py => packed_transfer.py} (95%) delete mode 100644 lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py delete mode 100644 lightllm/common/req_manager/deepseek4.py delete mode 100644 lightllm/common/state_cache_manager/deepseek4.py create mode 100644 lightllm/models/deepseek_v4/config.py delete mode 100644 lightllm/models/deepseek_v4/triton_kernel/cache_io.py create mode 100644 unit_tests/common/att_resources/test_dsv4_integration.py create mode 100644 unit_tests/common/att_resources/test_dsv4_state.py create mode 100644 unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py create mode 100644 unit_tests/common/dsv4_test_utils.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 4345f1fdf6..dce91a7fe3 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -322,8 +322,12 @@ lightllm/common/ │ │ ├── static_fp8_transfer.py # FP8 K/V 原始编码的定制传输 │ │ ├── create_mla_state.py # MLA 统一创建入口：选择普通、packed 或 packed FP8 │ │ ├── mla_state.py # MLAFullAttentionState：压缩 latent 存储 -│ │ ├── mla_transfer.py # latent KV 的定制传输 -│ │ ├── packed_mla_state.py # PackedMLAFullAttentionState：含 indexer key/scale 的 packed 布局 +│ │ ├── create_compressed_att_state.py # 序列压缩历史 KV 的独立创建入口 +│ │ ├── compressed_att_state.py # 独立配置及 V4 C4/C128 packed 压缩块存储 +│ │ ├── packed_transfer.py # 单字段 packed KV 的定制传输 +│ │ ├── packed_mla_state.py # 普通精度或 group128/FP32-scale packed MLA +│ │ ├── fp8_e8m0_packed_mla_state.py # 独立 FP8E8M0PackedMLAConfig / State：group64/E8M0 编码 +│ │ ├── fp8_e8m0_packed_mla_write.py # group64/E8M0 定制量化写入，不依赖具体模型 │ │ └── packed_mla_transfer.py # packed KV、RoPE 与 indexer 字段的定制传输 │ ├── bounded_att/ │ │ ├── __init__.py @@ -332,6 +336,13 @@ lightllm/common/ │ │ ├── transfer_utils.py # bounded_att_buffer_copies：生成快照字段的复制视图 │ │ ├── copy_kernel.py # copy_bounded_att_tensors：聚合多页、多字段的 Triton 拷贝 │ │ ├── conv_ssm_types.py # KDA/GDN 的 Conv/SSM 布局配置 +│ │ ├── create_sliding_window_state.py # 按 SWA 算法命名的创建入口，当前返回 FP8/E8M0 实现 +│ │ ├── fp8_e8m0_swa_state.py # FP8E8M0SlidingWindowConfig / State：仅管理近期 KV +│ │ ├── create_window_compressor_state.py # 窗口压缩器原子运行态创建入口 +│ │ ├── window_compressor_state.py # WindowCompressorStateConfig / State：候选与 gate +│ │ ├── window_pages_pool.py # WindowPagesPool：单一窗口的时间有序快照 +│ │ ├── fp8_e8m0_swa_transfer.py # FP8/E8M0 packed SWA 的 GPU 窗口搬运 +│ │ ├── window_compressor_transfer.py # 候选/gate ring 与时间有序快照的 GPU 搬运 │ │ ├── create_kda_state.py # KDA 统一创建入口 │ │ ├── kda_state.py # KDAState：Conv / SSM / MTP 与页操作 │ │ ├── kda_pages_pool.py # KDAPagesPool：KDA Conv 历史和方阵递推快照 @@ -399,15 +410,33 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 | 逻辑类型 | 创建入口 | 参数决定的实现 | | --- | --- | --- | | GQA（含 MHA/MQA） | `full_att.create_gqa_state(config, *, kv_type, page_size, device, ...)` | 普通浮点、INT4/INT8 分组量化、per-tensor/per-head 静态 FP8 | -| MLA | `full_att.create_mla_state(config, *, kv_type, packed, page_size, device, ...)` | 普通 MLA、带 indexer 的 packed MLA、packed FP8 MLA | +| MLA | `full_att.create_mla_state(config, kv_type, packed, page_size, device, fp8_format, ...)` | 普通 MLA、packed MLA、group128/FP32 或 group64/E8M0 packed FP8 MLA | +| 序列压缩 attention | `full_att.create_compressed_att_state(config, page_size, device)` | FP8/E8M0 packed 压缩 KV 与可选 indexer；V4 CSA/C4、HCA/C128 由模型分别配置 | | KDA | `bounded_att.create_kda_state(config, request_capacity, mtp_step=0, *, device)` | KDA Conv/SSM 运行态 | | GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | +| SWA | `bounded_att.create_sliding_window_state(config, request_capacity, mtp_step, device)` | 近期原始 KV 窗口，当前采用 FP8/E8M0 packed 编码 | +| 窗口压缩器 | `bounded_att.create_window_compressor_state(config, request_capacity, mtp_step, device)` | 按窗口长度和产出步长管理候选/gate，KV 和 indexer 分别创建实例 | | K-pool | `att_tail.create_kpool_state(config, request_capacity, mtp_step=0, *, device)` | 带池化尾部窗口及 MTP 历史槽位的 K-pool State | -每个算法的创建入口独立放在 `create_gqa_state.py`、`create_mla_state.py`、`create_kda_state.py`、`create_gdn_state.py` 和 `create_kpool_state.py`，通过各分类包导出。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 +每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器和 K-pool。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 +MLA 的 packed FP8 编码由 `fp8_format` 显式指定：默认 `"group128_fp32"` 使用原有 +`PackedMLAFullAttentionState`；`"group64_e8m0"` 使用独立的 `FP8E8M0PackedMLAState`， +其配置为 `FP8E8M0PackedMLAConfig`。创建入口不通过 `head_dim == 512` 等维度条件推断格式。 +GLM main/draft 在各自的模型工厂中传入 `fp8_format="group64_e8m0"`；未量化时仍返回普通 packed State。 + +`FP8E8M0PackedMLAState` 固定使用 448 维 FP8、64 维 BF16 和七个 E8M0 scale（补齐为八字节）， +KV 共 584 字节，随后保存 indexer key 及一个 FP32 scale；最后 64 维保留 BF16 不代表使用 RoPE。 +配置独立定义，不继承 `FullAttentionConfig`；只记录逻辑层分组、indexer 维度和 TP 坐标。 +固定 latent 维度、单 head、副本规则和 uint8 编码由实现确定，不作为重复输入。 +`token_bytes` 是含 16 字节对齐 padding 的物理行宽，`payload_bytes` 是不含行尾 padding 的有效字节数。 +例如 indexer 维度为 128 时，GPU 行宽为 720 字节，CPU/PD 每个 token 只传 716 字节。 +State 将连续的 KV/scale/indexer 字节作为一个传输区域，复用 `transfer_packed_pages` 的索引拷贝 kernel； +同卡 token 复制与 DP 迁移仍复制完整物理行。原 packed State 及其传输实现无需识别 E8M0 编码。 +定制量化写入放在 `fp8_e8m0_packed_mla_write.py`，State 不依赖 GLM 模型目录中的实现。 + 各模型的三类 State 工厂负责构造配置并调用对应算法入口；`lightllm/utils/cpu_cache_layout.py` 复用模型工厂创建 meta State，生成 CPU cache 联合页布局，不加载标定。full-attention 的实际存储由管理器后续统一分配；GPU 启动时 bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 ```python @@ -480,7 +509,7 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 #### 逻辑层与物理存储层 -`FullAttentionConfig.logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 +具体 State 配置中的 `logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。GQA/MLA 使用 `FullAttentionConfig` 及其具体配置，序列压缩历史使用独立的 `CompressedAttentionConfig`，不要求所有配置继承同一基类。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 ```python logical_layer_groups = ( @@ -518,7 +547,9 @@ FullAttentionState ├── GroupQuantizedFullAttentionState # INT4 / INT8 等 ├── StaticFP8FullAttentionState ├── MLAFullAttentionState -└── PackedMLAFullAttentionState # 附带 indexer 数据的 packed 布局 +├── PackedMLAFullAttentionState # 普通精度与 group128/FP32-scale 的 packed 布局 +├── FP8E8M0PackedMLAState # 独立 group64/E8M0 packed MLA，供 GLM main/draft 使用 +└── CompressedAttentionState # V4 的 C4/C128 压缩块 ``` 所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。`FullAttentionConfig` 与抽象接口共同定义在 `full_att/base_state.py`，量化和 packed 专属配置可扩展该配置；配置结构的继承不代表 State 实现的继承。 @@ -530,7 +561,7 @@ FullAttentionState #### 必须提供的公共接口 ```python -# ViewT / WriteT 为具体布局定义的类型，下面是公共签名草案。 +# ViewT 为具体布局定义的类型，下面是公共签名草案；构造参数由各具体 State 定义。 class FullAttentionState: # ======================================================================== # 一、资源定义与初始化 @@ -538,13 +569,6 @@ class FullAttentionState: @abstractmethod def get_state_type(self) -> str: ... - def __init__( - self, layout_config: FullAttentionConfig, - page_size: int = 1, *, - device: str | torch.device = "cuda", - ) -> None: - ... - def allocate_storage(self, token_capacity: int) -> None: ... @@ -561,7 +585,8 @@ class FullAttentionState: def get_layer_view(self, logical_layer_id: int) -> ViewT: ... - def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: WriteT) -> None: + @abstractmethod + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: ... def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -625,7 +650,7 @@ class FullAttentionState: | `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | | `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | | `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | -| `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | +| `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将 KV 写入已分配位置，完成该方案要求的编码及关联参数更新；融合压缩路径不调用本接口，具体实现显式拒绝误调用 | | `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | 容量、device、层映射及布局标识应从同一份稳定配置和描述中取得，避免多处维护互相矛盾的副本。配置相同且只改变容量时，容量计算和真实 buffer 创建应覆盖相同的数据组成。共享底层内存按实际分配去重记账。 @@ -670,19 +695,24 @@ manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字 带 indexer 的 packed 布局中，attention 与 indexer 可以分别读取 buffer 的不同区域。因此不能用 attention 的读取视图代替完整复制范围。类似地，静态 FP8 的物理存储 dtype 和计算时解释的 dtype 可以不同，视图必须明确相应参数。 -`write_full_att_kv` 更新其输入所描述的 KV 及关联字段；indexer 池化等专用 kernel 可以通过专用视图写入另一区域。管理器与执行流程在相关写入都满足同步条件后才允许复用或发布。融合写入路径可以直接使用预先绑定的布局视图，同时遵守相同的数据表示与完整性约定。 +`write_full_att_kv` 保留为 `FullAttentionState` 的抽象接口，各具体 State 实现该方法。 +压缩历史等复杂布局通过自己的计算视图或定制接口完成融合写入，不调用普通 KV 写入接口； +其实现显式抛出 `NotImplementedError`，避免误调用时静默丢失写入。 +indexer 池化等专用 kernel 可以通过专用视图写入另一区域。 +管理器与执行流程在相关写入都满足同步条件后才允许复用或发布。 #### 基类与子类的分工 -| 部分 | 基类可复用逻辑 | 子类必须负责的内容 | +| 部分 | 基类约定 | 具体 State 负责的内容 | | --- | --- | --- | -| 初始化 | 生命周期、配置与映射校验、稳定视图管理 | `_create_storage(...)`：实际 Tensor/区域和完整描述 | -| 层访问 | 逻辑层到存储层的映射、视图缓存 | `_make_layer_view(storage_layer_id)`：该布局的计算视图 | -| 写入 | 参数契约、层映射和公共调用入口 | `_write_full_att_kv(...)`：编码/量化/写入 kernel 及输入格式校验 | -| 完整复制 | 索引操作契约、公共调用入口 | `_copy_full_att_tokens(...)`：布局正确的复制，或显式复用与其描述兼容的公共工具 | -| 传输 | 定义 CPU/DP/PD 直接读写接口，未支持的操作显式报错 | 覆盖所支持的读写方法，直接使用自身 storage 调用相应 kernel | +| 初始化 | `allocate_storage`、`get_token_bytes` | 强类型构造配置、层映射和实际 Tensor 分配 | +| 层访问 | `contains_logical_layer`、`get_storage_layer_index`、`get_layer_view` | 自身布局的层映射和计算视图 | +| 写入 | `write_full_att_kv` | 实现普通写入；融合压缩路径使用定制方法或视图，普通入口显式拒绝误调用 | +| 完整复制 | `copy_full_att_tokens` 的索引与完整性契约 | 直接实现布局正确的复制，或调用兼容的公共 kernel | +| 传输 | CPU/DP/PD 直接读写契约 | 直接使用自身 storage 和具体 kernel，不从其他具体 State 继承行为 | -受保护方法名为候选。一个子类可以复用满足其布局要求的 kernel；通用复制工具只有在理解 token 维、stride、数据区域和别名关系时才适用。复杂 packed 布局应由专用实现处理，公共层不尝试仅凭 Tensor shape 推导复制语义。 +具体 State 可以复用满足其布局要求的 kernel；基类不规定构造参数或内部受保护方法。 +通用复制工具只有在理解 token 维、stride、数据区域和别名关系时才适用。复杂 packed 布局由具体 State 保证复制粒度与字段完整性。 #### 索引、异步执行与生命周期 @@ -1242,7 +1272,7 @@ DP 副本之间使用相同的 State 成员名。`copy_from_dp` 接收远端 man full-attention 使用 `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_pages_from_dp` 与 `copy_runtime_state_from_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 -full-attention DP 的 `dp_transfer.py` 提供显式存储入口：`copy_gqa_tokens` 接收 K/V，供 GQA 和静态 FP8 使用；`copy_group_quantized_tokens` 另接收 K/V 动态 scale；`copy_mla_tokens` 接收完整 latent/packed 存储。三种入口共用四维 token 布局的 Triton kernel，一次 launch 搬运当前 State 的全部字段；编译时展开 Tensor tuple 为指针参数，不创建或上传指针表。按两端 token 索引及各自真实 stride 直接寻址，无 `index_select`、跨设备 payload `.to()` 或 `index_copy_` 中间存储。每个字段独立选择对齐的整数搬运宽度，支持两端不同 token 容量、非连续索引和存储视图；默认 grid=24、BLOCK=4096、warps=8、stages=1，通过跳步循环覆盖全部元素，地址计算使用 int64。静态 FP8 标定参数不搬运；packed MLA 按原 DP 语义复制完整行，包括 padding，区别于 CPU/PD 的有效字段页格式。调用方保证目标索引唯一、源目标区域不重叠，并负责远端 peer/IPC 映射、源数据就绪及存储生命周期；算子在目标设备当前流执行，不在内部建立 peer 映射或同步。 +full-attention DP 的 `dp_transfer.py` 提供显式存储入口：`copy_gqa_tokens` 接收 K/V，供 GQA 和静态 FP8 使用；`copy_group_quantized_tokens` 另接收 K/V 动态 scale；`copy_packed_tokens` 接收完整 latent/packed 存储。三种入口共用四维 token 布局的 Triton kernel，一次 launch 搬运当前 State 的全部字段；编译时展开 Tensor tuple 为指针参数，不创建或上传指针表。按两端 token 索引及各自真实 stride 直接寻址，无 `index_select`、跨设备 payload `.to()` 或 `index_copy_` 中间存储。每个字段独立选择对齐的整数搬运宽度，支持两端不同 token 容量、非连续索引和存储视图；默认 grid=24、BLOCK=4096、warps=8、stages=1，通过跳步循环覆盖全部元素，地址计算使用 int64。静态 FP8 标定参数不搬运；packed MLA 按原 DP 语义复制完整行，包括 padding，区别于 CPU/PD 的有效字段页格式。调用方保证目标索引唯一、源目标区域不重叠，并负责远端 peer/IPC 映射、源数据就绪及存储生命周期；算子在目标设备当前流执行，不在内部建立 peer 映射或同步。 DP 拷贝在两张 H200 上对照旧 gather → `.to()` → scatter 路径测量：16 层、1025 个乱序 token，源/目标容量为 2048/3072；GQA/INT4/INT8 使用 8 个 KV head、head_dim=128。每种方案和方向预热后连续调用 20 次，三轮取中位数，计时包含 Python 入口及末尾设备同步，不含分配、初始化和编译；重复使用同一组存储，未测试与推理并发。 @@ -1778,6 +1808,13 @@ P/D 传输进程为每次 KV 任务创建普通 CPU token 索引 Tensor，在 `c CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_page_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 +各具体 State 的 `get_cpu_transfer_regions()` 使用 `@lru_cache(maxsize=None)` 缓存不可变区域元组， +并在其外层添加 `@deepcopy_result`，每次先取得缓存结果、再返回深拷贝，调用方不持有缓存原对象。 +装饰器保留 `cache_clear`、`cache_info` 和 `cache_parameters`，管理器仍可清理缓存。 +无需 CPU 快照时返回空元组。首次查询后 State 配置及全局 CPU 页 token 数保持不变； +当前各实现的区域不依赖存储容量或设备；为兼容定制布局受实际物理页规格影响的情况， +管理器完成存储分配后仍统一清除 State 与聚合布局缓存。子类覆盖此 getter 时须保留缓存能力。 + State 的全局存储字节数查询继续描述存储开销，不用于推导 CPU/PD 传输容量。CPU 联合页不预留页头、长度或版本字段。布局及其派生页大小均按需缓存，实际存储分配后一起清理，避免继续使用 meta 准备阶段的布局尺寸。 State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_size` 及具体分布规则，本地数量使用只读属性计算，不保存第二份可独立修改的数值。CPU/PD 全局页和 DP 运行态映射直接读取对应 State 的 config，管理器不维护 `_full_head_counts`、`_bounded_head_counts`。每个 draft 的 State 独立保存完整 config，本地形状相同也不意味着全局布局相同，容量规划和传输均读取各自配置。 @@ -1793,6 +1830,8 @@ full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例 DFlash 先合并 `dflash_config`，再同步 `num_hidden_layers` / `n_layer` 等别名，避免运行时与 CPU 页布局使用不同层数。GLM 原生 MTP 的有效层数描述 draft 自身；`mtp_checkpoint_layer_start` 单独保留 checkpoint 中接在 main 层后的权重起点，用于 transformer 和 pre/post 权重加载，重复规范化不会改变这个起点。 +V4 同样用 `num_hidden_layers` / `n_layer` 表示当前模型自有层数，DSpark 不另存一份层数，也不重复覆盖运行时的 `layers_num`。DSpark 的 `draft_layer_start_id` 表示整体执行序列中的 draft 起始逻辑层号；`stage_id` 由当前逻辑层号减去该起点得到，权重前缀使用 `mtp.{stage_id}`。普通 V4 MTP 的 `draft_weight_layer_id` 则用于权重构造和量化配置索引，各 draft 实例的执行逻辑层号仍由启动流程独立确定。两种起点语义分开命名，重复规范化不改变已记录的值。V4 模型组合配置统一使用 `compress_ratios` 表示逐层压缩倍率。 + ```mermaid flowchart LR A[模型启动读取配置] --> C[模型类 normalize_model_config] @@ -1881,6 +1920,79 @@ CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大 DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并在对齐边界重置目标 tail。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 +### 8.1 DeepSeek V4：压缩 KV 与有界窗口 + +V4 沿用共享 `AttentionStateManager`，不引入专用 ReqManager 或 MemoryManager。各模型通过 +`create_full_att_states` / `create_bounded_att_states` 创建自己的定义，注册完成后统一分配容量。 +main 与 draft0..draft35 使用相同 token 分配域，但分别持有存储；SWA prefill workspace 也按模型角色区分。 + +模型组合配置保留在 `models/deepseek_v4/config.py`，不放入原子 State 目录。 +组合配置的 `create_swa_config()` / `create_compressed_config()` 负责选择参与层及算法参数， +State 创建、backend 和提前预留 workspace 共用这些配置入口。编码字节数、packed 页布局及 +`get_attention_view()` 均由对应原子配置实现，backend 不再维护一套物理格式计算。 +SWA 使用独立的 `FP8E8M0SlidingWindowConfig`，只记录逻辑层、近期窗口、KV 编码与 TP 信息。 +压缩器使用独立的 `WindowCompressorStateConfig`，明确记录 `window_size`、`compress_stride`、 +`candidate_dim`、`target`、精度和 TP 信息；没有 C4/C128 层列表或 SWA 字段。 +`window_size` 是每次聚合覆盖的原始 token 数，`compress_stride` 是每隔多少原始 token 产出一条， +大小页需保留的跨步历史由 `overlap_tokens = window_size - compress_stride` 推导，不额外存储重复配置。 + +| V4 创建的运行态实例 | 窗口 / 产出步长 | 每 token 候选宽度 | 大小页快照窗口 | PD 窗口 | +| --- | --- | --- | --- | --- | +| SWA | 配置的 sliding window / 每 token 更新 | packed KV 字节 | 完整 sliding window | 完整 sliding window | +| CSA KV compressor | W8 / S4 | `2 * head_dim`，gate 另占同宽 | 上一组 4 个位置 | 8 个位置 | +| CSA indexer compressor | W8 / S4 | `2 * index_head_dim`，gate 另占同宽 | 上一组 4 个位置 | 8 个位置 | +| HCA KV compressor | W128 / S128 | `head_dim`，gate 另占同宽 | 组边界无跨步历史，不占页池和 CPU 页 | 128 个位置 | + +规范名称包含用途、窗口、步长及精度，例如 `main_compressor_kv_w8_s4_fp32_state`、 +`main_compressor_indexer_w8_s4_fp32_state`、`main_compressor_kv_w128_s128_fp32_state`， +SWA 为 `main_swa_w128_fp8_group64_e8m0_state`。同一个压缩器实现也可以保存 W12/S4 等其他窗口， +压缩计算仍由对应模型算子负责，不在 State 中根据窗口或比例猜测算法。 + +各 State 分别拥有自己的大小页池，共享同一组 slot ID 和 allocator；不保留 V4 聚合 State 或联合 pool。 +无快照的压缩器返回空 CPU regions，manager 不为其分配联合页区域，但保存/恢复仍校验已完成产出步长。 +缺少其他参与 State 的 CPU layout 仍会显式失败。标准 V4 MTP draft 只创建 SWA State。 + +| 组件 | 归属 | 数据与索引 | +| --- | --- | --- | +| `CompressedAttentionState` | full-attention | 分别保存 C4/C128 层的压缩 KV；C4 同块保存 indexer key/scale，外部索引仍以原始 token 为单位 | +| `FP8E8M0SlidingWindowState` | bounded-attention | 只保存近期原始 KV 的 SWA ring；不依赖压缩器或 V4 模型配置 | +| `WindowCompressorState` | bounded-attention | 每实例只保存一种窗口/步长/用途的候选与 gate；不保存压缩历史或 SWA | +| `WindowPagesPool` | bounded 大小页池 | 每个参与 State 各自保存单一窗口快照；共享 manager 的槽位分配域，不维护请求或 radix 引用 | +| `ReqManager` | 请求管理 | 只保存请求槽位和逻辑 token 到物理 token 的映射 | + +压缩历史通过独立的 `create_compressed_att_state()` 创建，不经过 MLA 创建入口。 +`CompressedAttentionConfig` 是独立的原子布局定义，不继承 `FullAttentionConfig`，也不持有 +`DSV4AttentionConfig`。它直接保存逻辑层分组 `logical_layer_groups`、压缩比例 `compress_ratio`、 +KV/RoPE 维度 `kv_dim` / `rope_dim`、每块覆盖的原始 token 数 `block_token_size`、可选 +`indexer_dim` 和当前 TP 信息。`indexer_dim=0` 表示不保存 indexer，不根据压缩比例推断。 +V4 模型负责将自己的组合配置转换为这些字段；State 不读取 SWA 窗口、全模型层列表或压缩器配置。 +当前实现固定使用一个完整复制的共享 KV head，使用 uint8 保存 FP8/BF16/scale 的原始编码， +不再要求调用方传入 `global_head_num=1`、`replicated=True`、`dtype=torch.uint8` 等固定参数。 +配置的 `token_slot_bytes` 按压缩比例推导每个原始 token 槽位摊销的物理字节数，包含 indexer 和 padding， +不借用 `head_dim` 表示字节数。State 的名称为 `compressed_c{压缩比例}_fp8_group64_e8m0`， +例如 `main_compressed_c4_fp8_group64_e8m0_state`。CSA/HCA 描述计算与历史选择流程， +不作为存储类的固定身份；同一实现按原子字段分别承载它们的压缩历史。 +它保存已完成的序列压缩结果，不保存 SWA 或 compressor 运行态；State 仍直接继承 `FullAttentionState`， +CPU/DP/PD 继续复用适用于该 packed 字节布局的索引传输 kernel，不改变物理格式。 + +压缩 KV 以 256 个原始 token 为一个物理块，块内数据与 scale 分区存放；C4 还包含 indexer。 +PD 的最后一个碎页必须携带整个已分配压缩块，避免遗漏有效 token 对应的 scale/indexer。 +CPU 页面、KV 分配页和前缀边界须遵守相应的压缩块对齐约束。 + +大小页保存/恢复接收绝对 `state_token_lens`，由 State 将请求 ring 与时间有序快照相互转换。 +CPU 读回仍先写入大页池，再恢复请求运行态；最后一个碎页使用管理器的 LOAD 暂存槽位。 +C128 在缓存边界已经关闭当前压缩组，无重叠历史需要保存；V4 的状态页边界因此必须是 128 的倍数。 +C4 保存上一组的四个位置，供下一组重叠压缩使用。运行态与大小页之间直接调用 V4 kernel， +不使用外部 `TransferRegion`；仅 CPU/PD 联合页使用 region 描述字段与 TP 写入权限。 + +PD 运行态采用与本地 MTP 容量无关的逻辑窗口：SWA 为配置窗口，C4/indexer 为八个位置， +C128 为 128 个位置。收发两端根据传入的绝对前缀长度各自计算 ring 位置，因此允许不同的 MTP 配置。 +SWA、压缩 KV 和 compressor 在 TP 内完整复制，仅 rank 0 写出共享外部区域，各 rank 可读回。 +V4 的这些窗口需要真实快照，不归入仅能在页边界清零的 `AttentionTailState`。 + +checkpoint 中的权重层号与 draft 自有逻辑层数分别记录；多个复用相同 MTP 权重的 draft 仍注册独立 State。 +模型和 backend 根据逻辑层查询 State 视图，不在 ReqManager 中附加 ring，不在 backend 中保存运行态副本。 + ## 9. 待细化的设计问题 | 问题 | 设计约定 | 待确定内容 | diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py index 19f96ce1a3..8185379fd5 100644 --- a/lightllm/common/attention_state/bounded_att/__init__.py +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -7,13 +7,24 @@ from .kda_state import KDAState from .base_state import BoundedAttentionState +from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState +from .window_compressor_state import WindowCompressorStateConfig, WindowCompressorState +from .create_sliding_window_state import create_sliding_window_state +from .create_window_compressor_state import create_window_compressor_state + __all__ = [ - "create_kda_state", - "create_gdn_state", + "BoundedAttentionState", "ConvSSMStateConfig", "GDNState", - "KDAState", - "BoundedAttentionState", "GDNPagesPool", + "create_gdn_state", + "KDAState", "KDAPagesPool", + "create_kda_state", + "FP8E8M0SlidingWindowConfig", + "FP8E8M0SlidingWindowState", + "create_sliding_window_state", + "WindowCompressorStateConfig", + "WindowCompressorState", + "create_window_compressor_state", ] diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index d65d6bba83..689b61b60b 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -2,12 +2,16 @@ from __future__ import annotations +from functools import lru_cache + + from typing import Any from abc import ABC, abstractmethod from typing import List import torch +from ..utils import deepcopy_result from lightllm.common.attention_state.types import ( StateTransferLayout, PageKind, @@ -114,11 +118,16 @@ def commit_bounded_att_state( # ======================================================================== @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回一份 CPU 快照的字段区间、全局形状与 TP 切片。 byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； 参与 CPU 页时，管理器按字段最大终点计算占用，不乘页内 token 数。 + 在合法缓存边界无需保留历史的 State 返回空元组，不占用 CPU 页面。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 管理器分配后统一清理，子类覆盖时须保留缓存和副本语义。 """ ... diff --git a/lightllm/common/attention_state/bounded_att/create_sliding_window_state.py b/lightllm/common/attention_state/bounded_att/create_sliding_window_state.py new file mode 100644 index 0000000000..c9f3af28bb --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/create_sliding_window_state.py @@ -0,0 +1,19 @@ +"""FP8E8M0SlidingWindowState 的独立创建入口。""" + +import torch + +from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState + + +def create_sliding_window_state( + config: FP8E8M0SlidingWindowConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> FP8E8M0SlidingWindowState: + return FP8E8M0SlidingWindowState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/create_window_compressor_state.py b/lightllm/common/attention_state/bounded_att/create_window_compressor_state.py new file mode 100644 index 0000000000..e80d802925 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/create_window_compressor_state.py @@ -0,0 +1,19 @@ +"""WindowCompressorState 的独立创建入口。""" + +import torch + +from .window_compressor_state import WindowCompressorStateConfig, WindowCompressorState + + +def create_window_compressor_state( + config: WindowCompressorStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> WindowCompressorState: + return WindowCompressorState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py new file mode 100644 index 0000000000..00a7862f1b --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py @@ -0,0 +1,365 @@ +"""单个共享 KV head 的 FP8/E8M0 滑动窗口运行态，不依赖模型或 compressor。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from math import prod + +import torch + +from ..utils import deepcopy_result +from ..types import PageKind, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .copy_kernel import copy_bounded_att_tensors +from .transfer_utils import bounded_att_buffer_copies +from .window_pages_pool import WindowPagesPool +from .fp8_e8m0_swa_transfer import copy_swa_window + + +@dataclass(frozen=True) +class FP8E8M0SlidingWindowConfig: + # 本 State 服务的逻辑层，顺序对应物理 layer 轴，各层拥有独立窗口。 + logical_layer_ids: tuple[int, ...] + # 每个 query 可见的近期 token 数，也是大小页和 PD 快照的逻辑窗口长度。 + window_size: int + # 未编码 KV 向量维度，包含末尾保留 BF16 的 RoPE 部分。 + kv_dim: int + # RoPE 维度；其余维度按 group64 FP8 编码，scale 使用 E8M0。 + rope_dim: int + # packed 物理页内的 token 数；页内 KV 数据与 scale 分区存放。 + page_token_size: int = 128 + # 各 TP rank 完整复制 KV；仅 rank 0 写出共享 CPU/PD 区域。 + tp_rank: int = 0 + tp_world_size: int = 1 + + @property + def kv_data_bytes(self) -> int: + return self.kv_dim + self.rope_dim + + @property + def kv_bytes(self) -> int: + scales = (self.kv_dim - self.rope_dim) // 64 + return self.kv_data_bytes + (scales + 7) // 8 * 8 + + def get_packed_kv_bytes(self, token_count: int) -> int: + """返回 packed 页字节数，包含数据行对齐。""" + return (token_count * self.kv_bytes + self.kv_data_bytes - 1) // self.kv_data_bytes * self.kv_data_bytes + + def get_attention_view(self, packed_kv: torch.Tensor) -> torch.Tensor: + """将运行态或 workspace 解释为 [page, token, head, byte]，不复制或重排数据。""" + pages = packed_kv.view(-1, self.get_packed_kv_bytes(self.page_token_size)) + return pages[:, : self.page_token_size * self.kv_bytes].view(-1, self.page_token_size, 1, self.kv_bytes) + + def __post_init__(self) -> None: + assert self.window_size > 0 and self.page_token_size > 0 + assert self.kv_dim > self.rope_dim >= 0 and (self.kv_dim - self.rope_dim) % 64 == 0 + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert 0 <= self.tp_rank < self.tp_world_size + assert self.get_packed_kv_bytes(self.page_token_size) % self.page_token_size == 0 + + +class FP8E8M0SlidingWindowState(BoundedAttentionState): + """仅保存近期原始 KV；prefill workspace 由计算 backend 管理。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: FP8E8M0SlidingWindowConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + self.config = config + self.state_name: str | None = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + page_tokens = config.page_token_size + self.swa_capacity_tokens = (config.window_size + 2 * mtp_step + page_tokens - 1) // page_tokens * page_tokens + self.HOLD_REQUEST_ID = request_capacity - 1 + self.HOLD_SWA_SLOT = self.HOLD_REQUEST_ID * self.swa_capacity_tokens + self.swa_buffer = torch.zeros( + size=( + len(config.logical_layer_ids), + request_capacity, + self.swa_capacity_tokens, + config.get_packed_kv_bytes(page_tokens) // page_tokens, + ), + dtype=torch.uint8, + device=device, + ) + self.big_pages_pool: WindowPagesPool | None = None + self.small_pages_pool: WindowPagesPool | None = None + + def get_state_type(self) -> str: + """窗口长度参与方案名称，允许同一模型组合不同滑窗。""" + return f"swa_w{self.config.window_size}_fp8_group64_e8m0" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """大小页只保存时间有序 KV，不保存 ring padding 或 MTP 候选位置。""" + shape = (len(self.config.logical_layer_ids), self.config.window_size, self.config.kv_bytes) + self.big_pages_pool = ( + WindowPagesPool( + size=big_page_capacity, + snapshot_shape=shape, + dtype=torch.uint8, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + WindowPagesPool( + size=small_page_capacity, + snapshot_shape=shape, + dtype=torch.uint8, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、运行态：层视图、重置与大小页保存恢复 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: + """返回该层的 [request, ring_token, packed_byte] 视图。""" + return self.swa_buffer[self._logical_to_storage_layer[logical_layer_id]] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + for req_idx in req_indices: + self.swa_buffer[:, req_idx].zero_() + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=True, + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=False, + ) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """ring 按绝对 token 位置寻址，无需提交额外游标。""" + pass + + # ======================================================================== + # 三、CPU Cache：仅搬运时间有序的大小页快照 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """CPU 快照保存完整近期窗口。""" + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定快照写入 CPU 页，已就绪页保持不变。""" + pool = self._pool(page_kind=page_kind) + fields, pages = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + fields.append({"swa": TensorField(tensor=pool.buffer[slot], axes=("layer", "window", "channel"))}) + pages.append(page) + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=pages, + write=True, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """CPU 页先恢复到大页池，随后再按前缀长度恢复请求 ring。""" + fields = [ + {"swa": TensorField(tensor=self.big_pages_pool.buffer[slot], axes=("layer", "window", "channel"))} + for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP：相同方案的快照及请求运行态复制 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: list[int], + dst_page_ids: list[int], + remote_dp_bounded_state: FP8E8M0SlidingWindowState, + ) -> None: + """复制已保存的大页，调用方保持槽位及 stream 依赖一致。""" + assert self.config == remote_dp_bounded_state.config + if self.big_pages_pool is not None: + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.big_pages_pool.buffer[src], self.big_pages_pool.buffer[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: FP8E8M0SlidingWindowState, + ) -> None: + """直接复制相同配置和 MTP 容量的物理 ring，不创建快照。""" + assert self.config == remote_dp_bounded_state.config and self.mtp_step == remote_dp_bounded_state.mtp_step + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.swa_buffer[:, src], self.swa_buffer[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + # ======================================================================== + # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回完整逻辑窗口，不包含本地 MTP 候选区。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按绝对长度将请求 ring 写入选定 PD 页。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + if region.write_enabled: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按接收端 ring 容量还原逻辑窗口，不写入大小页池。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=False, + ) + + # ======================================================================== + # 六、内部页池和窗口搬运 + # ======================================================================== + + def _pool(self, page_kind: PageKind) -> WindowPagesPool: + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled" + return pool + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + shape = (len(self.config.logical_layer_ids), self.config.window_size, self.config.kv_bytes) + return ( + TransferRegion( + field="swa", + byte_start=0, + byte_end=prod(shape), + dtype=torch.uint8, + axes=("layer", "window", "channel"), + shape=shape, + axis_mappings=(), + write_enabled=self.config.tp_rank == 0, + ), + ) + + def _copy_window( + self, + req_idx: int, + external: torch.Tensor, + state_token_len: int, + write: bool, + ) -> None: + copy_swa_window( + local=self.swa_buffer[:, req_idx], + external=external, + state_token_len=state_token_len, + config=self.config, + write=write, + ) diff --git a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py new file mode 100644 index 0000000000..49774a5c87 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py @@ -0,0 +1,97 @@ +"""FP8/E8M0 SWA packed ring 与时间有序快照的 GPU 搬运。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import triton +import triton.language as tl + +if TYPE_CHECKING: + from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig + + +def copy_swa_window( + local: torch.Tensor, + external: torch.Tensor, + state_token_len: int, + config: FP8E8M0SlidingWindowConfig, + write: bool, +) -> None: + """按绝对长度在 packed SWA ring 和 [layer, window, byte] 快照间双向复制。 + + local 是选定请求的视图。external 可为已 pinned/registered 的大小页或 CUDA PD 页； + 直接由 GPU 读写，不分配 gather Tensor，不在主机等待拷贝完成。 + SWA 的数据和 scale 分区存储，不能直接按 token 行 memcpy。 + """ + elements = external.numel() + _copy_swa_window[(min(24, triton.cdiv(elements, 1024)),)]( + LOCAL=local, + EXTERNAL=external, + window_start=state_token_len - external.shape[1], + LOCAL_LAYER_STRIDE=local.stride(0), + EXTERNAL_STRIDES=external.stride(), + LAYERS=local.shape[0], + RING_TOKENS=local.shape[1], + WINDOW_TOKENS=external.shape[1], + KV_BYTES=config.kv_bytes, + KV_DATA_BYTES=config.kv_data_bytes, + PAGE_TOKENS=config.page_token_size, + PAGE_BYTES=config.get_packed_kv_bytes(config.page_token_size), + WRITE=write, + BLOCK=1024, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["window_start"]) +def _copy_swa_window( + LOCAL, + EXTERNAL, + window_start, + LOCAL_LAYER_STRIDE: tl.constexpr, + EXTERNAL_STRIDES: tl.constexpr, + LAYERS: tl.constexpr, + RING_TOKENS: tl.constexpr, + WINDOW_TOKENS: tl.constexpr, + KV_BYTES: tl.constexpr, + KV_DATA_BYTES: tl.constexpr, + PAGE_TOKENS: tl.constexpr, + PAGE_BYTES: tl.constexpr, + WRITE: tl.constexpr, + BLOCK: tl.constexpr, +): + # 阶段一：按线性快照拆出 layer/window/byte，绝对窗口起点保持运行时参数。 + lanes = tl.arange(0, BLOCK).to(tl.int64) + count = LAYERS * WINDOW_TOKENS * KV_BYTES + window_start = window_start.to(tl.int64) + for start in range(tl.program_id(0).to(tl.int64) * BLOCK, count, tl.num_programs(0).to(tl.int64) * BLOCK): + index = start + lanes + byte = index % KV_BYTES + window = index // KV_BYTES % WINDOW_TOKENS + layer = index // (KV_BYTES * WINDOW_TOKENS) + position = window_start + window + slot = position % RING_TOKENS + + # 阶段二：SWA 每页先存全部 KV 数据，再存全部 scale，末尾保留物理 padding。 + packed_offset = slot // PAGE_TOKENS * PAGE_BYTES + tl.where( + byte < KV_DATA_BYTES, + slot % PAGE_TOKENS * KV_DATA_BYTES + byte, + PAGE_TOKENS * KV_DATA_BYTES + slot % PAGE_TOKENS * (KV_BYTES - KV_DATA_BYTES) + byte - KV_DATA_BYTES, + ) + local_offset = layer * tl.cast(LOCAL_LAYER_STRIDE, tl.int64) + packed_offset + external_offset = ( + layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) + + window * tl.cast(EXTERNAL_STRIDES[1], tl.int64) + + byte * tl.cast(EXTERNAL_STRIDES[2], tl.int64) + ) + valid = index < count + has_history = valid & (position >= 0) + # TODO：按数据/scale 两个分区扩大搬运单位，融合宽类型访问，保持 packed padding 不变。 + if WRITE: + values = tl.load(LOCAL + local_offset, mask=has_history, other=0) + tl.store(EXTERNAL + external_offset, values, mask=valid) + else: + values = tl.load(EXTERNAL + external_offset, mask=has_history, other=0) + tl.store(LOCAL + local_offset, values, mask=has_history) diff --git a/lightllm/common/attention_state/bounded_att/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn_state.py index 61506911dd..c7d8d75649 100644 --- a/lightllm/common/attention_state/bounded_att/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn_state.py @@ -1,5 +1,6 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod @@ -7,7 +8,7 @@ from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig @@ -180,6 +181,8 @@ def commit_bounded_att_state( # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 diff --git a/lightllm/common/attention_state/bounded_att/kda_state.py b/lightllm/common/attention_state/bounded_att/kda_state.py index 6f8cab3326..312f0dabee 100644 --- a/lightllm/common/attention_state/bounded_att/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda_state.py @@ -1,5 +1,6 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod @@ -7,7 +8,7 @@ from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState from .conv_ssm_types import ConvSSMStateConfig @@ -180,6 +181,8 @@ def commit_bounded_att_state( # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor_state.py new file mode 100644 index 0000000000..4c3c7f983d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor_state.py @@ -0,0 +1,377 @@ +"""固定窗口、固定产出步长的 attention 压缩器候选/gate 运行态。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from math import prod +from typing import Literal + +import torch + +from ..utils import deepcopy_result +from ..types import PageKind, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .copy_kernel import copy_bounded_att_tensors +from .transfer_utils import bounded_att_buffer_copies +from .window_pages_pool import WindowPagesPool +from .window_compressor_transfer import copy_compressor_window + + +@dataclass(frozen=True) +class WindowCompressorStateConfig: + # 本实例服务的逻辑层；KV 与 indexer 的压缩器分别创建实例。 + logical_layer_ids: tuple[int, ...] + # 每次压缩聚合的原始 token 窗口，例如 C4 为 W8，C128 为 W128。 + window_size: int + # 每执行多少个原始 token 产出一个条目，例如 C4 为 S4，C128 为 S128。 + compress_stride: int + # 每个 token 保存的候选向量宽度，gate 另占同样宽度。 + # W8/S4 可为前后半窗口分别保存候选，所以当前 V4 使用 2 * 输出向量维度。 + candidate_dim: int + # 产出的用途用于区分规范 State 成员名；不改变窗口保存和恢复算法。 + target: Literal["kv", "indexer"] = "kv" + # 候选和 gate 的精度，当前压缩算子使用 FP32。 + dtype: torch.dtype = torch.float32 + # 运行态在 TP 内完整复制，只有 rank 0 向共享外部页写出。 + tp_rank: int = 0 + tp_world_size: int = 1 + + @property + def overlap_tokens(self) -> int: + """组完成时仍需跨步保留的重叠历史：W8/S4 为 4，W128/S128 为 0。""" + return self.window_size - self.compress_stride + + def __post_init__(self) -> None: + assert self.window_size >= self.compress_stride > 0 and self.candidate_dim > 0 + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert self.target in ("kv", "indexer") and self.dtype in (torch.float32, torch.float16, torch.bfloat16) + assert 0 <= self.tp_rank < self.tp_world_size + + +class WindowCompressorState(BoundedAttentionState): + """只保存窗口候选/gate，不保存压缩后的 KV 或 SWA。 + + 大小页在产出步长的边界保存重叠历史，PD 则保存完整窗口以支持任意前缀长度。 + 无重叠的窗口在组边界不需要快照，因此不创建大小页池，也不占 CPU 联合页字节。 + """ + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: WindowCompressorStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + self.config = config + self.state_name: str | None = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + self.buffer = torch.zeros( + size=( + len(config.logical_layer_ids), + request_capacity, + config.window_size + mtp_step, + 2 * config.candidate_dim, + ), + dtype=config.dtype, + device=device, + ) + # HOLD 不得成为有效候选；正常请求在 reset 时初始化候选及 gate。 + self.buffer[:, request_capacity - 1, :, config.candidate_dim :].fill_(float("-inf")) + self.big_pages_pool: WindowPagesPool | None = None + self.small_pages_pool: WindowPagesPool | None = None + + def get_state_type(self) -> str: + """用途、窗口、步长和精度共同标识可复用的压缩器方案。""" + precision = {torch.float32: "fp32", torch.float16: "fp16", torch.bfloat16: "bf16"}[self.config.dtype] + return ( + f"compressor_{self.config.target}_w{self.config.window_size}" + f"_s{self.config.compress_stride}_{precision}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """只为跨步重叠历史创建快照池；无重叠时保持 None。""" + if not self.config.overlap_tokens: + return + shape = (len(self.config.logical_layer_ids), self.config.overlap_tokens, 2 * self.config.candidate_dim) + self.big_pages_pool = ( + WindowPagesPool( + size=big_page_capacity, + snapshot_shape=shape, + dtype=self.config.dtype, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + WindowPagesPool( + size=small_page_capacity, + snapshot_shape=shape, + dtype=self.config.dtype, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、运行态：层视图、重置与大小页保存恢复 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: + """返回该层 [request, window+MTP, candidate|gate] 的计算视图。""" + return self.buffer[self._logical_to_storage_layer[logical_layer_id]] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + for req_idx in req_indices: + self.buffer[:, req_idx].zero_() + self.buffer[:, req_idx, :, self.config.candidate_dim :].fill_(float("-inf")) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + """边界必须完成一个产出步长；仅保留下一步仍需读取的窗口重叠部分。""" + assert all( + length % self.config.compress_stride == 0 for length in state_token_lens + ), "prefix must close a compressor stride" + if not self.config.overlap_tokens: + return + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=True, + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + assert all( + length % self.config.compress_stride == 0 for length in state_token_lens + ), "prefix must close a compressor stride" + if not self.config.overlap_tokens: + return + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=False, + ) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" + pass + + # ======================================================================== + # 三、CPU Cache：仅搬运时间有序的大小页快照 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回快照布局；缓存边界无需保留历史的窗口返回空区域。""" + return self._get_transfer_regions(runtime=False) + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定快照写入 CPU 页，已就绪页保持不变。""" + pool = self._pool(page_kind=page_kind) + fields, pages = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + fields.append({"compressor": TensorField(tensor=pool.buffer[slot], axes=("layer", "window", "channel"))}) + pages.append(page) + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=pages, + write=True, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """CPU 页先恢复到大页池，随后再按前缀长度恢复请求 ring。""" + fields = [ + {"compressor": TensorField(tensor=self.big_pages_pool.buffer[slot], axes=("layer", "window", "channel"))} + for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP：相同方案的快照及请求运行态复制 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: list[int], + dst_page_ids: list[int], + remote_dp_bounded_state: WindowCompressorState, + ) -> None: + """复制已保存的大页，调用方保持槽位及 stream 依赖一致。""" + assert self.config == remote_dp_bounded_state.config + if self.big_pages_pool is not None: + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.big_pages_pool.buffer[src], self.big_pages_pool.buffer[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: WindowCompressorState, + ) -> None: + """直接复制相同配置和 MTP 容量的物理 ring，不创建快照。""" + assert self.config == remote_dp_bounded_state.config and self.mtp_step == remote_dp_bounded_state.mtp_step + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.buffer[:, src], self.buffer[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + # ======================================================================== + # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回完整逻辑窗口，不包含本地 MTP 候选区。""" + return self._get_transfer_regions(runtime=True) + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按绝对长度将请求 ring 写入选定 PD 页。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + if region.write_enabled: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按接收端 ring 容量还原逻辑窗口，不写入大小页池。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=False, + ) + + # ======================================================================== + # 六、内部页池和窗口搬运 + # ======================================================================== + + def _pool(self, page_kind: PageKind) -> WindowPagesPool: + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled" + return pool + + def _get_transfer_regions(self, runtime: bool) -> tuple[TransferRegion, ...]: + window = self.config.window_size if runtime else self.config.overlap_tokens + if not window: + return () + shape = (len(self.config.logical_layer_ids), window, 2 * self.config.candidate_dim) + return ( + TransferRegion( + field="compressor", + byte_start=0, + byte_end=prod(shape) * self.config.dtype.itemsize, + dtype=self.config.dtype, + axes=("layer", "window", "channel"), + shape=shape, + axis_mappings=(), + write_enabled=self.config.tp_rank == 0, + ), + ) + + def _copy_window( + self, + req_idx: int, + external: torch.Tensor, + state_token_len: int, + write: bool, + ) -> None: + copy_compressor_window( + local=self.buffer[:, req_idx], + external=external, + state_token_len=state_token_len, + write=write, + ) diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py b/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py new file mode 100644 index 0000000000..43fd11b0bf --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py @@ -0,0 +1,44 @@ +"""窗口压缩器候选/gate ring 与时间有序快照的 GPU 搬运。""" + +import torch +import triton + +from ..utils import view_max_copy_dtype +from ..att_tail.kpool_transfer import _transfer_kpool_pages + + +def copy_compressor_window( + local: torch.Tensor, + external: torch.Tensor, + state_token_len: int, + write: bool, +) -> None: + """在候选/gate ring 与时间有序快照之间复制，窗口长度由 external 的形状决定。 + + 大小页可以只保留跨步重叠历史，PD 保存完整窗口；两者均不包含 MTP 候选区。 + """ + if not local.numel(): + return + local_view, external_view = view_max_copy_dtype(copies=[(local, external)])[0] + window = external.shape[1] + elements = local_view.shape[0] * window * local_view.shape[2] + _transfer_kpool_pages[(min(24, triton.cdiv(elements, 1024)),)]( + LOCAL=local_view, + EXTERNAL=external_view, + local_window_start=(state_token_len - window) % local.shape[1], + external_window_start=0, + invalid_prefix=max(window - state_token_len, 0), + LOCAL_STRIDES=local_view.stride(), + EXTERNAL_STRIDES=external_view.stride(), + LAYERS=local_view.shape[0], + LOCAL_RING_SIZE=local.shape[1], + EXTERNAL_RING_SIZE=window, + WINDOW_SIZE=window, + CHANNELS=local_view.shape[2], + ELEMENT_BYTES=local_view.element_size(), + LOCAL_CIRCULAR=True, + EXTERNAL_CIRCULAR=False, + WRITE=write, + BLOCK_ELEMENTS=1024, + num_warps=4, + ) diff --git a/lightllm/common/attention_state/bounded_att/window_pages_pool.py b/lightllm/common/attention_state/bounded_att/window_pages_pool.py new file mode 100644 index 0000000000..957eef0613 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_pages_pool.py @@ -0,0 +1,27 @@ +"""单种时间有序窗口快照的页池，槽位由 manager 的共享 allocator 管理。""" + +import torch + +from .utils import create_page_pool_tensor, register_page_pool_tensor + + +class WindowPagesPool: + def __init__( + self, + size: int, + snapshot_shape: tuple[int, int, int], + dtype: torch.dtype, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + # [页槽位, 物理层, 时间窗口, 字段宽度]；只保存快照，不保存请求或长度元数据。 + self.buffer = create_page_pool_tensor( + size=(size, *snapshot_shape), + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + + def register_pin_memory(self) -> None: + """跨 DP 导入后，在接收进程当前设备上注册共享页池。""" + register_page_pool_tensor(tensor=self.buffer) diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index b51fc60359..b65340cfb3 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -29,3 +29,17 @@ "PackedMLAAttentionConfig", "PackedMLALayerView", ] + +from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionState, CompressedAttentionLayerView +from .create_compressed_att_state import create_compressed_att_state + +__all__ += [ + "create_compressed_att_state", + "CompressedAttentionConfig", + "CompressedAttentionState", + "CompressedAttentionLayerView", +] + +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + +__all__ += ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index a4772e52fa..83fe4efd8b 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -1,5 +1,8 @@ from __future__ import annotations +from functools import lru_cache + + from typing import List from abc import ABC, abstractmethod from dataclasses import dataclass @@ -7,6 +10,7 @@ import torch +from ..utils import deepcopy_result from lightllm.common.attention_state.types import ( StateTransferLayout, TransferRegion, @@ -98,22 +102,13 @@ def get_state_type(self) -> str: - GQA 分组量化：gqa_int{4|8}_group{分组大小}_{scale精度}_scale。 - GQA 静态 FP8：gqa_fp8_static_{per_tensor|per_head}。 - MLA 浮点：mla_{精度}；KV/indexer 联合存储：mla_packed_{精度}。 - - MLA packed FP8：mla_packed_fp8_group{分组大小}，当前内核固定为 group128。 + - MLA packed FP8：mla_packed_fp8_group128，或 mla_packed_fp8_group64_e8m0。 + - 序列压缩历史：compressed_c{压缩比例}_fp8_group64_e8m0。 例如 gqa_int8_group128_bf16_scale 对应 main_gqa_int8_group128_bf16_scale_state。 新量化实现必须区分影响方案选择的位宽、分组、scale 方式及精度，不编码标定数值。 """ ... - @abstractmethod - def __init__( - self, - layout_config: FullAttentionConfig, - page_size: int = 1, - device: str | torch.device = "cuda", - ) -> None: - """只准备 State 配置，不分配 KV 或 HOLD 页；存储由后续统一容量初始化。""" - raise NotImplementedError() - @abstractmethod def allocate_storage(self, token_capacity: int) -> None: """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" @@ -145,7 +140,10 @@ def get_layer_view(self, logical_layer_id: int) -> Any: @abstractmethod def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: - """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。""" + """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。 + + 融合压缩等专用写入路径可直接使用计算视图，不调用本接口。 + """ raise NotImplementedError() @abstractmethod @@ -160,12 +158,17 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # ======================================================================== @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 cpu_cache_token_page_size 返回 CPU 字段、全局形状和 TP 切片。 byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； 管理器按字段最大终点计算该 State 的 CPU 页占用。 本次有效 token 数由读写流程处理，不改变布局。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 CPU 页规格保持不变。 + 管理器实际分配存储后统一清除缓存，以兼容物理页规格变化；子类覆盖时须保留缓存能力。 """ ... diff --git a/lightllm/common/attention_state/full_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att_state.py new file mode 100644 index 0000000000..74722526b4 --- /dev/null +++ b/lightllm/common/attention_state/full_att/compressed_att_state.py @@ -0,0 +1,425 @@ +"""序列压缩后的历史 KV，采用 FP8/E8M0、BF16 RoPE 与可选 indexer 的 packed 编码。""" + +from __future__ import annotations + +from functools import lru_cache + +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import lcm, prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ..utils import deepcopy_result +from ..types import AxisMapping, TransferRegion +from .base_state import FullAttentionState +from .packed_transfer import transfer_packed_pages +from .dp_transfer import copy_packed_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +@dataclass(frozen=True) +class CompressedAttentionConfig: + # 外层顺序是物理 layer 轴；内层列出共享整份压缩 KV/indexer 存储的逻辑层。 + # 例如 ((0, 4), (8,)) 分配两份物理层存储，逻辑层 0/4 访问同一份数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 以原始 token 为单位的压缩比例，不改变外部请求映射的索引域。 + compress_ratio: int + # 每个压缩 KV 条目编码前的向量维度，包含末尾的 RoPE 部分；不是存储字节数。 + kv_dim: int + # 保留 BF16 的 RoPE 维度；其余维度按每 64 个元素一组编码为 FP8，scale 使用 E8M0。 + rope_dim: int + # 每个物理块覆盖的原始 token 数，必须能被 compress_ratio 整除。 + # 块内先保存全部 KV 数据，再保存 scale 和可选 indexer；不能按单个压缩条目拆开搬运。 + block_token_size: int + # 可选 indexer key 的 FP8 维度，每个 key 另有一个 FP32 scale；0 表示不保存 indexer。 + indexer_dim: int = 0 + # 当前 State 所属 TP 组内的 rank；当前实现是单个共享 KV head，在各 rank 完整复制。 + # rank 0 写出共享 CPU/PD 页，每个 rank 均可读回，不按 TP 切分压缩条目的向量。 + tp_rank: int = 0 + # 当前模型 TP 组的 rank 数，只描述复制关系，不影响单个 rank 的存储字节数。 + tp_world_size: int = 1 + + @property + def kv_entry_bytes(self) -> int: + """单条压缩 KV 的数据与 scale 字节数，不包含块末尾 padding 或 indexer。""" + scale_count = (self.kv_dim - self.rope_dim) // 64 + return self.kv_dim + self.rope_dim + (scale_count + 7) // 8 * 8 + + @property + def kv_block_bytes(self) -> int: + """块内 KV 数据及 E8M0 scale 区的字节跨度，包含数据行对齐，不包含 indexer。""" + entries = self.block_token_size // self.compress_ratio + # FP8 占一字节，BF16 RoPE 占两字节；每条目的 scale 区补齐到八字节。 + data_bytes = self.kv_dim + self.rope_dim + return (entries * self.kv_entry_bytes + data_bytes - 1) // data_bytes * data_bytes + + @property + def token_slot_bytes(self) -> int: + """每个原始 token 索引摊销的物理字节数，包含 KV、scale、indexer 和块内 padding。""" + block_bytes = self.kv_block_bytes + if self.indexer_dim: + entries = self.block_token_size // self.compress_ratio + block_bytes += entries * (self.indexer_dim + torch.float32.itemsize) + # 块既要能均摊到原始 token 索引域，也要保留 packed KV 数据行的对齐。 + alignment = lcm(self.block_token_size, self.kv_dim + self.rope_dim) + return (block_bytes + alignment - 1) // alignment * alignment // self.block_token_size + + def get_attention_view(self, packed_kv: torch.Tensor) -> torch.Tensor: + """返回算子使用的 [block, entry, head, byte] 视图，保留数据/scale 分区和块 stride。""" + entries = self.block_token_size // self.compress_ratio + return packed_kv[:, : entries * self.kv_entry_bytes].view(-1, entries, 1, self.kv_entry_bytes) + + def __post_init__(self) -> None: + assert ( + self.compress_ratio > 0 and self.block_token_size > 0 and self.block_token_size % self.compress_ratio == 0 + ) + assert self.kv_dim > self.rope_dim >= 0 and (self.kv_dim - self.rope_dim) % 64 == 0 + assert self.indexer_dim >= 0 and 0 <= self.tp_rank < self.tp_world_size + assert self.logical_layer_groups and all(self.logical_layer_groups), "storage layer groups must be nonempty" + logical_layer_ids = [layer for group in self.logical_layer_groups for layer in group] + assert len(logical_layer_ids) == len( + set(logical_layer_ids) + ), "logical layers cannot repeat across storage groups" + + +@dataclass(frozen=True) +class CompressedAttentionLayerView: + # [压缩块, packed KV 字节]，保留整个块中数据与 scale 的分区格式。 + kv: torch.Tensor + # 可选的 [压缩块, key/scale 字节]；未配置 indexer 时返回 None。 + indexer_k: torch.Tensor | None + + +class CompressedAttentionState(FullAttentionState): + """保存一种序列压缩比例的历史 KV 和可选 indexer，不依赖具体模型配置。 + + 本 State 不执行压缩，不保存 SWA 或压缩器运行态；这些运行态由 bounded State 管理。 + 不同压缩方案各自持有存储，沿用原始 token 索引及共享分配器。 + 当前实现固定为单个共享 KV head，各 TP rank 完整复制；不执行 CSA/HCA 的历史选择。 + """ + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: CompressedAttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只准备压缩 KV 配置和层映射，存储由统一容量分配阶段创建。""" + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回序列压缩比例、FP8 分组及 E8M0 scale 编码对应的方案标记。""" + return f"compressed_c{self.config.compress_ratio}_fp8_group64_e8m0" + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + # 外部 allocator 的每个物理页必须包含完整压缩块，不能拆开块内的数据和 scale。 + assert ( + self.page_size > 0 and self.page_size % self.config.block_token_size == 0 + ), "allocator page_size must be a positive multiple of block_token_size" + assert ( + token_capacity % self.config.block_token_size == 0 + ), "token capacity must contain complete compressed blocks" + self.capacity = token_capacity + # token 轴仍使用原始索引域；每 block_token_size 个槽位组成一个完整 packed 压缩块。 + self.storage = torch.zeros( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + 1, + self.config.token_slot_bytes, + ), + dtype=torch.uint8, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """返回每个原始 token 索引的存储开销，包含全部物理层及压缩块 padding。""" + config = self.config + return len(config.logical_layer_groups) * config.token_slot_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> CompressedAttentionLayerView: + """从同一物理压缩块取 KV 和 indexer 视图，不重排数据或创建副本。""" + config = self.config + layer = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + blocks = layer.view(-1, config.block_token_size * config.token_slot_bytes) + entries = config.block_token_size // config.compress_ratio + kv_bytes = config.kv_block_bytes + indexer = ( + blocks[:, kv_bytes : kv_bytes + entries * (config.indexer_dim + torch.float32.itemsize)] + if config.indexer_dim + else None + ) + return CompressedAttentionLayerView(kv=blocks[:, :kv_bytes], indexer_k=indexer) + + def write_full_att_kv( + self, + logical_layer_id: int, + token_indices: torch.Tensor, + kv: torch.Tensor, + ) -> None: + """压缩器通过计算视图融合写入；普通 KV 写入路径不应调用此方法。""" + raise NotImplementedError("compressed KV is written by the fused compressor") + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整压缩块；调用方保证块内索引完整、顺序对应且源目标不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + # 数据、scale 和 indexer 在块内分区保存，不能只复制部分原始 token 槽位。 + assert ( + src_indices.numel() % self.config.block_token_size == 0 + ), "token copies require complete compressed blocks" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + copy_packed_tokens( + src_storage=self.storage, + dst_storage=self.storage, + src_token_indices=src_indices, + dst_token_indices=dst_indices, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小描述完整压缩块，包含 KV、scale、indexer 和块内 padding。""" + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + page_tokens = layout.tokens_per_page + assert ( + page_tokens is not None and page_tokens > 0 and page_tokens % self.config.block_token_size == 0 + ), "CPU page token capacity must be a positive multiple of block_token_size" + # 联合页可包含其他 State；只校验本 State 的区域完整且不越过单页末尾。 + page_bytes = prod(cpu_page_buffers.shape[1:]) * cpu_page_buffers.element_size() + required_bytes = self.get_token_bytes() * page_tokens + assert ( + 0 <= layout.byte_start < layout.byte_start + required_bytes <= layout.byte_end <= page_bytes + ), "compressed State byte range must fit within one CPU page" + assert all( + 0 <= page < cpu_page_buffers.shape[0] for page in cpu_page_indices + ), "CPU page index is out of bounds" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + page_tokens = layout.tokens_per_page + assert ( + page_tokens is not None and page_tokens > 0 and page_tokens % self.config.block_token_size == 0 + ), "CPU page token capacity must be a positive multiple of block_token_size" + # 读回也允许联合页中的非零偏移；页面必须完整容纳当前 State 的全部压缩块。 + page_bytes = prod(cpu_page_buffers.shape[1:]) * cpu_page_buffers.element_size() + required_bytes = self.get_token_bytes() * page_tokens + assert ( + 0 <= layout.byte_start < layout.byte_start + required_bytes <= layout.byte_end <= page_bytes + ), "compressed State byte range must fit within one CPU page" + assert all( + 0 <= page < cpu_page_buffers.shape[0] for page in cpu_page_indices + ), "CPU page index is out of bounds" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: CompressedAttentionState, + ) -> None: + """复制其他 DP 的完整压缩块；调用方保证块内索引完整且顺序对应。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + assert ( + src_token_indices.numel() % self.config.block_token_size == 0 + ), "DP token copies require complete compressed blocks" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + # packed scale/indexer 可位于末尾有效 token 之后，PD 必须传完整的末压缩块。 + # TODO：将末块索引扩展融合到传输 kernel，省去这份很小的索引临时区。 + alignment = self.config.block_token_size + padding = (-token_indices.numel()) % alignment + if token_indices.numel() and padding: + tail = token_indices[-1:] + torch.arange( + start=1, + end=padding + 1, + dtype=token_indices.dtype, + device=token_indices.device, + ) + token_indices = torch.cat(tensors=(token_indices, tail)) + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + # packed scale/indexer 可位于末尾有效 token 之后，PD 必须传完整的末压缩块。 + # TODO：将末块索引扩展融合到传输 kernel，省去这份很小的索引临时区。 + alignment = self.config.block_token_size + padding = (-token_indices.numel()) % alignment + if token_indices.numel() and padding: + tail = token_indices[-1:] + torch.arange( + start=1, + end=padding + 1, + dtype=token_indices.dtype, + device=token_indices.device, + ) + token_indices = torch.cat(tensors=(token_indices, tail)) + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """整个物理压缩块按原始 token 槽位传输，KV、scale、indexer 和 padding 一起保留。""" + config = self.config + assert ( + page_tokens > 0 and page_tokens % config.block_token_size == 0 + ), "CPU/PD page token capacity must be a positive multiple of block_token_size" + shape = (len(config.logical_layer_groups), page_tokens, 1, config.token_slot_bytes) + return ( + TransferRegion( + field="kv", + byte_start=0, + byte_end=prod(shape), + dtype=torch.uint8, + axes=("layer", "token", "head", "dim"), + shape=shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=config.token_slot_bytes, + ), + ), + write_enabled=config.tp_rank == 0, + ), + ) diff --git a/lightllm/common/attention_state/full_att/create_compressed_att_state.py b/lightllm/common/attention_state/full_att/create_compressed_att_state.py new file mode 100644 index 0000000000..d01c080ae6 --- /dev/null +++ b/lightllm/common/attention_state/full_att/create_compressed_att_state.py @@ -0,0 +1,14 @@ +"""序列压缩 attention 历史 KV 的独立创建入口。""" + +import torch + +from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionState + + +def create_compressed_att_state( + config: CompressedAttentionConfig, + page_size: int = 256, + device: str | torch.device = "cuda", +) -> CompressedAttentionState: + """根据压缩比例创建历史 KV 定义；不分配存储，也不创建 SWA/compressor 运行态。""" + return CompressedAttentionState(layout_config=config, page_size=page_size, device=device) diff --git a/lightllm/common/attention_state/full_att/create_mla_state.py b/lightllm/common/attention_state/full_att/create_mla_state.py index 5ebd16c0a8..05e3459f4b 100644 --- a/lightllm/common/attention_state/full_att/create_mla_state.py +++ b/lightllm/common/attention_state/full_att/create_mla_state.py @@ -7,9 +7,10 @@ import torch -from .base_state import FullAttentionConfig +from .base_state import FullAttentionConfig, FullAttentionState from .mla_state import MLAFullAttentionState from .packed_mla_state import PackedMLAAttentionConfig, PackedMLAFullAttentionState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState def create_mla_state( @@ -19,11 +20,14 @@ def create_mla_state( indexer_dim: int = 128, page_size: int = 1, device: str | torch.device = "cuda", -) -> MLAFullAttentionState | PackedMLAFullAttentionState: + fp8_format: Literal["group128_fp32", "group64_e8m0"] = "group128_fp32", +) -> FullAttentionState: """创建普通或 packed MLA State，统一处理量化编码、TP 副本和物理行宽。 config.head_dim 始终是逻辑 KV latent 维度，尚未包含 indexer、scale 和 padding。 packed=True 时把 indexer 字段存入同一行；fp8kv_dsa 仅适用于 packed MLA。 + fp8_format 由模型显式选择，仅在 packed FP8 时生效；不根据 latent 维度推断编码。 + group128_fp32 与 group64_e8m0 使用独立 State。 """ if not packed: assert kv_type == "None", f"unsupported MLA encoding: {kv_type}" @@ -34,9 +38,22 @@ def create_mla_state( ) assert kv_type in ("None", "fp8kv_dsa"), f"unsupported packed MLA encoding: {kv_type}" fp8 = kv_type == "fp8kv_dsa" + assert not fp8 or fp8_format in ("group128_fp32", "group64_e8m0"), f"unknown packed MLA FP8 format: {fp8_format}" + if fp8 and fp8_format == "group64_e8m0": + assert config.head_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" + assert config.global_head_num == 1, "FP8/E8M0 packed MLA requires one replicated latent" + return FP8E8M0PackedMLAState( + layout_config=FP8E8M0PackedMLAConfig( + logical_layer_groups=config.logical_layer_groups, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + indexer_dim=indexer_dim, + ), + page_size=page_size, + device=device, + ) storage_dtype = torch.uint8 if fp8 else config.dtype - # FP8 latent 包含 512 字节编码、16 字节 scale 和 128 字节 BF16 RoPE。 - # indexer key 使用 INT8 并附带一个 FP32 scale；整行按 16 字节对齐。 + # 原有 DSA 格式：512 FP8 + 16 字节 FP32 scales + 64 BF16，另附 indexer key/scale。 kv_bytes = 656 if fp8 else config.head_dim * config.dtype.itemsize row_bytes = (kv_bytes + indexer_dim + 4 + 15) // 16 * 16 packed_config = PackedMLAAttentionConfig( diff --git a/lightllm/common/attention_state/full_att/dp_transfer.py b/lightllm/common/attention_state/full_att/dp_transfer.py index 8ded336af5..edcb8c2b5d 100644 --- a/lightllm/common/attention_state/full_att/dp_transfer.py +++ b/lightllm/common/attention_state/full_att/dp_transfer.py @@ -45,13 +45,16 @@ def copy_group_quantized_tokens( ) -def copy_mla_tokens( +def copy_packed_tokens( src_storage: torch.Tensor, dst_storage: torch.Tensor, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, ) -> None: - """复制 MLA 完整行；packed 方案同时保留 RoPE、scale、indexer 和行 padding。""" + """复制单字段存储的完整 token 槽位，包含编码数据、scale、indexer 和 padding。 + + 压缩历史等方案的完整块约束由具体 State 在调用前检查，本函数不识别压缩比例。 + """ _copy_token_storages( src_storages=(src_storage,), dst_storages=(dst_storage,), diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py new file mode 100644 index 0000000000..69e6c4c289 --- /dev/null +++ b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py @@ -0,0 +1,299 @@ +"""FP8/E8M0 packed MLA：448 维 FP8、64 维 BF16，并在同一行保存 indexer。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ..utils import deepcopy_result +from ..types import AxisMapping, StateTransferLayout, TransferRegion +from .base_state import FullAttentionState +from .dp_transfer import copy_packed_tokens +from .packed_transfer import transfer_packed_pages +from .packed_mla_state import PackedMLALayerView + + +@dataclass(frozen=True) +class FP8E8M0PackedMLAConfig: + """固定 group64/E8M0 编码，只记录调用方能够选择的原子布局信息。""" + + # 外层为物理存储层，内层为共享该层全部 KV/indexer 的逻辑层号。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 每个 token 的 indexer key 维度；其后另保存一个 FP32 scale，不包含在该维度中。 + indexer_dim: int = 128 + # 完整 latent 在各 TP rank 复制；rank 0 写出共享外部页，各 rank 均可读回。 + tp_rank: int = 0 + # 所属 TP 组大小，不增加物理 KV head，也不改变单 token 的存储字节数。 + tp_world_size: int = 1 + + # 算子输入的完整 latent 维度，不是量化后的字节数，也不包含 indexer。 + LATENT_DIM: ClassVar[int] = 512 + # 固定编码：448 个 FP8 元素 + 64 个 BF16 元素 + 7 个 E8M0 scale + 1 字节对齐。 + # 末尾 64 维保留 BF16 不代表启用了 RoPE，GLM 使用 NoPE attention。 + KV_BYTES: ClassVar[int] = 584 + + @property + def payload_bytes(self) -> int: + """一行有效数据的字节数：KV 编码、indexer key 及其 FP32 scale，不包含行尾 padding。""" + return self.KV_BYTES + self.indexer_dim + torch.float32.itemsize + + @property + def token_bytes(self) -> int: + """每个 token 的物理行字节数，包含末尾 16 字节对齐 padding。""" + return (self.payload_bytes + 15) // 16 * 16 + + def __post_init__(self) -> None: + assert self.logical_layer_groups and all(self.logical_layer_groups) + layers = [layer for group in self.logical_layer_groups for layer in group] + assert len(layers) == len(set(layers)), "logical layers cannot repeat across storage groups" + assert self.indexer_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + + +class FP8E8M0PackedMLAState(FullAttentionState): + """独立的 group64/E8M0 MLA 实现，不继承其他具体 State 或动态选择编码。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: FP8E8M0PackedMLAConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只准备布局和逻辑层映射，实际存储由管理器确定 token 容量后统一分配。""" + self.config = layout_config + self.state_name: str | None = None + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in group + } + + def get_state_type(self) -> str: + """返回固定算法、量化分组与 scale 编码标记，用于 main/draft 的规范成员名。""" + return "mla_packed_fp8_group64_e8m0" + + def allocate_storage(self, token_capacity: int) -> None: + """分配普通 token 及 HOLD 页；完整行包含 KV、indexer 和 16 字节对齐 padding。""" + self.capacity = token_capacity + self.storage = torch.empty( + size=(len(self.config.logical_layer_groups), token_capacity + self.page_size, 1, self.config.token_bytes), + dtype=torch.uint8, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """返回当前 TP rank 的单 token 开销，包含全部物理层和行尾 padding。""" + return len(self.config.logical_layer_groups) * self.config.token_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断该逻辑层是否使用当前 State，共享逻辑层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """将逻辑层映射为物理存储行；同组逻辑层返回相同下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: + """返回固定 584 字节 KV 和 indexer key/scale 视图，保留物理行 stride。""" + row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + return PackedMLALayerView( + kv=row[..., : self.config.KV_BYTES], + indexer_k=row[..., self.config.KV_BYTES : self.config.payload_bytes], + ) + + def write_full_att_kv( + self, + logical_layer_id: int, + token_indices: torch.Tensor, + kv: torch.Tensor, + ) -> None: + """使用固定 group64/E8M0 算子量化并写入 KV，保留 indexer 和行尾 padding。""" + from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import ( + write_fp8_e8m0_packed_kv, + ) + + assert token_indices.is_cuda, "token indices must reside on CUDA" + if token_indices.numel() == 0: + return + write_fp8_e8m0_packed_kv( + kv=kv, + token_indices=token_indices, + packed_kv=self.get_layer_view(logical_layer_id=logical_layer_id).kv, + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """按 CUDA 索引复制完整物理行，包含 indexer、scale 和 padding；两组索引不得重叠。""" + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + if src_indices.numel() == 0: + return + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache：按 token 搬运整份有效编码 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 CPU 页中的完整有效 token 编码，包含 indexer，不包含 GPU 行尾 padding。""" + return self._get_transfer_regions(page_tokens=get_env_start_args().cpu_cache_token_page_size) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """直接写出 KV、E8M0 scale 和 indexer；已就绪的 CPU 页面保持不变。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """从 CPU 页恢复完整有效 token 编码，支持末页只恢复部分 token。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移：复制相同编码的完整物理行 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: FP8E8M0PackedMLAState, + ) -> None: + """从其他 DP 的同一 TP 分片复制完整 token 数据，由调用方提供可访问的 CUDA 索引。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离：传输固定的完整 token 编码 + # ======================================================================== + + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回固定 FP8/E8M0 格式的 PD 页描述，由管理器安排联合页偏移。""" + return self._get_transfer_regions(page_tokens=get_env_start_args().pd_kv_page_size) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将完整有效 token 编码写入指定 PD 页，网络发送与同步由调用方负责。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """从指定 PD 页恢复完整有效 token 编码，保留其他 token 与行尾 padding。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助：固定 packed 编码作为一个区域传输 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + # 有效字节在每行内连续，KV、scale 和 indexer 无需分别发射 kernel。 + # GPU 行有 padding，通过实际 stride 跳过；外部页只保存 payload_bytes。 + return ( + TransferRegion( + field="kv", + byte_start=0, + byte_end=len(self.config.logical_layer_groups) * page_tokens * self.config.payload_bytes, + dtype=torch.uint8, + axes=("layer", "token", "head", "dim"), + shape=(len(self.config.logical_layer_groups), page_tokens, 1, self.config.payload_bytes), + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=self.config.payload_bytes, + ), + ), + write_enabled=self.config.tp_rank == 0, + ), + ) diff --git a/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_write.py similarity index 64% rename from lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py rename to lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_write.py index 7a28bc6614..495bbbf7cb 100644 --- a/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py +++ b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_write.py @@ -1,22 +1,45 @@ +"""group64/E8M0 packed MLA 的定制 KV 量化写入，不依赖具体模型。""" + import torch import triton import triton.language as tl +@torch.no_grad() +def write_fp8_e8m0_packed_kv( + kv: torch.Tensor, + token_indices: torch.Tensor, + packed_kv: torch.Tensor, +) -> None: + assert kv.shape[1:] == (1, 512) + assert packed_kv.dtype == torch.uint8 and packed_kv.shape[1] == 1 + assert packed_kv.shape[2] >= 584 and packed_kv.stride(2) == 1 + _write_fp8_e8m0_packed_kv[(token_indices.numel(),)]( + KV=kv, + TOKEN_INDICES=token_indices.contiguous(), + PACKED_KV=packed_kv, + kv_stride_s=kv.stride(0), + kv_stride_d=kv.stride(2), + packed_stride_s=packed_kv.stride(0), + num_warps=4, + num_stages=1, + ) + + @triton.jit -def _fwd_kernel_destindex_copy_kv_flashmla_fp8( +def _write_fp8_e8m0_packed_kv( KV, - MemIndex, - PackedKV, + TOKEN_INDICES, + PACKED_KV, kv_stride_s, kv_stride_d, packed_stride_s, ): token = tl.program_id(0) - mem_index = tl.load(MemIndex + token).to(tl.int64) - packed_row = PackedKV + mem_index * packed_stride_s + mem_index = tl.load(TOKEN_INDICES + token).to(tl.int64) + packed_row = PACKED_KV + mem_index * packed_stride_s - # Eight rows share one CTA; seven quantized groups are processed in parallel. + # 每个 token 按 group64 并行量化七组特征；末尾 64 维单独保留 BF16。 groups = tl.arange(0, 8) offsets = groups[:, None] * 64 + tl.arange(0, 64)[None, :] group_mask = groups < 7 @@ -33,20 +56,3 @@ def _fwd_kernel_destindex_copy_kv_flashmla_fp8( tail_values = tl.load(KV + token * kv_stride_s + (448 + tail_offsets) * kv_stride_d) tl.store((packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + tail_offsets, tail_values) tl.store(packed_row + 583, 0) - - -@torch.no_grad() -def destindex_copy_kv_flashmla_fp8(kv: torch.Tensor, mem_index: torch.Tensor, packed_kv: torch.Tensor): - assert kv.shape[1:] == (1, 512) - assert packed_kv.dtype == torch.uint8 and packed_kv.shape[1] == 1 - assert packed_kv.shape[2] >= 584 and packed_kv.stride(2) == 1 - _fwd_kernel_destindex_copy_kv_flashmla_fp8[(mem_index.numel(),)]( - kv, - mem_index.contiguous(), - packed_kv, - kv.stride(0), - kv.stride(2), - packed_kv.stride(0), - num_warps=4, - num_stages=1, - ) diff --git a/lightllm/common/attention_state/full_att/gqa_state.py b/lightllm/common/attention_state/full_att/gqa_state.py index 35583f276f..0086963514 100644 --- a/lightllm/common/attention_state/full_att/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa_state.py @@ -1,5 +1,6 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod @@ -8,7 +9,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .gqa_transfer import transfer_gqa_pages @@ -144,6 +145,8 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 diff --git a/lightllm/common/attention_state/full_att/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa_transfer.py index 4c6d1bb8a8..ccd7edc990 100644 --- a/lightllm/common/attention_state/full_att/gqa_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa_transfer.py @@ -28,7 +28,7 @@ def transfer_gqa_pages( ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 - 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 @@ -39,8 +39,6 @@ def transfer_gqa_pages( assert isinstance(page_indices, list), "page indices must be a list" if page_buffers.device.type == "cpu": assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" - else: - assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" page_tokens = layout.tokens_per_page assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" assert token_indices.is_cuda, "token indices must reside on CUDA" @@ -59,7 +57,7 @@ def transfer_gqa_pages( # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -68,7 +66,7 @@ def transfer_gqa_pages( page_readies_tensor = None if write and page_readies is not None: page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) diff --git a/lightllm/common/attention_state/full_att/group_quantized_state.py b/lightllm/common/attention_state/full_att/group_quantized_state.py index f72db0940a..172af40d61 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/group_quantized_state.py @@ -1,5 +1,6 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from dataclasses import dataclass from math import prod @@ -9,7 +10,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .group_quantized_transfer import transfer_group_quantized_pages @@ -193,6 +194,8 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 diff --git a/lightllm/common/attention_state/full_att/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/group_quantized_transfer.py index 5fec71c382..a92176982d 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_transfer.py +++ b/lightllm/common/attention_state/full_att/group_quantized_transfer.py @@ -30,7 +30,7 @@ def transfer_group_quantized_pages( ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 - 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 @@ -41,8 +41,6 @@ def transfer_group_quantized_pages( assert isinstance(page_indices, list), "page indices must be a list" if page_buffers.device.type == "cpu": assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" - else: - assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" page_tokens = layout.tokens_per_page assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" assert token_indices.is_cuda, "token indices must reside on CUDA" @@ -61,7 +59,7 @@ def transfer_group_quantized_pages( # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -70,7 +68,7 @@ def transfer_group_quantized_pages( page_readies_tensor = None if write and page_readies is not None: page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) diff --git a/lightllm/common/attention_state/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla_state.py index aabd08eadc..fed5d918b8 100644 --- a/lightllm/common/attention_state/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla_state.py @@ -1,30 +1,26 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod import torch from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState -from .mla_transfer import transfer_mla_pages -from .dp_transfer import copy_mla_tokens +from .packed_transfer import transfer_packed_pages +from .dp_transfer import copy_packed_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout -logger = init_logger(name=__name__) - - class MLAFullAttentionState(FullAttentionState): """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" - # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -37,7 +33,6 @@ def __init__( ) -> None: """按 MLA 布局配置准备压缩 KV 存储。""" self.config = layout_config - self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { @@ -57,15 +52,23 @@ def get_state_type(self) -> str: def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity - shape = self._storage_shape() - self.storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.storage = torch.empty( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ), + dtype=self.config.dtype, + device=self.device, + ) def get_token_bytes(self) -> int: - """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" + """由配置计算当前 TP rank 的单 token latent 字节数，无需分配存储。""" config = self.config return ( len(config.logical_layer_groups) * config.head_num * config.head_dim - * self.storage_dtype.itemsize + * config.dtype.itemsize ) # ======================================================================== @@ -122,6 +125,8 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 @@ -140,7 +145,7 @@ def write_to_cpu( ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - transfer_mla_pages( + transfer_packed_pages( kv_storage=self.storage, token_indices=token_indices, page_buffers=cpu_page_buffers, @@ -159,7 +164,7 @@ def read_from_cpu( ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - transfer_mla_pages( + transfer_packed_pages( kv_storage=self.storage, token_indices=token_indices, page_buffers=cpu_page_buffers, @@ -184,7 +189,7 @@ def copy_tokens_from_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - copy_mla_tokens( + copy_packed_tokens( src_storage=remote_dp_full_att_state.storage, dst_storage=self.storage, src_token_indices=src_token_indices, @@ -209,7 +214,7 @@ def write_to_pd_page( ) -> None: """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" - transfer_mla_pages( + transfer_packed_pages( kv_storage=self.storage, token_indices=token_indices, page_buffers=pd_page_buffers, @@ -227,7 +232,7 @@ def read_from_pd_page( ) -> None: """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" - transfer_mla_pages( + transfer_packed_pages( kv_storage=self.storage, token_indices=token_indices, page_buffers=pd_page_buffers, @@ -241,88 +246,25 @@ def read_from_pd_page( # ======================================================================== def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: - """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + """描述唯一的 MLA latent 字段，外部页保存全局 heads，副本只由首 rank 写出。""" config = self.config - byte_offset = 0 - regions = [] - logger.debug( - "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", - self.state_name, - config.tp_rank, - config.tp_world_size, - page_tokens, - ) - # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 - field_specs = (("kv", self.storage_dtype, config.head_dim, "dim"),) - for field_name, field_dtype, field_width, last_axis in field_specs: - axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) - global_shape = list(local_shape) - global_shape[axes.index("token")] = page_tokens - head_axis = axes.index("head") - # field 已按 K/V 拆开：这里读取的是单个字段的本地 head 数，不是合并 buffer 的 2H。 - local_head_count = global_shape[head_axis] - # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 - # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 - # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 - global_shape[head_axis] = config.global_head_num - - # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 - # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 - replica_count = config.tp_world_size * local_head_count // config.global_head_num - # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 - head_start = config.tp_rank // replica_count * local_head_count - # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 - write_enabled = config.tp_rank % replica_count == 0 - axis_mappings = tuple( - AxisMapping( - axis=axis, - local_start=0, - external_start=head_start if axis == "head" else 0, - length=dimension, - ) - for axis, dimension in zip(axes, local_shape) - if axis not in ("layer", "token") - ) - - # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 - element_bytes = field_dtype.itemsize - byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes - field_region = TransferRegion( - field=field_name, - byte_start=byte_offset, - byte_end=byte_offset + prod(global_shape) * element_bytes, - dtype=field_dtype, - axes=axes, - shape=tuple(global_shape), - axis_mappings=axis_mappings, - write_enabled=write_enabled, - ) - regions.append(field_region) - byte_offset = field_region.byte_end - logger.debug( - "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.state_name, - field_name, - head_start, - local_head_count, - replica_count, - write_enabled, - field_region.byte_start, - field_region.byte_end, - ) - logger.debug( - "%s transfer regions built: fields=%d, total_bytes=%d", - self.state_name, - len(regions), - byte_offset, - ) - return tuple(regions) - - def _storage_shape(self) -> tuple[int, ...]: + shape = (len(config.logical_layer_groups), page_tokens, config.global_head_num, config.head_dim) + # 本地分片可在多个 rank 复制；先去掉副本编号，再定位外部页的 head 起点。 + local_head_count = config.head_num + replica_count = config.tp_world_size * local_head_count // config.global_head_num + head_start = config.tp_rank // replica_count * local_head_count return ( - len(self.config.logical_layer_groups), - self.capacity + self.page_size, - self.config.head_num, - self.config.head_dim, + TransferRegion( + field="kv", + byte_start=0, + byte_end=prod(shape) * config.dtype.itemsize, + dtype=config.dtype, + axes=("layer", "token", "head", "dim"), + shape=shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=local_head_count), + AxisMapping(axis="dim", local_start=0, external_start=0, length=config.head_dim), + ), + write_enabled=config.tp_rank % replica_count == 0, + ), ) diff --git a/lightllm/common/attention_state/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla_state.py index 95d9c7344f..6e1e21fe8f 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla_state.py @@ -1,5 +1,6 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, List from dataclasses import dataclass from math import prod @@ -9,11 +10,11 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import dtype_name +from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .packed_mla_transfer import transfer_packed_mla_pages -from .dp_transfer import copy_mla_tokens +from .dp_transfer import copy_packed_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -195,6 +196,8 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 @@ -272,7 +275,7 @@ def copy_tokens_from_dp( ) assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - copy_mla_tokens( + copy_packed_tokens( src_storage=remote_dp_full_att_state.storage, dst_storage=self.storage, src_token_indices=src_token_indices, diff --git a/lightllm/common/attention_state/full_att/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla_transfer.py index 22e39de04e..5d437d78cd 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla_transfer.py @@ -30,7 +30,7 @@ def transfer_packed_mla_pages( ) -> None: """直接搬运 latent 及配套字段；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 - 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 @@ -41,8 +41,6 @@ def transfer_packed_mla_pages( assert isinstance(page_indices, list), "page indices must be a list" if page_buffers.device.type == "cpu": assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" - else: - assert page_buffers.device == kv_storage.device, "CUDA pages must reside on the State device" page_tokens = layout.tokens_per_page assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" assert token_indices.is_cuda, "token indices must reside on CUDA" @@ -61,7 +59,7 @@ def transfer_packed_mla_pages( # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( - device=kv_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -70,7 +68,7 @@ def transfer_packed_mla_pages( page_readies_tensor = None if write and page_readies is not None: page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( - device=kv_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) diff --git a/lightllm/common/attention_state/full_att/mla_transfer.py b/lightllm/common/attention_state/full_att/packed_transfer.py similarity index 95% rename from lightllm/common/attention_state/full_att/mla_transfer.py rename to lightllm/common/attention_state/full_att/packed_transfer.py index 2172f65c69..e7bdb50f92 100644 --- a/lightllm/common/attention_state/full_att/mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_transfer.py @@ -1,4 +1,4 @@ -"""MLA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 latent 与外部页。""" +"""单字段 packed 存储的 CPU/PD 页传输，供 MLA 和压缩历史等方案复用。""" import torch import triton @@ -13,7 +13,7 @@ # ======================================================================== -def transfer_mla_pages( +def transfer_packed_pages( kv_storage: torch.Tensor, token_indices: torch.Tensor, page_buffers: torch.Tensor, @@ -25,9 +25,9 @@ def transfer_mla_pages( num_warps: int = 8, num_stages: int = 1, ) -> None: - """直接搬运 latent；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + """直接搬运 packed 存储；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 - 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 @@ -38,8 +38,6 @@ def transfer_mla_pages( assert isinstance(page_indices, list), "page indices must be a list" if page_buffers.device.type == "cpu": assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" - else: - assert page_buffers.device == kv_storage.device, "CUDA pages must reside on the State device" page_tokens = layout.tokens_per_page assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" assert token_indices.is_cuda, "token indices must reside on CUDA" @@ -52,13 +50,13 @@ def transfer_mla_pages( if token_count == 0: return - # 阶段一：准备索引，不预选或复制 latent。 + # 阶段一：准备索引，不预选或复制 payload。 # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( - device=kv_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -67,7 +65,7 @@ def transfer_mla_pages( page_readies_tensor = None if write and page_readies is not None: page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( - device=kv_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -120,7 +118,7 @@ def transfer_mla_pages( )[0] local, external = view_max_copy_dtype(copies=[(local, external)])[0] elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] - _copy_mla_pages[(grid_num,)]( + _copy_packed_pages[(grid_num,)]( LOCAL=local, EXTERNAL=external, TOKEN_IDS=token_indices, @@ -149,7 +147,7 @@ def transfer_mla_pages( @triton.jit -def _copy_mla_pages( +def _copy_packed_pages( # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 diff --git a/lightllm/common/attention_state/full_att/static_fp8_state.py b/lightllm/common/attention_state/full_att/static_fp8_state.py index 14350cde90..458ebae0c0 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/static_fp8_state.py @@ -1,5 +1,7 @@ from __future__ import annotations +from functools import lru_cache + from typing import TYPE_CHECKING, List from collections.abc import Mapping from dataclasses import dataclass @@ -10,6 +12,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger +from ..utils import deepcopy_result from ..types import AxisMapping, TransferRegion from .base_state import FullAttentionConfig, FullAttentionState from .static_fp8_transfer import transfer_static_fp8_pages @@ -231,6 +234,8 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 三、CPU Cache 数据交换 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 diff --git a/lightllm/common/attention_state/full_att/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/static_fp8_transfer.py index eefbb15e41..e8bd29d220 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_transfer.py +++ b/lightllm/common/attention_state/full_att/static_fp8_transfer.py @@ -28,7 +28,7 @@ def transfer_static_fp8_pages( ) -> None: """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 - 外部页为 pinned/注册 CPU 内存或当前设备的 CUDA PD buffer。 + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 @@ -39,8 +39,6 @@ def transfer_static_fp8_pages( assert isinstance(page_indices, list), "page indices must be a list" if page_buffers.device.type == "cpu": assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" - else: - assert page_buffers.device == k_storage.device, "CUDA pages must reside on the State device" page_tokens = layout.tokens_per_page assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" assert token_indices.is_cuda, "token indices must reside on CUDA" @@ -59,7 +57,7 @@ def transfer_static_fp8_pages( # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) @@ -68,7 +66,7 @@ def transfer_static_fp8_pages( page_readies_tensor = None if write and page_readies is not None: page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( - device=k_storage.device, + device=torch.cuda.current_device(), non_blocking=True, ) diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index b9af2e2c73..b895383da9 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -265,6 +265,8 @@ def allocate_storage( self.cpu_tail_restore = bool(self.get_attention_tail_states()) # 存储分配可能改变 State 的物理页规格，清除准备阶段的布局及其派生 CPU/PD 页大小缓存。 + for state in (*full_states, *bounded_states): + state.get_cpu_transfer_regions.cache_clear() self.get_cpu_page_layouts.cache_clear() self.get_cpu_page_bytes.cache_clear() self.get_pd_token_page_layouts.cache_clear() @@ -292,10 +294,9 @@ def create_dp_ipc_copy(self) -> AttentionStateManager: for state in self.get_bounded_attention_states(): for pool in (state.big_pages_pool, state.small_pages_pool): if pool is not None: - assert ( - pool.conv_state.device.type == "cpu" and pool.conv_state.is_shared() - and pool.ssm_state.device.type == "cpu" and pool.ssm_state.is_shared() - ), "DP IPC requires shared CPU page pools" + for tensor in vars(pool).values(): + if isinstance(tensor, torch.Tensor): + assert tensor.device.type == "cpu" and tensor.is_shared(), "DP IPC requires shared CPU page pools" ipc_copy = copy(self) ipc_copy.full_att_allocator = None ipc_copy.bounded_att_big_page_allocator = None @@ -484,6 +485,9 @@ def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) for region in state.get_cpu_transfer_regions() ) + # 无跨步重叠的压缩器在组边界不需要快照，不占用 CPU 页。 + if not regions: + continue # 以实际字段终点确定 State 占用，保留内部 padding，并对齐下一个 State 的起点。 state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 layouts[state.state_name] = StateTransferLayout( @@ -542,7 +546,12 @@ def write_req_to_cpu( # 阶段二：分别从大小页池取快照，写入对应 CPU 页面，不修改请求运行态。 for state in bounded_states: - layout = layouts[state.state_name] + layout = layouts.get(state.state_name) + if layout is None: + # 部分 bounded 状态在缓存边界没有需保留的历史（如 W128/S128 压缩器），不写入 CPU cache。 + # 仅允许未声明 CPU 传输区域的 State 跳过，避免必要快照因布局遗漏而丢失。 + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue if big_page_slot_ids: state.write_state_pages_to_cpu( page_kind="big", @@ -585,8 +594,9 @@ def read_req_from_cpu( assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" if bounded_states: assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" + # CPU 页的 token 数与 bounded 大页边界在启动时统一，直接读取全局配置。 + page_tokens = get_env_start_args().cpu_cache_token_page_size # token 索引只用于判断本批末页是否完整，不读取 GPU 上的索引值。 - page_tokens = layouts[bounded_states[0].state_name].tokens_per_page has_partial_page = token_indices.numel() % page_tokens != 0 assert isinstance(big_page_slot_ids, list), "slot indices must be a list" assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), "slots must cover complete CPU pages only" @@ -606,6 +616,11 @@ def read_req_from_cpu( # 阶段二：完整页和末尾碎页统一读入大页池，碎页使用追加的 LOAD 槽位。 for state in bounded_states: + if state.state_name not in layouts: + # 无跨步历史的 bounded 状态不占 CPU 页，因此无需读回快照；运行态恢复仍在下一阶段处理。 + # 有 CPU 传输区域的 State 必须具备布局，不能静默跳过。 + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue state.read_state_pages_from_cpu( cpu_page_buffers=cpu_page_buffers, cpu_page_indices=cpu_page_indices, diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py index a1181d1d24..c2f331f5bf 100644 --- a/lightllm/common/attention_state/utils.py +++ b/lightllm/common/attention_state/utils.py @@ -2,20 +2,41 @@ from __future__ import annotations +from collections.abc import Callable +from copy import deepcopy +from functools import WRAPPER_ASSIGNMENTS, wraps +from typing import ParamSpec, TypeVar + import torch +_P = ParamSpec("_P") +_R = TypeVar("_R") + MAX_DRAFT_MODELS = 36 MODEL_ROLES = ("main", *(f"draft{i}" for i in range(MAX_DRAFT_MODELS))) # ======================================================================== -# 一、State 命名与成员管理 +# 一、State 命名与返回值副本 # ======================================================================== def dtype_name(dtype: torch.dtype): return {torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32"}[dtype] +def deepcopy_result(func: Callable[_P, _R]) -> Callable[_P, _R]: + """每次返回结果的独立深拷贝，放在 lru_cache 外层以保护缓存原对象。 + + 保留被装饰函数的缓存管理接口，调用方仍可使用 cache_clear/cache_info/cache_parameters。 + """ + + @wraps(func, assigned=WRAPPER_ASSIGNMENTS + ("cache_clear", "cache_info", "cache_parameters")) + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: + return deepcopy(func(*args, **kwargs)) + + return wrapped + + # ======================================================================== # 二、传输视图：合并两端共同支持的尾部维度 # ======================================================================== diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py index 536db3d4fe..448efcf979 100644 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py @@ -18,42 +18,69 @@ class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - self.cache_config = model.req_manager.cache_config + 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(model.config["compress_ratios"])) + 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 model.args.run_mode == "decode": return None, 0 - config = model.req_manager.cache_config - page_token_num = config.swa_page_token_num + 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.sliding_window + 2 * page_token_num), page_token_num + max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num ) * page_token_num ) 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 = cls.get_gpu_workspace_buffer( - key_name="deepseek_v4_swa_prefill", - workspace_size=config.layer_num * token_num * config.swa_slot_bytes, + key_name=f"deepseek_v4_swa_prefill_{model.att_state_role}", + workspace_size=layer_count * token_num * slot_bytes, dtype=torch.uint8, ) - return buffer.view(config.layer_num, token_num, config.swa_slot_bytes), capacity - - def get_att_input_params(self, infer_state, layer_index): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import view_swa_kv - - swa, compressed_kv = infer_state.req_manager.get_att_input_params(layer_index) + return buffer.view(layer_count, token_num, slot_bytes), capacity + + def get_att_input_params(self, infer_state, logical_layer_id: int): + state = infer_state.swa_state + swa_buffer = state.get_layer_view(logical_layer_id=logical_layer_id) + config = state.config + swa = swa_buffer.view(-1, config.get_packed_kv_bytes(config.page_token_size)) + ratio = self.layer_compress_ratios[logical_layer_id] + compressed_kv = infer_state.get_compressed_kv_view(logical_layer_id=logical_layer_id).kv if ratio else None if infer_state.is_prefill: - swa = view_swa_kv(infer_state.prefill_att_state.swa_buffer[layer_index], self.cache_config) + swa = infer_state.prefill_att_state.swa_buffer[config.logical_layer_ids.index(logical_layer_id)].view( + -1, config.get_packed_kv_bytes(config.page_token_size), + ) return swa, compressed_kv def _flashmla_att( @@ -66,20 +93,12 @@ def _flashmla_att( ) -> torch.Tensor: """Attend to SWA plus optional compressed KV and return all padded Q heads.""" swa_packed, compressed_packed = packed_kv - config = self.cache_config - bytes_per_token = config.kv_bytes - # Expose FlashMLA's [page, token, KV head, byte] axes while retaining padded page strides. + # 原子配置提供计算视图；backend 不再计算编码字节数或解释压缩块布局。 compressed_cache = None if compressed_packed is not None: - page_size = config.compressed_kv_token_alignment // nsa_dict["compress_ratio"] - bytes_per_page = page_size * bytes_per_token - compressed_cache = compressed_packed[:, :bytes_per_page].view( - compressed_packed.shape[0], page_size, 1, bytes_per_token - ) - swa_bytes_per_page = config.swa_page_token_num * bytes_per_token - swa_cache = swa_packed[:, :swa_bytes_per_page].view( - swa_packed.shape[0], config.swa_page_token_num, 1, bytes_per_token - ) + config = self.compressed_configs[nsa_dict["compress_ratio"]] + compressed_cache = config.get_attention_view(packed_kv=compressed_packed) + swa_cache = self.swa_config.get_attention_view(packed_kv=swa_packed) padded_out, _ = flashmla.flash_mla_with_kvcache( q=q.unsqueeze(1), @@ -116,22 +135,22 @@ class _PrefillAttState(BasePrefillAttState): def prepare_swa_workspace(self, infer_state): from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace - requests = infer_state.req_manager + swa_state = infer_state.swa_state self.swa_buffer, capacity = self.backend.get_swa_prefill_workspace(self.backend.model) - self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - self.backend.cache_config.swa_page_token_num + self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - swa_state.config.page_token_size self.prefill_token_ranges = torch.zeros( - (requests.max_request_num + 1, 2), dtype=torch.int32, device=self.swa_buffer.device + (swa_state.request_capacity, 2), dtype=torch.int32, device=self.swa_buffer.device ) self.swa_prefill_workspace_offsets = torch.full( - (requests.max_request_num + 1,), -1, dtype=torch.int32, device=self.swa_buffer.device + (swa_state.request_capacity,), -1, dtype=torch.int32, device=self.swa_buffer.device ) prepare_swa_prefill_workspace( - requests, - self, - infer_state.b_req_idx, - infer_state.b_ready_cache_len, - infer_state.b_seq_len, - infer_state.microbatch_index * capacity, + swa_state=swa_state, + prefill_state=self, + b_req_idx=infer_state.b_req_idx, + b_ready_cache_len=infer_state.b_ready_cache_len, + b_seq_len=infer_state.b_seq_len, + workspace_base=infer_state.microbatch_index * capacity, ) def init_state(self): diff --git a/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py deleted file mode 100644 index 6794446508..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py +++ /dev/null @@ -1,175 +0,0 @@ -import math - -import torch -import triton - -from .mem_manager import MemoryManager -from .operator import DeepseekV4MemOperator -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig, DeepseekV4StateCacheManager -from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_compressed_kv_pages -from lightllm.utils.envs_utils import get_env_start_args - - -class DeepseekV4MemoryManager(MemoryManager): - operator_class = DeepseekV4MemOperator - - def __init__( - self, - size, - dtype, - head_num, - head_dim, - layer_num, - cache_config: DeepseekV4CacheConfig, - always_copy=False, - mem_fraction=0.9, - ): - self.cache_config = cache_config - super().__init__( - size, - torch.uint8, - head_num, - max(cache_config.compressed_kv_slot_bytes), - len(cache_config.compressed_kv_slot_bytes), - always_copy, - mem_fraction, - ) - - def get_cell_size(self): - return sum(self.cache_config.compressed_kv_slot_bytes) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - self.kv_buffer = [ - torch.zeros((size + self.page_size, head_num, slot_bytes), dtype=dtype, device="cuda") - for slot_bytes in self.cache_config.compressed_kv_slot_bytes - ] - args = get_env_start_args() - big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - # Reserve the same CPU-cache load/offload staging slots as the other hybrid models. - self.big_page_buffers = DeepseekV4StateCacheManager( - triton.cdiv(size, big_page_tokens) + 2, self.cache_config, keep_num=2 - ) - self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 - self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 - - def _free_buffers(self): - self.kv_buffer = self.big_page_buffers = None - - def free_all(self): - super().free_all() - self.big_page_buffers.clear_to_init_state() - - def get_att_input_params(self, layer_index): - index = self.cache_config.layer_to_compressed_kv_idx.get(layer_index) - return self.kv_buffer[index] if index is not None else None - - def get_index_kv_buffer(self, index): - return {f"layer_{layer}": buffer[index].unsqueeze(0) for layer, buffer in enumerate(self.kv_buffer)} - - def load_index_kv_buffer(self, index, load_tensor_dict): - for layer, buffer in enumerate(self.kv_buffer): - buffer[index] = load_tensor_dict[f"layer_{layer}"][0].to(buffer.device) - - def __getstate__(self): - state = self.__dict__.copy() - state["big_page_buffers"] = None # Keep process-local pinned checkpoints out of CUDA IPC. - return state - - def write_to_shm(self, req_manager): - # Transport receives GPU views, just as Qwen's mem manager exports linear state. - for name in ( - "req_to_swa", - "req_to_c4_state", - "req_to_c4_indexer_state", - "req_to_c128_state", - ): - setattr(self, name, getattr(req_manager, name)) - return super().write_to_shm(req_manager) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - compressed_kv_bytes = page_size * self.get_cell_size() - swa_shape, *compressor_shapes = self.cache_config.get_runtime_state_shapes(get_env_start_args().mtp_step) - state_bytes = ( - math.prod(swa_shape) + sum(math.prod(shape) for shape in compressor_shapes) * torch.float32.itemsize - ) - return (page_num, 1, 1, 1, max(compressed_kv_bytes, state_bytes)) - - def alloc_paged_kv_move_buffer(self, page_num, page_size): - self.pd_kv_page_size = page_size - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - self.att_state_page_helper = DeepseekV4AttStatePageHelper(self) - return kv_move_buffer - - def write_mem_to_page_kv_move_buffer( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] - return self.att_state_page_helper.write_req_to_page(page_index, req_idx, dp_mems) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - page = self.kv_move_buffer[page_index] - copy_compressed_kv_pages( - mem_manager=mem_managers[dp_index * dp_world_size], - mem_indexes=pin_indexes.cuda(non_blocking=True), - page_indexes=torch.zeros(1, dtype=torch.int32, device=page.device), - kv_move_buffer=page.view(1, -1), - to_pages=True, - page_size=self.pd_kv_page_size, - round_tail=True, - ) - - def read_page_kv_move_buffer_to_mem( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] - return self.att_state_page_helper.read_page_to_req(page_index, req_idx, dp_mems) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - indexes = pin_indexes.cuda(non_blocking=True) - page = self.kv_move_buffer[page_index] - page_ids = torch.zeros(1, dtype=torch.int32, device=page.device) - for mem in mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size]: - copy_compressed_kv_pages( - mem_manager=mem, - mem_indexes=indexes, - page_indexes=page_ids, - kv_move_buffer=page.view(1, -1), - to_pages=False, - page_size=self.pd_kv_page_size, - round_tail=True, - ) - - -class DeepseekV4AttStatePageHelper: - """整块迁移请求的 SWA 和 compressor ring，保持物理槽位顺序。""" - - def __init__(self, mem_manager): - self.mem_manager = mem_manager - - def write_req_to_page(self, page_index, req_idx, dp_mems): - page = self.mem_manager.kv_move_buffer[page_index] - self._copy_req_state(page, dp_mems[0], req_idx, write=True) - - def read_page_to_req(self, page_index, req_idx, dp_mems): - page = self.mem_manager.kv_move_buffer[page_index] - for mem in dp_mems: - self._copy_req_state(page, mem, req_idx, write=False) - - def _copy_req_state(self, page, mem, req_idx, *, write): - data = page.reshape(-1) - offset = 0 - for buffer in (mem.req_to_swa, mem.req_to_c4_state, mem.req_to_c4_indexer_state, mem.req_to_c128_state): - if buffer is None: - continue - state = buffer[:, req_idx] - nbytes = state.numel() * state.element_size() - page_state = data[offset : offset + nbytes].view(state.dtype).view(state.shape) - if write: - page_state.copy_(state, non_blocking=True) - else: - state.copy_(page_state, non_blocking=True) - offset += nbytes diff --git a/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py deleted file mode 100644 index f28f400331..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py +++ /dev/null @@ -1,19 +0,0 @@ -from .glm5_next_mem_manager import Glm5NextMemManager -from .operator import LinearAttMemOperator - - -class FP8Glm5NextMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): - from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import ( - destindex_copy_kv_flashmla_fp8, - ) - - packed_kv = self.mem_manager.get_att_input_params(layer_index) - destindex_copy_kv_flashmla_fp8(kv, mem_index, packed_kv) - - -class FP8Glm5NextMemManager(Glm5NextMemManager): - operator_class = FP8Glm5NextMemOperator - - def get_att_input_params(self, layer_index): - return self._layer_buffer(layer_index)[:, :, : self.linear_config.FP8_MLA_KV_BYTES] diff --git a/lightllm/common/req_manager/deepseek4.py b/lightllm/common/req_manager/deepseek4.py deleted file mode 100644 index f90d9fc4b9..0000000000 --- a/lightllm/common/req_manager/deepseek4.py +++ /dev/null @@ -1,88 +0,0 @@ -import torch - -from .hybrid_base import HybridAttentionReqManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4StateCacheManager - - -class DeepseekV4ReqManager(HybridAttentionReqManager): - """持有跨 forward 的请求映射、SWA ring 和 compressor ring，并负责状态保存和恢复。 - - 当前 forward 的派生索引、页表和元数据由 infer_state 构建并持有， - 模型 workspace 提供投影与量化的计算 scratch。 - """ - - def __init__( - self, - max_request_num, - max_sequence_length, - mem_manager, - cache_config, - *, - mtp_step=0, - ): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.cache_config = cache_config - self.sliding_window = cache_config.sliding_window - swa_shape, c4_shape, indexer_shape, c128_shape = cache_config.get_runtime_state_shapes(mtp_step) - self.swa_capacity_tokens = swa_shape[1] - self.HOLD_SWA_SLOT = self.HOLD_REQUEST_ID * self.swa_capacity_tokens - self.c4_state_tokens = c4_shape[1] - self.c128_state_tokens = c128_shape[1] - self.req_to_swa = torch.zeros( - (cache_config.layer_num, max_request_num + 1, self.swa_capacity_tokens, cache_config.swa_slot_bytes), - dtype=torch.uint8, - device="cuda", - ) - self.req_to_c4_state = self._alloc_compressor_state(*c4_shape) - self.req_to_c4_indexer_state = self._alloc_compressor_state(*indexer_shape) - self.req_to_c128_state = self._alloc_compressor_state(*c128_shape) - - def _alloc_compressor_state(self, layer_count, token_capacity, state_dim): - buffer = torch.zeros( - (layer_count, self.max_request_num + 1, token_capacity, state_dim), dtype=torch.float32, device="cuda" - ) - buffer[:, self.HOLD_REQUEST_ID, :, state_dim // 2 :].fill_(float("-inf")) - return buffer - - def get_att_input_params(self, layer_index): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer, view_swa_kv - - return view_swa_kv(self.req_to_swa[layer_index], self.cache_config), get_compressed_kv_buffer( - self.mem_manager, layer_index - ) - - def create_small_page_cache_manager(self, size): - self.small_page_buffers = DeepseekV4StateCacheManager(size, self.cache_config) - return self.small_page_buffers - - def save_state(self, req, buffer_idx, state_cache_manager): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import gather_swa_state - - config = self.cache_config - swa, c4, indexer = state_cache_manager.get_state_cache(buffer_idx) - # 保存时 cur_kv_len 还是 chunk 起点；运行态对应本轮 prefill 终点。 - req_idx = req.req_idx - checkpoint_len = req.get_chuncked_input_token_len() - positions = torch.arange(checkpoint_len - config.sliding_window, checkpoint_len, device="cuda") - swa.copy_(gather_swa_state(self, req_idx, positions), non_blocking=True) - if not config.n_c4: - return - positions = torch.arange(checkpoint_len - config.short_compress_ratio, checkpoint_len, device="cuda") - state_slots = positions % self.c4_state_tokens - # compressor 已将 chunk 尾部写回 ring；大小页只保存当前终点。 - for target, runtime in ((c4, self.req_to_c4_state), (indexer, self.req_to_c4_indexer_state)): - target.copy_(runtime[:, req_idx].index_select(1, state_slots), non_blocking=True) - - def restore_state(self, req, state_cache_manager, buffer_idx): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import scatter_swa_state - - config = self.cache_config - swa, c4, indexer = state_cache_manager.get_state_cache(buffer_idx) - positions = torch.arange(req.cur_kv_len - config.sliding_window, req.cur_kv_len, device="cuda") - scatter_swa_state(self, req.req_idx, positions, swa.cuda(non_blocking=True)) - if config.n_c4: - positions = torch.arange(req.cur_kv_len - config.short_compress_ratio, req.cur_kv_len, device="cuda") - slots = positions % self.c4_state_tokens - self.req_to_c4_state[:, req.req_idx].index_copy_(1, slots, c4.cuda(non_blocking=True)) - self.req_to_c4_indexer_state[:, req.req_idx].index_copy_(1, slots, indexer.cuda(non_blocking=True)) - # Aligned checkpoints close C128 groups; the next group overwrites its tail. diff --git a/lightllm/common/state_cache_manager/deepseek4.py b/lightllm/common/state_cache_manager/deepseek4.py deleted file mode 100644 index 368bd8e222..0000000000 --- a/lightllm/common/state_cache_manager/deepseek4.py +++ /dev/null @@ -1,146 +0,0 @@ -import math -from dataclasses import dataclass - -import torch -import triton - -from lightllm.utils.envs_utils import get_env_start_args - -from .base import StateCacheManager - - -DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT = 256 - - -@dataclass -class DeepseekV4CacheConfig: - """V4 KV 布局和大小页状态缓存配置。""" - - compress_rates: tuple[int, ...] - head_dim: int - rope_head_dim: int - indexer_head_dim: int - sliding_window: int - - # Packed FP8 KV format shared by Flash and Pro. - quant_group_dim = 64 - scale_alignment_bytes = 8 - # 每个 SWA 物理页的 token 数，用于 KV/scale 打包和 ring 容量对齐。 - swa_page_token_num = 128 - compressed_kv_token_alignment = DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT - short_compress_ratio = 4 - long_compress_ratio = 128 - - def __post_init__(self): - self.compress_rates = tuple(self.compress_rates) - self.layer_num = len(self.compress_rates) - c4_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate == self.short_compress_ratio] - c128_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate == self.long_compress_ratio] - self.layer_to_c4_idx = {layer: slot for slot, layer in enumerate(c4_layers)} - self.layer_to_c128_idx = {layer: slot for slot, layer in enumerate(c128_layers)} - self.n_c4, self.n_c128 = len(self.layer_to_c4_idx), len(self.layer_to_c128_idx) - compressed_kv_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate] - self.layer_to_compressed_kv_idx = {layer: slot for slot, layer in enumerate(compressed_kv_layers)} - self.kv_data_bytes = self.head_dim + self.rope_head_dim - self.kv_scale_bytes = ( - triton.cdiv((self.head_dim - self.rope_head_dim) // self.quant_group_dim, self.scale_alignment_bytes) - * self.scale_alignment_bytes - ) - self.kv_bytes = self.kv_data_bytes + self.kv_scale_bytes - self.indexer_bytes = self.indexer_head_dim + torch.float32.itemsize - self.swa_slot_bytes = self.get_packed_kv_bytes(self.swa_page_token_num) // self.swa_page_token_num - self.compressed_kv_slot_bytes = self._get_compressed_kv_slot_bytes() - - @classmethod - def from_model_config(cls, config, args): - from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num - - layer_num = config["num_hidden_layers"] - compress_rates = list(config["compress_ratios"]) - if args.mtp_mode is not None: - layer_num += get_added_mtp_kv_layer_num() - # Runtime-added draft layers use sliding-window attention. - compress_rates.extend([0] * (layer_num - len(compress_rates))) - return cls( - compress_rates=tuple(compress_rates[:layer_num]), - head_dim=config["head_dim"], - rope_head_dim=config["qk_rope_head_dim"], - indexer_head_dim=config["index_head_dim"], - sliding_window=config["sliding_window"], - ) - - def get_packed_kv_bytes(self, token_count): - return triton.cdiv(token_count * self.kv_bytes, self.kv_data_bytes) * self.kv_data_bytes - - def _get_compressed_kv_slot_bytes(self): - block_token_num = self.compressed_kv_token_alignment - # 整块按 KV 数据行对齐，同时能均分到 token 槽位。 - block_bytes_alignment = math.lcm(block_token_num, self.kv_data_bytes) - slot_bytes_list = [] - for ratio in self.compress_rates: - if ratio == 0: - continue - compressed_token_num = block_token_num // ratio - block_bytes = self.get_packed_kv_bytes(compressed_token_num) - if ratio == self.short_compress_ratio: - block_bytes += compressed_token_num * self.indexer_bytes - block_bytes = triton.cdiv(block_bytes, block_bytes_alignment) * block_bytes_alignment - slot_bytes_list.append(block_bytes // block_token_num) - return tuple(slot_bytes_list) - - def get_swa_state_shape(self): - return (self.layer_num, self.sliding_window, self.kv_bytes) - - def get_c4_state_shape(self): - return (self.n_c4, self.short_compress_ratio, 4 * self.head_dim) - - def get_indexer_state_shape(self): - return (self.n_c4, self.short_compress_ratio, 4 * self.indexer_head_dim) - - def get_runtime_state_shapes(self, mtp_step): - """每个请求的完整 ring；P/D 使用相同布局，可直接复制槽位。""" - swa_tokens = triton.cdiv(self.sliding_window + 2 * mtp_step, self.swa_page_token_num) * self.swa_page_token_num - c4_tokens = 2 * self.short_compress_ratio + mtp_step - c128_tokens = self.long_compress_ratio + mtp_step - return ( - (self.layer_num, swa_tokens, self.swa_slot_bytes), - (self.n_c4, c4_tokens, 4 * self.head_dim), - (self.n_c4, c4_tokens, 4 * self.indexer_head_dim), - (self.n_c128, c128_tokens, 2 * self.head_dim), - ) - - def get_cpu_cache_full_att_bytes(self): - args = get_env_start_args() - return args.cpu_cache_token_page_size * sum(self.compressed_kv_slot_bytes) - - def get_cpu_cache_big_page_bytes(self): - # MLA KV and compressor state are replicated across TP ranks. - swa_bytes = math.prod(self.get_swa_state_shape()) - c4_bytes = math.prod(self.get_c4_state_shape()) * torch.float32.itemsize - indexer_bytes = math.prod(self.get_indexer_state_shape()) * torch.float32.itemsize - return self.get_cpu_cache_full_att_bytes() + swa_bytes + c4_bytes + indexer_bytes - - -class DeepseekV4StateCacheManager(StateCacheManager): - """Pinned continuation checkpoints; compressed KV stays in token pages.""" - - def __init__(self, size, cache_config: DeepseekV4CacheConfig, keep_num=0): - super().__init__(size, keep_num) - self.cache_config = cache_config - self.swa_shape = cache_config.get_swa_state_shape() - self.c4_shape = cache_config.get_c4_state_shape() - self.indexer_shape = cache_config.get_indexer_state_shape() - self.swa_nbytes = math.prod(self.swa_shape) - self.c4_nbytes = math.prod(self.c4_shape) * torch.float32.itemsize - indexer_nbytes = math.prod(self.indexer_shape) * torch.float32.itemsize - self.buffer = torch.empty( - (size, self.swa_nbytes + self.c4_nbytes + indexer_nbytes), dtype=torch.uint8, pin_memory=True - ) - - def get_state_cache(self, buffer_idx): - data = self.buffer[buffer_idx] - swa = data[: self.swa_nbytes].view(self.swa_shape) - indexer_start = self.swa_nbytes + self.c4_nbytes - c4 = data[self.swa_nbytes : indexer_start].view(torch.float32).view(self.c4_shape) - indexer = data[indexer_start:].view(torch.float32).view(self.indexer_shape) - return swa, c4, indexer diff --git a/lightllm/models/deepseek_v4/config.py b/lightllm/models/deepseek_v4/config.py new file mode 100644 index 0000000000..a57e89bf4e --- /dev/null +++ b/lightllm/models/deepseek_v4/config.py @@ -0,0 +1,85 @@ +"""模型选择 V4 的逻辑层和压缩方案，State 配置只记录确定的几何参数。""" + + +from dataclasses import dataclass +from typing import ClassVar + +from lightllm.common.attention_state.bounded_att.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig +from lightllm.common.attention_state.full_att.compressed_att_state import CompressedAttentionConfig + + +@dataclass(frozen=True) +class DSV4AttentionConfig: + # 当前模型拥有的逻辑层 ID，顺序对应 SWA 的物理 layer 轴；main/draft 分别创建。 + logical_layer_ids: tuple[int, ...] + # 与 logical_layer_ids 逐项对应；0 为纯 SWA，4/128 为同时启用对应压缩器。 + compress_ratios: tuple[int, ...] + # 未编码的共享 KV 向量宽度，包含以 BF16 存储的 RoPE 部分。 + head_dim: int + # RoPE 特征宽度；其余特征按每 64 个元素一组编码为 FP8。 + rope_head_dim: int + # C4 indexer key 的特征宽度；每个 key 另有一个 FP32 scale。 + indexer_head_dim: int + # 每份可恢复 SWA 快照包含的逻辑 token 数，不包含 MTP 候选位置。 + sliding_window: int + # 存储在各 TP rank 完整复制，只有 rank 0 向共享的 CPU/PD 区域写出。 + tp_rank: int = 0 + tp_world_size: int = 1 + + # 压缩 KV 按原始 token 的 256-token 块存储，KV 与 indexer 不得跨块拆开。 + compressed_kv_token_alignment: ClassVar[int] = 256 + + def __post_init__(self) -> None: + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert len(self.logical_layer_ids) == len(self.compress_ratios) + assert all(ratio in (0, 4, 128) for ratio in self.compress_ratios) + assert self.head_dim > self.rope_head_dim > 0 and (self.head_dim - self.rope_head_dim) % 64 == 0 + assert self.indexer_head_dim > 0 and self.sliding_window > 0 + assert 0 <= self.tp_rank < self.tp_world_size + + def create_swa_config(self) -> FP8E8M0SlidingWindowConfig: + """模型选择窗口及向量维度，具体编码和物理页布局由原子配置负责。""" + return FP8E8M0SlidingWindowConfig( + logical_layer_ids=self.logical_layer_ids, + window_size=self.sliding_window, + kv_dim=self.head_dim, + rope_dim=self.rope_head_dim, + tp_rank=self.tp_rank, + tp_world_size=self.tp_world_size, + ) + + def create_compressed_config(self, compress_ratio: int) -> CompressedAttentionConfig: + """按模型层组合创建一种压缩历史配置，供 State 创建和 backend 使用。""" + return CompressedAttentionConfig( + logical_layer_groups=tuple( + (layer,) + for layer, ratio in zip(self.logical_layer_ids, self.compress_ratios) + if ratio == compress_ratio + ), + compress_ratio=compress_ratio, + kv_dim=self.head_dim, + rope_dim=self.rope_head_dim, + block_token_size=self.compressed_kv_token_alignment, + indexer_dim=self.indexer_head_dim if compress_ratio == 4 else 0, + tp_rank=self.tp_rank, + tp_world_size=self.tp_world_size, + ) + + +def create_attention_config( + model_config: dict[str, object], + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, +) -> DSV4AttentionConfig: + ratios = model_config["compress_ratios"] + return DSV4AttentionConfig( + logical_layer_ids=model_logical_layer_ids, + compress_ratios=tuple(ratios[layer] if layer < len(ratios) else 0 for layer in model_logical_layer_ids), + head_dim=model_config["head_dim"], + rope_head_dim=model_config["qk_rope_head_dim"], + indexer_head_dim=model_config["index_head_dim"], + sliding_window=model_config["sliding_window"], + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index 14f751394e..b0abc2f24b 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -1,9 +1,13 @@ import dataclasses +from typing import Literal import torch from lightllm.common.basemodel import InferStateInfo -from lightllm.common.req_manager import DeepseekV4ReqManager -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState +from lightllm.common.attention_state.full_att.compressed_att_state import ( + CompressedAttentionState, + CompressedAttentionLayerView, +) C4_PREFILL_LOGITS_BUDGET_BYTES = 512 * 1024 * 1024 @@ -23,11 +27,39 @@ class DeepseekV4InferStateInfo(InferStateInfo): """Per-forward inputs, derived indices and attention metadata. Indices and metadata are built for this batch and reused by its layers. - Request-owned SWA and compressor rings remain in the request manager. + SWA and compressor rings are owned by the model's bounded State. """ - req_manager: DeepseekV4ReqManager - mem_manager: DeepseekV4MemoryManager + @property + def swa_state(self) -> FP8E8M0SlidingWindowState: + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter="swa_", + logical_layer_id=self.first_logical_layer_id, + state_class=FP8E8M0SlidingWindowState, + ) + + def get_compressor_state( + self, logical_layer_id: int, target: Literal["kv", "indexer"] = "kv" + ) -> WindowCompressorState: + """按当前层和用途选择独立的窗口压缩器。""" + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter=f"compressor_{target}_", + logical_layer_id=logical_layer_id, + state_class=WindowCompressorState, + ) + + def get_compressed_kv_state(self, logical_layer_id: int) -> CompressedAttentionState: + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter="compressed", + logical_layer_id=logical_layer_id, + state_class=CompressedAttentionState, + ) + + def get_compressed_kv_view(self, logical_layer_id: int) -> CompressedAttentionLayerView: + return self.get_compressed_kv_state(logical_layer_id).get_layer_view(logical_layer_id) def __init__(self): super().__init__() @@ -56,9 +88,21 @@ def __init__(self): self.c4_indexer_metadata: C4IndexerMetadata | None = None def init_some_extra_state(self, model): + self.first_logical_layer_id = model.layers_infer[0].layer_num_ + self.first_c4_logical_layer_id = None + has_c128_layer = False + ratios = model.config["compress_ratios"] + for layer in model.layers_infer: + compress_ratio = ratios[layer.layer_num_] if layer.layer_num_ < len(ratios) else 0 + if compress_ratio == 4 and self.first_c4_logical_layer_id is None: + self.first_c4_logical_layer_id = layer.layer_num_ + if compress_ratio == 128: + has_c128_layer = True self.c4_indexer_metadata = None # Rebuild for each eager forward or graph capture. self.dsv4_c4_indices = None self.dsv4_c4_lengths = None + self.dsv4_c128_indices = None + self.dsv4_c128_lengths = None super().init_some_extra_state(model) # sets position_ids, b_q_seq_len, b_q_start_loc (prefill) if self.is_prefill: self.prefill_att_state.prepare_swa_workspace(self) @@ -84,14 +128,18 @@ def init_some_extra_state(self, model): self.dsv4_workspace = model.dsv4_workspace self._init_swa_metadata(model) + # 当前模型没有 C128 层时，不分配或构建这组索引；draft 只检查自身拥有的层。 + if not has_c128_layer: + return + from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices - long_ratio = self.req_manager.cache_config.long_compress_ratio + c128_compress_ratio = 128 max_kv_seq_len = self.max_kv_seq_len # Replay copies C128 indices into the captured state, so their widths must match. if not self.is_prefill and model.graph is not None and model.graph.can_run(self.batch_size, max_kv_seq_len): max_kv_seq_len = model.graph.graph_max_len_in_batch - c128_entries = max(1, max_kv_seq_len // long_ratio) + c128_entries = max(1, max_kv_seq_len // c128_compress_ratio) c128_cap = ((c128_entries + 63) // 64) * 64 self.dsv4_c128_indices = torch.empty((pos.numel(), c128_cap), dtype=torch.int32, device=pos.device) self.dsv4_c128_lengths = torch.empty_like(pos, dtype=torch.int32) @@ -99,7 +147,7 @@ def init_some_extra_state(self, model): self.dsv4_sparse_req_idx, self.position_ids, self.req_manager.req_to_token_indexs, - long_ratio, + c128_compress_ratio, self.dsv4_c128_indices, self.dsv4_c128_lengths, ) @@ -113,7 +161,7 @@ def _init_swa_metadata(self, model): pos = self.position_ids has_vision = model.has_vision - swa_width = self.req_manager.cache_config.sliding_window + swa_width = self.swa_state.config.window_size if has_vision and self.is_prefill: swa_width += model.config["vision_max_n_token"] self.dsv4_swa_indices = torch.empty((pos.numel(), swa_width), dtype=torch.int32, device=pos.device) @@ -141,11 +189,11 @@ def _init_swa_metadata(self, model): build_swa_indices( token_req_idx=self.dsv4_sparse_req_idx, position_ids=self.position_ids, - req_manager=self.req_manager, + swa_state=self.swa_state, swa_indices=self.dsv4_swa_indices, swa_lengths=self.dsv4_swa_lengths, swa_write_slots=self.dsv4_swa_write_slots, - sliding_window=self.req_manager.cache_config.sliding_window, + sliding_window=self.swa_state.config.window_size, image_left=self.dsv4_image_left, image_right=self.dsv4_image_right, prefill_state=self.prefill_att_state if self.is_prefill else None, @@ -180,7 +228,7 @@ def get_c4_indexer_metadata(self, c4_cap): self.req_manager.req_to_token_indexs, self.req_manager.HOLD_REQUEST_ID, ) - page_size = self.req_manager.cache_config.compressed_kv_token_alignment // 4 + page_size = self.get_compressed_kv_state(self.first_c4_logical_layer_id).config.block_token_size // 4 rows_per_chunk = max(1, self.position_ids.numel()) if self.is_prefill: # DeepGEMM rounds the logits width to a multiple of 256. diff --git a/lightllm/models/deepseek_v4/layer_infer/compressor.py b/lightllm/models/deepseek_v4/layer_infer/compressor.py index f9841302cf..4938d1f823 100644 --- a/lightllm/models/deepseek_v4/layer_infer/compressor.py +++ b/lightllm/models/deepseek_v4/layer_infer/compressor.py @@ -18,7 +18,9 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int, is_ self.layer_idx_ = layer_idx self.tp_world_size_ = tp_world_size self.is_in_indexer = is_in_indexer - self.compress_ratio = network_config["compress_ratios"][layer_idx] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_idx] if layer_idx < len(ratios) else 0 self.head_dim = network_config["index_head_dim" if is_in_indexer else "head_dim"] self.qk_rope_head_dim = network_config["qk_rope_head_dim"] self.eps = network_config["rms_norm_eps"] @@ -92,35 +94,25 @@ def fused_compress( is_in_indexer: bool = False, out_buffer: torch.Tensor = None, ): - mem_manager = infer_state.mem_manager - req_manager = infer_state.req_manager + state = infer_state.get_compressor_state(logical_layer_id=layer_idx, target="indexer" if is_in_indexer else "kv") + req_to_compressor_state = state.get_layer_view(logical_layer_id=layer_idx) compressed_kv_slots = torch.where( (infer_state.position_ids + 1) % compress_ratio == 0, infer_state.mem_index.reshape(-1) // compress_ratio, -1, ) if is_in_indexer: - req_to_compressor_state = req_manager.req_to_c4_indexer_state[ - req_manager.cache_config.layer_to_c4_idx[layer_idx] - ] if out_buffer is None: out_buffer = torch.empty( - (infer_state.mem_index.numel(), req_manager.cache_config.indexer_head_dim), + (infer_state.mem_index.numel(), head_dim), dtype=torch.bfloat16, device=infer_state.mem_index.device, ) out_page_size = 1 else: - if compress_ratio == 4: - req_to_compressor_state = req_manager.req_to_c4_state[req_manager.cache_config.layer_to_c4_idx[layer_idx]] - else: - req_to_compressor_state = req_manager.req_to_c128_state[ - req_manager.cache_config.layer_to_c128_idx[layer_idx] - ] - out_page_size = req_manager.cache_config.compressed_kv_token_alignment // compress_ratio - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer - - out_buffer = get_compressed_kv_buffer(mem_manager, layer_idx) + compressed = infer_state.get_compressed_kv_state(logical_layer_id=layer_idx) + out_page_size = compressed.config.block_token_size // compress_ratio + out_buffer = compressed.get_layer_view(logical_layer_id=layer_idx).kv if kv_score.shape[0] == 0: return out_buffer @@ -160,7 +152,7 @@ def fused_compress( CANDIDATE_DIM=candidate_dim, STATE_DIM=state_dim, COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=compress_ratio * (2 if is_c4 else 1), + WINDOW_SIZE=state.config.window_size, IS_C4=is_c4, IS_PREFILL=infer_state.is_prefill, STATE_CAPACITY_TOKENS=state_capacity_tokens, diff --git a/lightllm/models/deepseek_v4/layer_infer/indexer.py b/lightllm/models/deepseek_v4/layer_infer/indexer.py index cdec83bfdb..6c0d8b14aa 100644 --- a/lightllm/models/deepseek_v4/layer_infer/indexer.py +++ b/lightllm/models/deepseek_v4/layer_infer/indexer.py @@ -18,7 +18,9 @@ class DeepseekV4IndexInfer(BaseLayerInfer): def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.layer_idx_ = layer_idx - self.compress_ratio = network_config["compress_ratios"][layer_idx] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_idx] if layer_idx < len(ratios) else 0 self.index_topk = network_config["index_topk"] self.index_head_dim = network_config["index_head_dim"] self.index_n_heads = network_config["index_n_heads"] @@ -64,7 +66,6 @@ def write_indexer_k( if use_custom_tensor_manager: hadamard_out = self.alloc_tensor(scratch.shape, scratch.dtype, device=scratch.device) scratch = hadamard_transform(scratch, scale=self.index_head_dim ** -0.5, out=hadamard_out) - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_indexer_k_buffer from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k import ( destindex_copy_indexer_k, ) @@ -73,8 +74,8 @@ def write_indexer_k( scratch, infer_state.mem_index.reshape(-1), infer_state.position_ids, - get_indexer_k_buffer(infer_state.mem_manager, self.layer_idx_), - infer_state.mem_manager.cache_config.compressed_kv_token_alignment // self.compress_ratio, + infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k, + infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size // self.compress_ratio, ) def build_metadata( @@ -170,11 +171,9 @@ def _c4_indices(self, infer_state: DeepseekV4InferStateInfo, idx_q_fp8, weights) return slots.unsqueeze(1), lengths paged_meta = infer_state.get_c4_indexer_metadata(c4_cap) - mem_manager = infer_state.mem_manager - page_size = mem_manager.cache_config.compressed_kv_token_alignment // self.compress_ratio - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_indexer_k_buffer + page_size = infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size // self.compress_ratio - packed = get_indexer_k_buffer(mem_manager, self.layer_idx_) + packed = infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k indexer_k_cache = packed.view( packed.shape[0], page_size, diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index f52cd019fb..c5fe0495ee 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -30,7 +30,9 @@ def __init__(self, layer_num, network_config): self.hc_mult = network_config["hc_mult"] self.sinkhorn_iters = network_config["hc_sinkhorn_iters"] self.hc_eps = network_config["hc_eps"] - self.compress_ratio = network_config["compress_ratios"][layer_num] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_num] if layer_num < len(ratios) else 0 self.is_hash = layer_num < network_config["num_hash_layers"] self.has_vision = network_config.get("vision_n_layers", 0) > 0 self.vocab_size = network_config["vocab_size"] @@ -334,7 +336,7 @@ def _select_rope(self, infer_state: DeepseekV4InferStateInfo): def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa - req_manager = infer_state.req_manager + state = infer_state.swa_state att_state = infer_state.prefill_att_state if infer_state.is_prefill else infer_state.decode_att_state swa, _ = att_state.backend.get_att_input_params(infer_state, self.layer_num_) # 一次完成 RMSNorm、RoPE、FP8 打包和 SWA 写入。 @@ -347,8 +349,8 @@ def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): swa_write_slots=infer_state.dsv4_swa_write_slots, swa_runtime_write_slots=infer_state.dsv4_swa_runtime_write_slots, swa_kv_buffer=swa, - swa_page_token_num=req_manager.cache_config.swa_page_token_num, - req_to_swa=req_manager.req_to_swa[self.layer_num_], + swa_page_token_num=state.config.page_token_size, + req_to_swa=state.get_layer_view(logical_layer_id=self.layer_num_), ) def _get_qkv( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index 9a6b75cc2d..5cd25f2466 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -1,10 +1,10 @@ +from __future__ import annotations + import os +from typing import TYPE_CHECKING import torch from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.req_manager import DeepseekV4ReqManager -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import ( DeepseekV4PreAndPostLayerWeight, ) @@ -37,6 +37,12 @@ from lightllm.utils.log_utils import init_logger from lightllm.distributed.communication_op import dist_group_manager +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + logger = init_logger(__name__) @@ -61,8 +67,7 @@ def add(self, layer_index: int, hidden): class DeepseekV4TpPartModel(LlamaTpPartModel): - req_manager: DeepseekV4ReqManager - mem_manager: DeepseekV4MemoryManager + is_hybrid_att_model = True has_vision = False pre_and_post_weight_class = DeepseekV4PreAndPostLayerWeight @@ -90,30 +95,110 @@ def _verify_params(self): raise ValueError("DeepSeek-V4 preselected MoE routing does not yet support redundant experts") return - def _init_req_manager(self): - cache_config = DeepseekV4CacheConfig.from_model_config(self.config, self.args) - self.req_manager = DeepseekV4ReqManager( - self.max_req_num, - self.max_seq_length, - None, - cache_config, - mtp_step=self.args.mtp_step, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list[FullAttentionState]: + """分别创建 C4/C128 定义；逻辑索引保留原始 token 域，不提前分配存储。""" + from lightllm.common.attention_state.full_att.create_compressed_att_state import create_compressed_att_state + from .config import create_attention_config + + assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" + config = create_attention_config( + model_config=model_config, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ) + states = [] + for ratio in (4, 128): + if ratio not in config.compress_ratios: + continue + states.append(create_compressed_att_state( + config=config.create_compressed_config(compress_ratio=ratio), + page_size=args.page_size, + device="meta", + )) + return states - def _init_mem_manager(self): - DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(self) - cache_config = self.req_manager.cache_config - self.mem_manager = DeepseekV4MemoryManager( - self.max_total_token_num, - dtype=self.data_type, - head_num=1, - head_dim=cache_config.head_dim, - layer_num=cache_config.layer_num, - mem_fraction=self.mem_fraction, - cache_config=cache_config, + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list[BoundedAttentionState]: + """创建本模型的 SWA/compressor 请求运行态，大小页容量由管理器后续确定。""" + from lightllm.common.attention_state.bounded_att import ( + WindowCompressorStateConfig, + create_window_compressor_state, + create_sliding_window_state, ) - self.req_manager.bind_mem_manager(self.mem_manager) - return + from .config import create_attention_config + + config = create_attention_config( + model_config=model_config, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + states = [create_sliding_window_state( + config=config.create_swa_config(), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + )] + # 模型决定窗口组合：CSA 使用 W8/S4，HCA 使用 W128/S128；State 本身不识别 C4/C128。 + for ratio, window_size in ((4, 8), (128, 128)): + layers = tuple( + layer for layer, compress_ratio in zip(model_logical_layer_ids, config.compress_ratios) + if compress_ratio == ratio + ) + if not layers: + continue + targets = ( + (("kv", config.head_dim), ("indexer", config.indexer_head_dim)) + if ratio == 4 else (("kv", config.head_dim),) + ) + for target, head_dim in targets: + states.append(create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=layers, + window_size=window_size, + compress_stride=ratio, + candidate_dim=head_dim * (2 if window_size > ratio else 1), + target=target, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + )) + return states + + def prepare_attention_states(self) -> None: + super().prepare_attention_states() + # 在统一容量规划前预留长 prefill 的 SWA 临时区，避免 KV 分配挤占计算所需空间。 + DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(self) + + def _init_quant(self): + if not self.is_mtp_draft_model: + return super()._init_quant() + from lightllm.common.quantization import Quantcfg + + # 权重量化按 checkpoint 层号查询；draft 的资源层数仍只描述自身。 + config = dict(self.config, n_layer=len(self.config["compress_ratios"])) + self.quant_cfg = Quantcfg(config, self.quant_type, self.quant_cfg_path, self.expert_dtype) def _init_att_backend(self): args = get_env_start_args() diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py index ed2e129396..5e0c6dc492 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py +++ b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py @@ -53,14 +53,11 @@ def _build_dspark_swa_indices_kernel( def build_dspark_swa_indices( token_req_idx: torch.Tensor, position_ids: torch.Tensor, - req_manager, + swa_state, swa_indices: torch.Tensor, swa_lengths: torch.Tensor, swa_write_slots: torch.Tensor, - sliding_window: int, block_size: int, - hold_req_id: int, - hold_swa_slot: int, ): """Build ``history SWA + complete draft block`` indices for every DSpark query row.""" @@ -76,10 +73,10 @@ def build_dspark_swa_indices( swa_indices.stride(0), swa_lengths, swa_write_slots, - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - HOLD_REQUEST_ID=hold_req_id, - HOLD_SWA_SLOT=hold_swa_slot, - SLIDING_WINDOW=sliding_window, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + HOLD_SWA_SLOT=swa_state.HOLD_SWA_SLOT, + SLIDING_WINDOW=swa_state.config.window_size, BLOCK_SIZE=block_size, WIDTH=width, BLOCK_W=triton.next_power_of_2(width), diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py b/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py index 11704de3c9..f46eaaf1cb 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py +++ b/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py @@ -138,7 +138,7 @@ def _build_swa_indices_kernel( def build_swa_indices( token_req_idx: torch.Tensor, position_ids: torch.Tensor, - req_manager, + swa_state, swa_indices: torch.Tensor, swa_lengths: torch.Tensor, swa_write_slots: torch.Tensor, @@ -173,14 +173,14 @@ def build_swa_indices( token_to_batch_idx, b_seq_len, WRITE_SWA_RUNTIME=swa_runtime_write_slots is not None, - SWA_PAGE_TOKEN_NUM=req_manager.cache_config.swa_page_token_num, + SWA_PAGE_TOKEN_NUM=swa_state.config.page_token_size, SLIDING_WINDOW=sliding_window, WIDTH=width, HAS_IMAGE=image_left is not None, BLOCK_W=triton.next_power_of_2(width), - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, - HOLD_SWA_SLOT=prefill_state.HOLD_SWA_SLOT if prefill_state is not None else req_manager.HOLD_SWA_SLOT, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + HOLD_SWA_SLOT=prefill_state.HOLD_SWA_SLOT if prefill_state is not None else swa_state.HOLD_SWA_SLOT, IS_PREFILL=prefill_state is not None, num_warps=4, ) diff --git a/lightllm/models/deepseek_v4/triton_kernel/cache_io.py b/lightllm/models/deepseek_v4/triton_kernel/cache_io.py deleted file mode 100644 index c012ca1b92..0000000000 --- a/lightllm/models/deepseek_v4/triton_kernel/cache_io.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Model-side packed views and token-indexed cache transfers.""" - -import triton -import triton.language as tl - -from lightllm.utils.envs_utils import get_env_start_args - - -def view_swa_kv(swa_buffer, config): - return swa_buffer.view(-1, config.get_packed_kv_bytes(config.swa_page_token_num)) - - -def get_compressed_kv_buffer(mem_manager, layer_index): - kv_buffer = mem_manager.get_att_input_params(layer_index) - if kv_buffer is None: - return None - config = mem_manager.cache_config - compress_ratio = config.compress_rates[layer_index] - packed_kv_pages = kv_buffer.view(-1, config.compressed_kv_token_alignment * kv_buffer.shape[-1]) - return packed_kv_pages[:, : config.get_packed_kv_bytes(config.compressed_kv_token_alignment // compress_ratio)] - - -def get_indexer_k_buffer(mem_manager, layer_index): - config = mem_manager.cache_config - kv_buffer = mem_manager.get_att_input_params(layer_index) - packed_kv_pages = kv_buffer.view(-1, config.compressed_kv_token_alignment * kv_buffer.shape[-1]) - entries = config.compressed_kv_token_alignment // config.short_compress_ratio - begin = config.get_packed_kv_bytes(entries) - return packed_kv_pages[:, begin : begin + entries * config.indexer_bytes] - - -@triton.jit -def _copy_compressed_kv_kernel( - kv_buffer, - kv_move_buffer, - mem_indexes, - page_indexes, - layer_byte_offset, - TOKEN_NUM: tl.constexpr, - COPY_NUM: tl.constexpr, - PAGE_TOKENS: tl.constexpr, - PAGE_STRIDE: tl.constexpr, - SLOT_BYTES: tl.constexpr, - TO_PAGES: tl.constexpr, - ROUND_TAIL: tl.constexpr, - BLOCK_BYTES: tl.constexpr, -): - kv_buffer = kv_buffer.to(tl.pointer_type(tl.uint8)) - tokens = tl.program_id(0) * 32 + tl.arange(0, 32) - slot = tl.load(mem_indexes + tokens, tokens < TOKEN_NUM, -1).to(tl.int64) - if ROUND_TAIL: - # Packed page tails include scales for earlier tokens. PD also carries - # the unused token slots of its final allocated compression block. - last = tl.load(mem_indexes + TOKEN_NUM - 1).to(tl.int64) - slot = tl.where(tokens < TOKEN_NUM, slot, last + tokens - TOKEN_NUM + 1) - page = tokens // PAGE_TOKENS - valid = (tokens < COPY_NUM) & (slot >= 0) - page_id = tl.load(page_indexes + page, tokens < COPY_NUM, 0).to(tl.int64) - columns = tl.arange(0, BLOCK_BYTES) - kv_ptr = kv_buffer + slot[:, None] * SLOT_BYTES + columns[None, :] - page_ptr = ( - kv_move_buffer - + page_id[:, None] * PAGE_STRIDE - + layer_byte_offset - + (tokens % PAGE_TOKENS)[:, None] * SLOT_BYTES - + columns[None, :] - ) - mask = valid[:, None] & (columns[None, :] < SLOT_BYTES) - if TO_PAGES: - tl.store(page_ptr, tl.load(kv_ptr, mask, 0), mask) - else: - tl.store(kv_ptr, tl.load(page_ptr, mask, 0), mask) - - -def copy_compressed_kv_pages( - mem_manager, mem_indexes, page_indexes, kv_move_buffer, *, to_pages, page_size, round_tail=False -): - if not mem_indexes.numel(): - return - config = mem_manager.cache_config - copy_num = ( - triton.cdiv(mem_indexes.numel(), config.compressed_kv_token_alignment) * config.compressed_kv_token_alignment - if round_tail - else mem_indexes.numel() - ) - layer_byte_offset = 0 - for kv_buffer in mem_manager.kv_buffer: - slot_bytes = kv_buffer.shape[-1] - _copy_compressed_kv_kernel[(triton.cdiv(copy_num, 32),)]( - kv_buffer.data_ptr(), - kv_move_buffer, - mem_indexes, - page_indexes, - layer_byte_offset, - TOKEN_NUM=mem_indexes.numel(), - COPY_NUM=copy_num, - PAGE_TOKENS=page_size, - PAGE_STRIDE=kv_move_buffer.stride(0), - SLOT_BYTES=slot_bytes, - TO_PAGES=to_pages, - ROUND_TAIL=round_tail, - BLOCK_BYTES=triton.next_power_of_2(slot_bytes), - num_warps=4, - ) - layer_byte_offset += page_size * slot_bytes - - -@triton.jit -def _copy_cpu_cache_kernel( - kv_buffer, - big_page_state_buffer, - mem_indexes, - big_page_buffer_ids, - page_indexes, - page_readies, - cpu_cache_tensor, - TOKEN_NUM: tl.constexpr, - PAGE_TOKENS: tl.constexpr, - PAGE_BYTES: tl.constexpr, - STATE_BYTES: tl.constexpr, - SLOT_BYTES: tl.constexpr, - TO_CPU_CACHE: tl.constexpr, - HAS_READIES: tl.constexpr, - BLOCK: tl.constexpr, -): - page = tl.program_id(0) - if HAS_READIES: - if tl.load(page_readies + page): - return - - cpu_page = cpu_cache_tensor + tl.load(page_indexes + page).to(tl.int64) * PAGE_BYTES - block_start, block_step = tl.program_id(1), tl.num_programs(1) - layer_offset = 0 - # CPU 页前半段按层保存压缩 KV 和 indexer K，碎页只复制实际存在的 token 槽位。 - for layer_index in tl.static_range(len(SLOT_BYTES)): - slot_bytes = SLOT_BYTES[layer_index] - layer_bytes = PAGE_TOKENS * slot_bytes - for block in range(block_start, tl.cdiv(layer_bytes, BLOCK), block_step): - offsets = block * BLOCK + tl.arange(0, BLOCK) - tokens = page * PAGE_TOKENS + offsets // slot_bytes - mask = (offsets < layer_bytes) & (tokens < TOKEN_NUM) - slots = tl.load(mem_indexes + tokens, mask, -1).to(tl.int64) - mask &= slots >= 0 - gpu = kv_buffer[layer_index] + slots * slot_bytes + offsets % slot_bytes - cpu = cpu_page + layer_offset + offsets - if TO_CPU_CACHE: - tl.store(cpu, tl.load(gpu, mask, 0), mask) - else: - tl.store(gpu, tl.load(cpu, mask, 0), mask) - layer_offset += layer_bytes - - # 页尾的 SWA、C4 和 indexer 状态已连续存储，作为一整块复制。 - state_id = tl.load(big_page_buffer_ids + page).to(tl.int64) - state = big_page_state_buffer + state_id * STATE_BYTES - for block in range(block_start, tl.cdiv(STATE_BYTES, BLOCK), block_step): - offsets = block * BLOCK + tl.arange(0, BLOCK) - mask = offsets < STATE_BYTES - cpu = cpu_page + layer_offset + offsets - if TO_CPU_CACHE: - tl.store(cpu, tl.load(state + offsets, mask, 0), mask) - else: - tl.store(state + offsets, tl.load(cpu, mask, 0), mask) - - -def _copy_cpu_cache_pages( - mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor, *, to_cpu_cache, page_readies=None -): - assert len(big_page_buffer_ids) == len(page_indexes) - cpu_cache_pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1) - _copy_cpu_cache_kernel[(len(page_indexes), 12)]( - tuple(mem_manager.kv_buffer), - mem_manager.big_page_buffers.buffer, - mem_indexes, - big_page_buffer_ids, - page_indexes, - page_readies, - cpu_cache_pages, - TOKEN_NUM=len(mem_indexes), - PAGE_TOKENS=get_env_start_args().cpu_cache_token_page_size, - PAGE_BYTES=cpu_cache_pages.stride(0), - STATE_BYTES=mem_manager.big_page_buffers.buffer.shape[-1], - SLOT_BYTES=mem_manager.cache_config.compressed_kv_slot_bytes, - TO_CPU_CACHE=to_cpu_cache, - HAS_READIES=page_readies is not None, - BLOCK=4096, - num_warps=4, - ) - - -def copy_kv_buffer_to_cpu_cache( - mem_manager, mem_indexes, page_indexes, page_readies, big_page_buffer_ids, cpu_cache_tensor -): - _copy_cpu_cache_pages( - mem_manager, - mem_indexes, - big_page_buffer_ids, - page_indexes, - cpu_cache_tensor, - to_cpu_cache=True, - page_readies=page_readies, - ) - - -def copy_cpu_cache_to_kv_buffer(mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor): - _copy_cpu_cache_pages( - mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor, to_cpu_cache=False - ) - - -@triton.jit -def _copy_state_cache_kernel(src_state, dst_state, NBYTES: tl.constexpr, BLOCK: tl.constexpr): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < NBYTES - tl.store(dst_state + offsets, tl.load(src_state + offsets, mask, 0), mask) - - -def copy_state_cache(src_state, dst_state): - """在当前 CUDA stream 上整块复制 pinned 状态，保持与 offload/load 的执行顺序。""" - assert src_state.shape == dst_state.shape - _copy_state_cache_kernel[(triton.cdiv(src_state.numel(), 4096),)]( - src_state, dst_state, NBYTES=src_state.numel(), BLOCK=4096, num_warps=4 - ) diff --git a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py index e94e92c381..9742e7cae8 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py +++ b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py @@ -1,4 +1,4 @@ -"""SWA 物理索引、prefill 临时区，以及大小页 SWA 状态的保存和恢复。""" +"""SWA 物理索引与 prefill workspace 的布局准备和前缀搬运。""" import torch import triton @@ -27,21 +27,21 @@ def get_swa_token_index( return tl.where(req_idx == HOLD_REQUEST_ID, HOLD_SWA_SLOT + position % SWA_PAGE_TOKEN_NUM, slot) -def get_swa_token_indices(req_manager, req_idx, position_ids, *, prefill_state=None): - config = req_manager.cache_config - swa_token_indices = req_idx * req_manager.swa_capacity_tokens + position_ids % req_manager.swa_capacity_tokens - hold_slot = req_manager.HOLD_SWA_SLOT +def get_swa_token_indices(swa_state, req_idx, position_ids, *, prefill_state=None): + config = swa_state.config + swa_token_indices = req_idx * swa_state.swa_capacity_tokens + position_ids % swa_state.swa_capacity_tokens + hold_slot = swa_state.HOLD_SWA_SLOT if prefill_state is not None: ready = prefill_state.prefill_token_ranges[req_idx, 0] first = ( - (ready - config.sliding_window + 1).clamp_min(0) // config.swa_page_token_num * config.swa_page_token_num + (ready - config.window_size + 1).clamp_min(0) // config.page_token_size * config.page_token_size ) workspace = prefill_state.swa_prefill_workspace_offsets[req_idx] swa_token_indices = workspace + position_ids - first hold_slot = prefill_state.HOLD_SWA_SLOT return torch.where( - torch.as_tensor(req_idx, device=position_ids.device) == req_manager.HOLD_REQUEST_ID, - hold_slot + position_ids % config.swa_page_token_num, + torch.as_tensor(req_idx, device=position_ids.device) == swa_state.HOLD_REQUEST_ID, + hold_slot + position_ids % config.page_token_size, swa_token_indices, ) @@ -76,8 +76,8 @@ def _prepare_swa_prefill_workspace_kernel( @triton.jit def _copy_swa_prefix_to_workspace_kernel( - req_to_swa, - swa_buffer, + runtime_swa_buffer, + workspace_swa_buffer, swa_prefill_workspace_offsets, b_req_idx, b_ready_cache_len, @@ -85,7 +85,7 @@ def _copy_swa_prefix_to_workspace_kernel( SLIDING_WINDOW: tl.constexpr, SWA_PAGE_TOKEN_NUM: tl.constexpr, HOLD_REQUEST_ID: tl.constexpr, - REQ_LAYER_STRIDE: tl.constexpr, + RUNTIME_LAYER_STRIDE: tl.constexpr, WORKSPACE_LAYER_STRIDE: tl.constexpr, PACKED_BYTES: tl.constexpr, BLOCK: tl.constexpr, @@ -101,18 +101,18 @@ def _copy_swa_prefix_to_workspace_kernel( if position >= ready: return source_slot = req_idx * SWA_CAPACITY_TOKENS + position % SWA_CAPACITY_TOKENS - source = req_to_swa + layer * REQ_LAYER_STRIDE + source_slot // SWA_PAGE_TOKEN_NUM * PACKED_BYTES - target = swa_buffer + layer * WORKSPACE_LAYER_STRIDE + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES + source = runtime_swa_buffer + layer * RUNTIME_LAYER_STRIDE + source_slot // SWA_PAGE_TOKEN_NUM * PACKED_BYTES + target = workspace_swa_buffer + layer * WORKSPACE_LAYER_STRIDE + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES lanes = tl.arange(0, BLOCK) for chunk in range(0, PACKED_BYTES, BLOCK): offset = chunk + lanes tl.store(target + offset, tl.load(source + offset, offset < PACKED_BYTES, 0), offset < PACKED_BYTES) -def prepare_swa_prefill_workspace(req_manager, prefill_state, b_req_idx, b_ready_cache_len, b_seq_len, workspace_base): +def prepare_swa_prefill_workspace(swa_state, prefill_state, b_req_idx, b_ready_cache_len, b_seq_len, workspace_base): if not b_req_idx.numel(): return - config = req_manager.cache_config + config = swa_state.config _prepare_swa_prefill_workspace_kernel[(1,)]( prefill_state.prefill_token_ranges, prefill_state.swa_prefill_workspace_offsets, @@ -120,106 +120,33 @@ def prepare_swa_prefill_workspace(req_manager, prefill_state, b_req_idx, b_ready b_ready_cache_len, b_seq_len, BATCH=b_req_idx.numel(), - SLIDING_WINDOW=config.sliding_window, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, + SLIDING_WINDOW=config.window_size, + SWA_PAGE_TOKEN_NUM=config.page_token_size, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, SWA_WORKSPACE_BASE=workspace_base, BLOCK=triton.next_power_of_2(b_req_idx.numel()), num_warps=4, ) - swa_buffer = prefill_state.swa_buffer + workspace_swa_buffer = prefill_state.swa_buffer _copy_swa_prefix_to_workspace_kernel[ ( b_req_idx.numel(), - triton.cdiv(config.sliding_window - 1, config.swa_page_token_num) + 1, - swa_buffer.shape[0], + triton.cdiv(config.window_size - 1, config.page_token_size) + 1, + workspace_swa_buffer.shape[0], ) ]( - req_manager.req_to_swa, - swa_buffer, - prefill_state.swa_prefill_workspace_offsets, - b_req_idx, - b_ready_cache_len, - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - SLIDING_WINDOW=config.sliding_window, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, - REQ_LAYER_STRIDE=req_manager.req_to_swa.stride(0), - WORKSPACE_LAYER_STRIDE=swa_buffer.stride(0), - PACKED_BYTES=config.get_packed_kv_bytes(config.swa_page_token_num), + runtime_swa_buffer=swa_state.swa_buffer, + workspace_swa_buffer=workspace_swa_buffer, + swa_prefill_workspace_offsets=prefill_state.swa_prefill_workspace_offsets, + b_req_idx=b_req_idx, + b_ready_cache_len=b_ready_cache_len, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + SLIDING_WINDOW=config.window_size, + SWA_PAGE_TOKEN_NUM=config.page_token_size, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + RUNTIME_LAYER_STRIDE=swa_state.swa_buffer.stride(0), + WORKSPACE_LAYER_STRIDE=workspace_swa_buffer.stride(0), + PACKED_BYTES=config.get_packed_kv_bytes(config.page_token_size), BLOCK=4096, num_warps=4, ) - - -@triton.jit -def _copy_swa_state_kernel( - swa_buffer, - swa_state, - swa_token_indices, - position_ids, - TOKEN_COUNT: tl.constexpr, - SWA_LAYER_STRIDE: tl.constexpr, - KV_BYTES: tl.constexpr, - SWA_PAGE_TOKEN_NUM: tl.constexpr, - SWA_PAGE_BYTES: tl.constexpr, - KV_DATA_BYTES: tl.constexpr, - GATHER: tl.constexpr, - BLOCK: tl.constexpr, -): - token, layer = tl.program_id(0), tl.program_id(1) - slot = tl.load(swa_token_indices + token).to(tl.int64) - position = tl.load(position_ids + token) - byte = tl.arange(0, BLOCK) - offset = slot // SWA_PAGE_TOKEN_NUM * SWA_PAGE_BYTES - packed_offset = tl.where( - byte < KV_DATA_BYTES, - slot % SWA_PAGE_TOKEN_NUM * KV_DATA_BYTES + byte, - SWA_PAGE_TOKEN_NUM * KV_DATA_BYTES - + slot % SWA_PAGE_TOKEN_NUM * (KV_BYTES - KV_DATA_BYTES) - + byte - - KV_DATA_BYTES, - ) - packed = swa_buffer + layer * SWA_LAYER_STRIDE + offset + packed_offset - dense = swa_state + (layer * TOKEN_COUNT + token) * KV_BYTES + byte - valid = (byte < KV_BYTES) & (position >= 0) - if GATHER: - tl.store(dense, tl.load(packed, valid, 0), byte < KV_BYTES) - else: - tl.store(packed, tl.load(dense, valid, 0), valid) - - -def _copy_swa_state(req_manager, req_idx, position_ids, swa_state, *, gather): - config = req_manager.cache_config - with torch.cuda.device(req_manager.req_to_swa.device): - swa_token_indices = get_swa_token_indices(req_manager, req_idx, position_ids) - _copy_swa_state_kernel[(position_ids.numel(), config.layer_num)]( - req_manager.req_to_swa, - swa_state, - swa_token_indices, - position_ids, - TOKEN_COUNT=position_ids.numel(), - SWA_LAYER_STRIDE=req_manager.req_to_swa.stride(0), - KV_BYTES=config.kv_bytes, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - SWA_PAGE_BYTES=config.get_packed_kv_bytes(config.swa_page_token_num), - KV_DATA_BYTES=config.kv_data_bytes, - GATHER=gather, - BLOCK=triton.next_power_of_2(config.kv_bytes), - num_warps=4, - ) - - -def gather_swa_state(req_manager, req_idx, position_ids): - """将 SWA packed 存储中的窗口读为 [层, token, kv_bytes]，供大小页保存。""" - swa_state = torch.empty( - (req_manager.cache_config.layer_num, position_ids.numel(), req_manager.cache_config.kv_bytes), - dtype=torch.uint8, - device=req_manager.req_to_swa.device, - ) - _copy_swa_state(req_manager, req_idx, position_ids, swa_state, gather=True) - return swa_state - - -def scatter_swa_state(req_manager, req_idx, position_ids, swa_state): - _copy_swa_state(req_manager, req_idx, position_ids, swa_state, gather=False) diff --git a/lightllm/models/deepseek_v4_dspark/infer_struct.py b/lightllm/models/deepseek_v4_dspark/infer_struct.py index 42d7eaf37e..5c4fe88d26 100644 --- a/lightllm/models/deepseek_v4_dspark/infer_struct.py +++ b/lightllm/models/deepseek_v4_dspark/infer_struct.py @@ -19,7 +19,7 @@ def _init_swa_metadata(self, model): return super()._init_swa_metadata(model) pos = self.position_ids - swa_width = ((self.req_manager.cache_config.sliding_window + model.block_size + 127) // 128) * 128 + swa_width = ((self.swa_state.config.window_size + model.block_size + 127) // 128) * 128 self.dsv4_swa_indices = torch.empty((pos.numel(), swa_width), dtype=torch.int32, device=pos.device) self.dsv4_swa_lengths = torch.empty_like(pos, dtype=torch.int32) self.dsv4_swa_write_slots = torch.empty_like(pos, dtype=torch.int32) @@ -27,12 +27,9 @@ def _init_swa_metadata(self, model): build_dspark_swa_indices( token_req_idx=self.dsv4_sparse_req_idx, position_ids=self.position_ids, - req_manager=self.req_manager, + swa_state=self.swa_state, swa_indices=self.dsv4_swa_indices, swa_lengths=self.dsv4_swa_lengths, swa_write_slots=self.dsv4_swa_write_slots, - sliding_window=model.config["sliding_window"], block_size=model.block_size, - hold_req_id=self.req_manager.HOLD_REQUEST_ID, - hold_swa_slot=self.req_manager.HOLD_SWA_SLOT, ) diff --git a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py index 5b696e5f30..24bc22aab5 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py @@ -14,9 +14,8 @@ class DeepseekV4DSparkTransformerLayerInfer(DeepseekV4TransformerLayerInfer): def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) - final_layer = network_config["n_layer"] + network_config["dspark_layer_num"] - 1 - self.stage_id = layer_num - network_config["n_layer"] - self.is_last_layer = layer_num == final_layer + self.stage_id = layer_num - network_config["draft_layer_start_id"] + self.is_last_layer = self.stage_id == network_config["n_layer"] - 1 def context_forward( self, diff --git a/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py index c0a7efc8fa..64c79b4e40 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py @@ -24,8 +24,8 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg): hc_mult = network_config["hc_mult"] markov_rank = network_config["dspark_markov_rank"] target_hidden_size = hidden * len(network_config["dspark_target_layer_ids"]) - first_layer_idx = network_config["n_layer"] - last_stage = network_config["dspark_layer_num"] - 1 + first_layer_idx = network_config["draft_layer_start_id"] + last_stage = network_config["n_layer"] - 1 last_prefix = f"mtp.{last_stage}" # Replaced with the already-loaded target weights by the model. diff --git a/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py index 1d37b95601..02ba7d055a 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py @@ -8,5 +8,5 @@ class DeepseekV4DSparkTransformerLayerWeight(DeepseekV4TransformerLayerWeight): def _parse_config(self): super()._parse_config() - stage_id = self.layer_num_ - self.network_config_["n_layer"] + stage_id = self.layer_num_ - self.network_config_["draft_layer_start_id"] self.prefix = f"mtp.{stage_id}" diff --git a/lightllm/models/deepseek_v4_dspark/model.py b/lightllm/models/deepseek_v4_dspark/model.py index ee7f40b42b..b2f324e75d 100644 --- a/lightllm/models/deepseek_v4_dspark/model.py +++ b/lightllm/models/deepseek_v4_dspark/model.py @@ -53,21 +53,25 @@ def __init__(self, kvargs: dict): assert not self.mtp_previous_draft_models, "DeepSeek-V4 DSpark uses one parallel draft model" super().__init__(kvargs) - def _init_config(self): - super()._init_config() - draft_layer_num = len(self.config["compress_ratios"]) - self.config["n_layer"] + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # draft 接在 main 的逻辑层之后；实际 checkpoint 权重使用 mtp.{stage_id} 前缀。 + config.setdefault("draft_layer_start_id", config["num_hidden_layers"]) + draft_layer_num = len(config["compress_ratios"]) - config["draft_layer_start_id"] assert draft_layer_num > 0 - assert self.config["compress_ratios"][-draft_layer_num:] == [0] * draft_layer_num - assert len(self.config["dspark_target_layer_ids"]) == draft_layer_num - self.config["dspark_layer_num"] = draft_layer_num + assert config["compress_ratios"][-draft_layer_num:] == [0] * draft_layer_num + assert len(config["dspark_target_layer_ids"]) == draft_layer_num + config.update(num_hidden_layers=draft_layer_num, n_layer=draft_layer_num) # Aliases consumed by the shared block/Markov implementation. - self.config["block_size"] = int(self.config["dspark_block_size"]) - self.config["mask_token_id"] = int(self.config["dspark_noise_token_id"]) - self.config["markov_rank"] = int(self.config["dspark_markov_rank"]) - self.config["markov_head_type"] = "vanilla" - self.config["enable_confidence_head"] = True - self.config["confidence_head_with_markov"] = True + config["block_size"] = int(config["dspark_block_size"]) + config["mask_token_id"] = int(config["dspark_noise_token_id"]) + config["markov_rank"] = int(config["dspark_markov_rank"]) + config["markov_head_type"] = "vanilla" + config["enable_confidence_head"] = True + config["confidence_head_with_markov"] = True + return config def _verify_params(self): super()._verify_params() @@ -82,8 +86,8 @@ def _verify_params(self): def _init_quant(self): super()._init_quant() - expert_quant_type = self.quant_cfg.get_quant_type(self.config["n_layer"] - 1, "fused_moe") - for layer_idx in range(self.config["n_layer"], len(self.config["compress_ratios"])): + expert_quant_type = self.quant_cfg.get_quant_type(self.config["draft_layer_start_id"] - 1, "fused_moe") + for layer_idx in range(self.config["draft_layer_start_id"], len(self.config["compress_ratios"])): self.quant_cfg.quant_cfg[layer_idx]["fused_moe"] = expert_quant_type def _init_weights(self, start_layer_index=None): @@ -94,7 +98,7 @@ def _init_weights(self, start_layer_index=None): ) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ - first_layer = self.config["n_layer"] + first_layer = self.config["draft_layer_start_id"] self.trans_layers_weight = [ self.transformer_weight_class( first_layer + stage_id, @@ -102,7 +106,7 @@ def _init_weights(self, start_layer_index=None): network_config=self.config, quant_cfg=self.quant_cfg, ) - for stage_id in range(self.config["dspark_layer_num"]) + for stage_id in range(self.config["n_layer"]) ] self.context_wkv_weight = ROWMMWeight( in_dim=self.config["hidden_size"], @@ -114,12 +118,6 @@ def _init_weights(self, start_layer_index=None): tp_world_size=1, ) - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_size: int): # A DSpark logical request always occupies one complete physical block. # Keep the generic HOLD padding semantics, but make its requested size @@ -134,16 +132,12 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s def _init_infer_layer(self, start_layer_index=None): self.pre_infer = self.pre_layer_infer_class(network_config=self.config) self.post_infer = self.post_layer_infer_class(network_config=self.config) - first_layer = self.config["n_layer"] + first_layer = self.config["draft_layer_start_id"] self.layers_infer = [ self.transformer_layer_infer_class(first_layer + stage_id, network_config=self.config) - for stage_id in range(self.config["dspark_layer_num"]) + for stage_id in range(self.config["n_layer"]) ] - def _init_some_value(self): - super()._init_some_value() - self.layers_num = self.config["dspark_layer_num"] - def _init_custom(self): self._freqs_cis_sliding = self.main_model._freqs_cis_sliding self._freqs_cis_compress = self.main_model._freqs_cis_compress diff --git a/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py index bd1a5a3929..a8f124b6f7 100644 --- a/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py @@ -20,7 +20,7 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg): hidden = network_config["hidden_size"] vocab = network_config["vocab_size"] hc_mult = network_config["hc_mult"] - layer_idx = network_config["n_layer"] + layer_idx = network_config["draft_weight_layer_id"] prefix = "mtp.0" self.wte_weight_ = EmbeddingWeight( diff --git a/lightllm/models/deepseek_v4_mtp/model.py b/lightllm/models/deepseek_v4_mtp/model.py index 95c3f183fb..36ac375cef 100644 --- a/lightllm/models/deepseek_v4_mtp/model.py +++ b/lightllm/models/deepseek_v4_mtp/model.py @@ -39,6 +39,14 @@ def __init__(self, kvargs: dict): super().__init__(kvargs) return + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # 权重层号用于量化配置及权重构造，与各 draft 实例的执行逻辑层号分别记录。 + config.setdefault("draft_weight_layer_id", config["num_hidden_layers"]) + config.update(num_hidden_layers=1, n_layer=1) + return config + def _init_custom(self): self._freqs_cis_sliding = self.main_model._freqs_cis_sliding self._freqs_cis_compress = self.main_model._freqs_cis_compress @@ -53,16 +61,8 @@ def _init_custom(self): layer.sin_compress_table = self._sin_cached_compress return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): - mtp_layer_index = self.config["n_layer"] + draft_weight_layer_id = self.config["draft_weight_layer_id"] self.pre_post_weight = self.pre_and_post_weight_class( self.data_type, network_config=self.config, quant_cfg=self.quant_cfg ) @@ -70,7 +70,7 @@ def _init_weights(self, start_layer_index=None): self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ self.trans_layers_weight = [ self.transformer_weight_class( - mtp_layer_index, + draft_weight_layer_id, self.data_type, network_config=self.config, quant_cfg=self.quant_cfg, @@ -89,11 +89,6 @@ def _init_infer_layer(self, start_layer_index=None): assert self.layers_infer[0].compress_ratio == 0, "DeepSeek-V4 MTP draft layer must be SWA-only" return - def _init_some_value(self): - super()._init_some_value() - self.layers_num = 1 - return - def _gen_special_model_input(self, token_num: int): return { "mtp_draft_input_hiddens": torch.randn( diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 3b03fc0386..3146ba3b30 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -120,7 +120,6 @@ def create_full_att_states( from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) - assert args.llm_kv_type == "None", "GLM K-pool requires unquantized KV" cfg = Glm5NextResourceConfig.from_model_config(config, args) logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) if not logical_layer_groups: @@ -136,6 +135,7 @@ def create_full_att_states( ), kv_type=args.llm_kv_type, packed=True, + fp8_format="group64_e8m0", page_size=args.page_size, device="meta", ) @@ -178,7 +178,10 @@ def create_att_tail_states( return [state] def _init_att_backend(self): - self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + + backend_class = Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + self.prefill_att_backend = backend_class(model=self) self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 9e96fd74f2..57df11fc88 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -99,6 +99,7 @@ def create_full_att_states( ), kv_type=args.llm_kv_type, packed=True, + fp8_format="group64_e8m0", page_size=args.page_size, device="meta", ) diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py index 27c1d455d6..3a0ef250fc 100644 --- a/lightllm/models/qwen3_5_mtp/model.py +++ b/lightllm/models/qwen3_5_mtp/model.py @@ -33,7 +33,7 @@ def _init_custom(self): def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: """原生 MTP 只包含一个 full-attention 层，与 main 共用 checkpoint。""" config = super().normalize_model_config(model_config) - config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1) + config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1, layer_types=["full_attention"]) return config def _init_some_value(self): diff --git a/lightllm/models/qwen3next/config.py b/lightllm/models/qwen3next/config.py index c40faa9e15..d647096564 100644 --- a/lightllm/models/qwen3next/config.py +++ b/lightllm/models/qwen3next/config.py @@ -42,6 +42,8 @@ class Qwen3NextResourceConfig: full_attention_interval: int # 当前模型总层数，包含线性与 full-attention 两类层，不合并其他模型的层。 all_layer_num: int + # 显式执行层类型，支持 Loop 等非周期排列；为空时沿用周期配置。 + layer_types: tuple[str, ...] = () def get_full_att_layer_indices(self) -> tuple[int, ...]: """返回当前模型内部的 full-attention 层序号，不包含其他模型的层。 @@ -49,6 +51,8 @@ def get_full_att_layer_indices(self) -> tuple[int, ...]: 调用方通过 model_logical_layer_ids[index] 转换为实际逻辑层 ID； 物理存储层映射由 State 管理，draft 使用自身配置独立创建 State。 """ + if self.layer_types: + return tuple(index for index, kind in enumerate(self.layer_types) if kind != "linear_attention") return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) @classmethod @@ -85,4 +89,5 @@ def from_model_config( ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=llm_config["full_attention_interval"], all_layer_num=n_layer, + layer_types=tuple(llm_config.get("layer_types", ())), ) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 899d0ddf8d..36460f6d28 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -57,7 +57,8 @@ def create_bounded_att_states( assert config["num_attention_heads"] % tp_world_size == 0 assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Qwen3NextResourceConfig.from_model_config(config, args) - logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) + full_layer_indices = set(cfg.get_full_att_layer_indices()) + logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if i not in full_layer_indices) if not logical_layer_ids: return [] state = create_gdn_state( diff --git a/lightllm/models/sensenova_loop/model.py b/lightllm/models/sensenova_loop/model.py index ce6d36a7d5..ea882ba450 100644 --- a/lightllm/models/sensenova_loop/model.py +++ b/lightllm/models/sensenova_loop/model.py @@ -1,7 +1,6 @@ """SenseNova Loop: shared checkpoint weights and independent execution caches.""" from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel -from lightllm.utils.config_utils import get_config_json from .config import get_loop_layer_indices, get_loop_layer_types from .layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer from .layer_weights.transformer_layer_weight import SenseNovaLoopTransformerLayerWeight @@ -13,20 +12,26 @@ class SenseNovaLoopTpPartModel(Qwen3_5TpPartModel): transformer_weight_class = SenseNovaLoopTransformerLayerWeight transformer_layer_infer_class = SenseNovaLoopTransformerLayerInfer + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + if "loop_physical_indices" not in config: + settings = model_config["loop_settings"] + indices = get_loop_layer_indices(settings, config["num_hidden_layers"]) + config.update( + layer_types=get_loop_layer_types(settings, config), + n_layer=len(indices), + num_hidden_layers=len(indices), + loop_physical_indices=indices, + ) + if model_config.get("mhc") is not None: + config["mhc"] = model_config["mhc"] + return config + def _init_config(self): super()._init_config() - model_config = get_config_json(self.weight_dir_) - settings = model_config["loop_settings"] - self.num_physical_layers = self.config["n_layer"] - self.loop_layer_indices = get_loop_layer_indices(settings, self.num_physical_layers) - self.config.update( - layer_types=get_loop_layer_types(settings, self.config), - n_layer=len(self.loop_layer_indices), - num_hidden_layers=len(self.loop_layer_indices), - loop_physical_indices=self.loop_layer_indices, - ) - if model_config.get("mhc") is not None: - self.config["mhc"] = model_config["mhc"] + self.loop_layer_indices = self.config["loop_physical_indices"] + self.num_physical_layers = max(self.loop_layer_indices) + 1 def _init_weights(self, start_layer_index=0): # Allocate and preprocess each checkpoint layer once, then alias weights. diff --git a/unit_tests/common/att_resources/conftest.py b/unit_tests/common/att_resources/conftest.py index aac6d6de02..32bb70d93e 100644 --- a/unit_tests/common/att_resources/conftest.py +++ b/unit_tests/common/att_resources/conftest.py @@ -17,6 +17,16 @@ def allocator_environment(kv_allocator_environment: None, monkeypatch: pytest.Mo monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) get_env_start_args.cache_clear() yield + # 方法缓存以 State 实例为 key；测试会反复创建实例和替换参数，结束后释放缓存引用。 + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + + state_classes = [FullAttentionState, BoundedAttentionState] + for state_class in state_classes: + state_classes.extend(state_class.__subclasses__()) + getter = state_class.get_cpu_transfer_regions + if hasattr(getter, "cache_clear"): + getter.cache_clear() get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index c33c9d1eba..b6679a9835 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -12,13 +12,14 @@ from typing import TYPE_CHECKING from dataclasses import replace -from functools import partial +from functools import lru_cache, partial from math import prod import pytest import torch from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.utils import deepcopy_result from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig @@ -694,6 +695,8 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( def test_cpu_bounded_capacity_includes_state_field_padding(cpu_args: EasyDict): class PaddedGDNState(GDNState): + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self): return tuple( replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) @@ -922,6 +925,8 @@ def test_manager_uses_state_regions_and_adds_each_resource_offset(transfer: str, if not torch.cuda.is_available(): pytest.skip("CUDA required") class ReversedCpuRegionsState(GQAFullAttentionState): + @deepcopy_result + @lru_cache(maxsize=None) def get_cpu_transfer_regions(self): regions = super().get_cpu_transfer_regions() # 本实现的 CPU 页先保存 V 再保存 K，PD 仍沿用 K/V 顺序。 diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py new file mode 100644 index 0000000000..f620b230ef --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -0,0 +1,280 @@ +"""通过模型工厂和生产 fused kernel 验证 V4 State 接入。""" + +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.fixture +def cache(monkeypatch): + for key in ("LIGHTLLM_DP_WORLD_SIZE", "LIGHTLLM_CURRENT_RANK_IN_DP", "LIGHTLLM_CURRENT_DEVICE_ID"): + monkeypatch.setenv(key, "1" if key.endswith("WORLD_SIZE") else "0") + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", + json.dumps( + dict( + page_size=256, + pd_kv_page_size=256, + run_mode="normal", + enable_dp_prompt_cache_fetch=False, + linear_att_hash_page_size=256, + linear_att_page_block_num=8, + cpu_cache_token_page_size=2048, + ) + ), + ) + get_env_start_args.cache_clear() + config = dict( + num_hidden_layers=3, + compress_ratios=[4, 128, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + num_attention_heads=64, + ) + args = SimpleNamespace(llm_kv_type="None", page_size=256, mtp_step=3) + manager = AttentionStateManager(request_capacity=3) + for state in DeepseekV4TpPartModel.create_full_att_states( + model_config=config, + args=args, + model_logical_layer_ids=(0, 1, 2), + ): + state.device = torch.device("cuda") + manager.register_state("main", state) + bounded = DeepseekV4TpPartModel.create_bounded_att_states( + model_config=config, + args=args, + model_logical_layer_ids=(0, 1, 2), + request_capacity=3, + device="cuda", + ) + for state in bounded: + manager.register_state("main", state) + window = next(state for state in bounded if isinstance(state, FP8E8M0SlidingWindowState)) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=8192, page_size=256, big_page_tokens=2048, small_page_capacity=2) + + class TestModel(SimpleNamespace): + __hash__ = object.__hash__ + + model = TestModel( + config=config, + tp_rank_=0, + tp_world_size_=1, + max_req_num=2, + batch_max_tokens=2048, + layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], + att_state_role="main", + args=SimpleNamespace(run_mode="normal", enable_prefill_microbatch_overlap=True), + ) + from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + + backend = DeepseekV4FlashMlaFp8SparseAttBackend(model=model) + yield manager, window, backend + torch.cuda.synchronize() + get_env_start_args.cache_clear() + + +def make_infer(manager, window): + infer = DeepseekV4InferStateInfo() + infer.att_state_manager = manager + infer.att_state_role = "main" + infer.first_logical_layer_id = window.config.logical_layer_ids[0] + return infer + + +def read_packed_rows(buffer, slots, page_tokens): + """独立参考读取 packed 数据/scale 分区，不复用 State 搬运算子。""" + page_ids, rows = slots // page_tokens, slots % page_tokens + data = rows[:, None] * 576 + torch.arange(576, device=slots.device) + scale = page_tokens * 576 + rows[:, None] * 8 + torch.arange(8, device=slots.device) + return torch.cat((buffer[page_ids[:, None], data], buffer[page_ids[:, None], scale]), dim=-1) + + +def test_backend_reads_owned_state_views(cache): + manager, window, backend = cache + infer = make_infer(manager, window) + assert infer.swa_state is window + infer.is_prefill = False + for layer in range(3): + swa, compressed = backend.get_att_input_params(infer, layer) + assert swa.shape[1] == 74880 + assert (compressed is None) == (layer == 2) + if layer != 2: + assert compressed.data_ptr() == infer.get_compressed_kv_view(layer).kv.data_ptr() + + +@pytest.mark.parametrize("ratio,layer,is_indexer", [(4, 0, False), (4, 0, True), (128, 1, False)]) +def test_real_compressor_continues_after_checkpoint_restore(cache, ratio, layer, is_indexer): + from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress + + manager, window, _ = cache + length, boundary = 1024, 512 + head_dim = 128 if is_indexer else 512 + scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), device="cuda") + cos = torch.ones((length, 32), device="cuda") + sin, weight = torch.zeros_like(cos), torch.ones(head_dim, device="cuda") + slots = [torch.arange(offset, offset + length, dtype=torch.int32, device="cuda") for offset in (0, 2048)] + + def compress(req_idx, start, end, decode=False): + infer = make_infer(manager, window) + positions = torch.arange(start, end, dtype=torch.int32, device="cuda") + width = end - start + infer.mem_index, infer.position_ids, infer.is_prefill = slots[req_idx][start:end], positions, not decode + infer._dsv4_token_to_batch_idx = None if decode else torch.zeros(width, dtype=torch.int32, device="cuda") + infer.b_req_idx = torch.full((width if decode else 1,), req_idx, dtype=torch.int32, device="cuda") + infer.b_mtp_index = ( + torch.arange(width, dtype=torch.int32, device="cuda") + if decode + else torch.zeros(1, dtype=torch.int32, device="cuda") + ) + infer.b_seq_len = positions + 1 if decode else torch.tensor([end], dtype=torch.int32, device="cuda") + infer.b_ready_cache_len = None if decode else torch.tensor([start], dtype=torch.int32, device="cuda") + infer.b_q_start_loc = None if decode else torch.zeros(1, dtype=torch.int32, device="cuda") + return fused_compress( + kv_score=scores[start:end], + infer_state=infer, + layer_idx=layer, + norm_weight=weight, + eps=1e-6, + head_dim=head_dim, + qk_rope_head_dim=64, + compress_ratio=ratio, + cos_table=cos, + sin_table=sin, + is_in_indexer=is_indexer, + ) + + reference = compress(0, 0, length) + reference_kv = None if is_indexer else make_infer(manager, window).get_compressed_kv_view(layer).kv.clone() + manager.reset_bounded_att_state([0, 1]) + compress(0, 0, boundary - 2) + compress(0, boundary - 2, boundary) + manager.save_bounded_att_runtime_to_pages([0], [boundary], "small", [0]) + manager.restore_bounded_att_runtime_from_pages([1], [boundary], "small", [0]) + for start, end in ((boundary, 768), (768, length - 4), (length - 4, length)): + output = compress(1, start, end, decode=end == length) + if is_indexer: + closed = (torch.arange(start, end, device="cuda") + 1) % ratio == 0 + torch.testing.assert_close(output[closed], reference[start:end][closed], rtol=0, atol=0) + if not is_indexer: + actual = make_infer(manager, window).get_compressed_kv_view(layer).kv + source = slots[0][boundary + ratio - 1 : length : ratio].long() // ratio + target = slots[1][boundary + ratio - 1 : length : ratio].long() // ratio + torch.testing.assert_close( + read_packed_rows(reference_kv, source, 256 // ratio), + read_packed_rows(actual, target, 256 // ratio), + rtol=0, + atol=0, + ) + + +def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + + manager, window, backend = cache + reference = torch.zeros((17, 74880), dtype=torch.uint8, device="cuda") + kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") + weight, angles = torch.randn(512, dtype=torch.bfloat16, device="cuda"), torch.randn((2048, 32), device="cuda") + freqs = torch.complex(angles.cos(), angles.sin()) + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs + for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): + infer = make_infer(manager, window) + positions = torch.arange(start, end, device="cuda", dtype=torch.int32) + infer.is_prefill, infer.position_ids, infer.microbatch_index = True, positions, 0 + infer.b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") + infer.b_ready_cache_len, infer.b_seq_len = ( + torch.tensor([value], dtype=torch.int32, device="cuda") for value in (start, end) + ) + prefill = backend.create_att_prefill_state(infer) + prefill.prepare_swa_workspace(infer) + infer.prefill_att_state = prefill + indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") + lengths, writes, runtime = (torch.empty_like(positions) for _ in range(3)) + build_swa_indices( + token_req_idx=torch.zeros_like(positions), + position_ids=positions, + swa_state=window, + swa_indices=indices, + swa_lengths=lengths, + swa_write_slots=writes, + prefill_state=prefill, + swa_runtime_write_slots=runtime, + token_to_batch_idx=torch.zeros_like(positions), + b_seq_len=infer.b_seq_len, + ) + infer.dsv4_swa_write_slots, infer.dsv4_swa_runtime_write_slots = writes, runtime + layer._norm_rope_and_cache_kv(kv[start:end], infer, SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight))) + fused_kv_norm_rope_to_swa(kv[start:end], weight, 1e-6, freqs, positions, positions, reference, 128) + visible = torch.arange(max(0, start - 127), end, device="cuda") + offset = prefill.swa_prefill_workspace_offsets[0] + visible - max(0, start - 127) // 128 * 128 + actual = read_packed_rows(prefill.swa_buffer[2].view(-1, 74880), offset.long(), 128) + torch.testing.assert_close(actual[:, :-1], read_packed_rows(reference, visible, 128)[:, :-1], rtol=0, atol=0) + tail = torch.arange(max(0, end - window.config.window_size), end, device="cuda") + actual = read_packed_rows(window.swa_buffer[2, 0].view(-1, 74880), tail % window.swa_capacity_tokens, 128) + torch.testing.assert_close(actual[:, :-1], read_packed_rows(reference, tail, 128)[:, :-1], rtol=0, atol=0) + + +@pytest.mark.parametrize("draft_kind", ["mtp", "dspark"]) +def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): + from contextlib import nullcontext + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel + from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + + _, _, backend = cache + config = dict( + backend.model.config, + compress_ratios=[4, 128, 0, 0, 0], + dspark_target_layer_ids=[0, 1], + dspark_block_size=4, + dspark_noise_token_id=1, + dspark_markov_rank=4, + ) + args = SimpleNamespace( + llm_kv_type="None", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", run_mode="decode" + ) + main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + main.config, main.args, main.max_req_num, main.tp_rank_, main.tp_world_size_ = config, args, 2, 0, 1 + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] + main.torch_memory_saver = SimpleNamespace(region=lambda **kwargs: nullcontext()) + main.prepare_attention_states() + cls = DeepseekV4MTPModel if draft_kind == "mtp" else DeepseekV4DSparkModel + drafts = [] + for index in range(2 if draft_kind == "mtp" else 1): + draft = cls.__new__(cls) + draft.main_model, draft.mtp_previous_draft_models = main, drafts.copy() + draft.config = cls.normalize_model_config(config) + assert draft.config == cls.normalize_model_config(draft.config) + draft.args, draft.max_req_num, draft.tp_rank_, draft.tp_world_size_ = args, 2, 0, 1 + draft.torch_memory_saver = main.torch_memory_saver + first_layer = 3 + sum(len(model.layers_infer) for model in drafts) + draft.layers_infer = [ + SimpleNamespace(layer_num_=first_layer + i) for i in range(draft.config["num_hidden_layers"]) + ] + draft.prepare_attention_states() + drafts.append(draft) + manager = main.att_state_manager + assert all(model.att_state_manager is manager for model in drafts) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=8192, page_size=256, big_page_tokens=2048, small_page_capacity=2) + assert len(manager.get_full_attention_states()) == 2 + assert len(manager.get_bounded_attention_states()) == len(drafts) + 4 + for model in (main, *drafts): + states = manager.get_bounded_attention_states(model.att_state_role) + state = next(item for item in states if isinstance(item, FP8E8M0SlidingWindowState)) + assert state.config.logical_layer_ids == tuple(layer.layer_num_ for layer in model.layers_infer) + assert state.swa_buffer.is_cuda + if model is not main: + assert len(states) == 1 diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py new file mode 100644 index 0000000000..87e0133801 --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -0,0 +1,509 @@ +"""V4 State 独立验证：packed 字段、逻辑快照及不同 MTP ring 的 PD 恢复。""" + +from dataclasses import replace + +import pytest +import torch + +from lightllm.models.deepseek_v4.config import DSV4AttentionConfig +from lightllm.common.attention_state.full_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, +) +from lightllm.common.attention_state.bounded_att import ( + BoundedAttentionState, + FP8E8M0SlidingWindowConfig, + FP8E8M0SlidingWindowState, + WindowCompressorStateConfig, + WindowCompressorState, +) +from lightllm.common.attention_state.types import StateTransferLayout + + +def make_config() -> DSV4AttentionConfig: + return DSV4AttentionConfig( + logical_layer_ids=(2, 5, 8), + compress_ratios=(0, 4, 128), + head_dim=512, + rope_head_dim=64, + indexer_head_dim=128, + sliding_window=128, + ) + + +def make_full(ratio: int, device: str = "meta") -> CompressedAttentionState: + config = make_config() + layout = CompressedAttentionConfig( + logical_layer_groups=((5 if ratio == 4 else 8,),), + compress_ratio=ratio, + kv_dim=config.head_dim, + rope_dim=config.rope_head_dim, + block_token_size=config.compressed_kv_token_alignment, + indexer_dim=config.indexer_head_dim if ratio == 4 else 0, + ) + return CompressedAttentionState(layout_config=layout, page_size=256, device=device) + + +def pd_layout(state: BoundedAttentionState | CompressedAttentionState) -> StateTransferLayout: + regions = state.get_pd_transfer_regions() + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None if isinstance(state, BoundedAttentionState) else 256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + + +def gather_swa_reference(state: FP8E8M0SlidingWindowState, req_idx: int, length: int) -> torch.Tensor: + """按格式的独立参考计算读回窗口，覆盖数据/scale 分区和环绕。""" + config = state.config + packed = state.swa_buffer[:, req_idx].cpu().reshape(len(config.logical_layer_ids), -1) + result = torch.zeros((len(config.logical_layer_ids), config.window_size, config.kv_bytes), dtype=torch.uint8) + for window, position in enumerate(range(length - config.window_size, length)): + if position < 0: + continue + slot = position % state.swa_capacity_tokens + page = slot // 128 * config.get_packed_kv_bytes(128) + token = slot % 128 + data = page + token * config.kv_data_bytes + scale = page + 128 * config.kv_data_bytes + token * (config.kv_bytes - config.kv_data_bytes) + result[:, window, : config.kv_data_bytes] = packed[:, data : data + config.kv_data_bytes] + result[:, window, config.kv_data_bytes :] = packed[:, scale : scale + config.kv_bytes - config.kv_data_bytes] + return result + + +def test_compressed_meta_geometry() -> None: + config = make_config() + swa_config = config.create_swa_config() + assert swa_config.kv_bytes == 584 and swa_config.get_packed_kv_bytes(128) // 128 == 585 + for ratio, slot_bytes in ((4, 180), (128, 9)): + state = make_full(ratio=ratio) + assert state.config.token_slot_bytes == slot_bytes + assert state.storage is None + state.allocate_storage(token_capacity=512) + assert state.storage.shape == (1, 768, 1, slot_bytes) + view = state.get_layer_view(logical_layer_id=5 if ratio == 4 else 8) + assert view.kv.shape[0] == 3 + assert (view.indexer_k is not None) == (ratio == 4) + attention_view = state.config.get_attention_view(packed_kv=view.kv) + assert attention_view.shape == (3, 256 // ratio, 1, 584) + assert attention_view.stride(0) == view.kv.stride(0) + with pytest.raises(NotImplementedError, match="fused compressor"): + state.write_full_att_kv( + logical_layer_id=5 if ratio == 4 else 8, + token_indices=torch.empty(0, dtype=torch.int64, device="meta"), + kv=torch.empty(0, device="meta"), + ) + + +def test_compressed_factory_uses_independent_config() -> None: + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_compressed_att_state + + definition = make_full(ratio=4) + assert not isinstance(definition.config, FullAttentionConfig) + state = create_compressed_att_state(config=definition.config, page_size=256, device="meta") + assert isinstance(state, CompressedAttentionState) + assert state.storage is None and state.get_token_bytes() == 180 + assert state.get_state_type() == "compressed_c4_fp8_group64_e8m0" + + +@pytest.mark.parametrize("page_tokens", [0, 128, 384]) +def test_compressed_transfer_pages_require_complete_blocks(page_tokens: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + monkeypatch.setattr(args, "cpu_cache_token_page_size", page_tokens) + monkeypatch.setattr(args, "pd_kv_page_size", page_tokens) + state = make_full(ratio=4) + with pytest.raises(AssertionError, match="page token capacity"): + state.get_cpu_transfer_regions() + with pytest.raises(AssertionError, match="page token capacity"): + state.get_pd_transfer_regions() + + +@pytest.mark.parametrize("write", [True, False]) +@pytest.mark.parametrize("invalid", ["page_tokens", "page_bytes", "byte_start", "page_index"]) +def test_compressed_cpu_page_boundaries( + write: bool, + invalid: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """错误的布局或物理页在提交 GPU kernel 前失败，避免跨页写入其他缓存数据。""" + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 256) + state = make_full(ratio=4) + regions = state.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + pages = torch.empty((2, layout.byte_end), dtype=torch.uint8) + page_indices = [0] + match invalid: + case "page_tokens": + layout = replace(layout, tokens_per_page=128) + case "page_bytes": + pages = pages[:, :-1] + case "byte_start": + layout = replace(layout, byte_start=-1) + case "page_index": + page_indices = [2] + # 只需索引数量，不创建 CUDA Tensor；所有非法输入必须在读取实际存储前被拦截。 + indices = torch.empty(256, dtype=torch.int32, device="meta") + with pytest.raises(AssertionError, match="CPU page|CPU pages"): + if write: + state.write_to_cpu( + token_indices=indices, + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + layout=layout, + ) + else: + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + token_indices=indices, + layout=layout, + ) + + +@pytest.mark.parametrize("ratio,block_tokens,indexer_dim,slot_bytes", [(4, 256, 0, 153), (8, 512, 32, 81)]) +def test_atomic_compressed_config_controls_layout( + ratio: int, + block_tokens: int, + indexer_dim: int, + slot_bytes: int, +) -> None: + """无需模型配置即可创建 State；indexer 是否存在独立于压缩比例。""" + from lightllm.common.attention_state.full_att import create_compressed_att_state + + config = CompressedAttentionConfig( + logical_layer_groups=((0, 4), (8,)), + compress_ratio=ratio, + kv_dim=512, + rope_dim=64, + block_token_size=block_tokens, + indexer_dim=indexer_dim, + tp_rank=1, + tp_world_size=2, + ) + state = create_compressed_att_state(config=config, page_size=block_tokens, device="meta") + assert not hasattr(config, "cache_config") + assert config.token_slot_bytes == slot_bytes + assert state.get_token_bytes() == 2 * slot_bytes + state.allocate_storage(token_capacity=2 * block_tokens) + assert state.storage.shape == (2, 3 * block_tokens, 1, slot_bytes) + assert state.get_storage_layer_index(0) == state.get_storage_layer_index(4) + view = state.get_layer_view(logical_layer_id=4) + assert view.kv.shape == (3, config.kv_block_bytes) + if indexer_dim: + assert view.indexer_k.shape == (3, block_tokens // ratio * (indexer_dim + 4)) + else: + assert view.indexer_k is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +def test_pd_compressed_kv_preserves_partial_block(ratio: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 256) + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=512) + source.storage.random_(0, 256) + target.storage.zero_() + layout = pd_layout(source) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, device="cuda") + indices = torch.arange(17, dtype=torch.int32, device="cuda") + source.write_to_pd_page(token_indices=indices, pd_page_buffers=pages, pd_page_index=0, layout=layout) + target.read_from_pd_page(pd_page_buffers=pages, token_indices=indices, pd_page_index=0, layout=layout) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, :256], target.storage[:, :256]) + assert not target.storage[:, 256:].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +def test_compressed_cpu_and_token_copy(ratio: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 256) + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=768) + source.storage.random_(0, 256) + regions = source.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + cpu_pages = torch.zeros((3, layout.byte_end), dtype=torch.uint8, pin_memory=True) + src_indices = torch.arange(256, 512, dtype=torch.int32, device="cuda") + dst_indices = torch.arange(256, dtype=torch.int64, device="cuda") + source.write_to_cpu( + token_indices=src_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[2], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[2], + token_indices=dst_indices, + layout=layout, + ) + target.copy_tokens_from_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices + 512, + remote_dp_full_att_state=source, + ) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, 256:512], target.storage[:, :256]) + assert torch.equal(source.storage[:, 256:512], target.storage[:, 512:768]) + assert not target.storage[:, 256:512].any() + assert not cpu_pages[:2].any() + target.copy_full_att_tokens(src_indices=dst_indices, dst_indices=dst_indices + 256) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, 256:512], target.storage[:, 256:512]) + assert not target.storage[:, 768:].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("token_count", [128, 257]) +@pytest.mark.parametrize("from_dp", [False, True]) +def test_compressed_token_copy_rejects_partial_blocks(ratio: int, token_count: int, from_dp: bool) -> None: + """即使长度已按压缩比例对齐，也不能拆开含 KV/scale/indexer 的物理块。""" + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=512) + indices = torch.arange(token_count, dtype=torch.int64, device="cuda") + with pytest.raises(AssertionError, match="complete compressed blocks"): + if from_dp: + target.copy_tokens_from_dp( + src_token_indices=indices, + dst_token_indices=indices, + remote_dp_full_att_state=source, + ) + else: + target.copy_full_att_tokens(src_indices=indices, dst_indices=indices) + + +def make_atomic_state(kind: str, mtp_step: int = 0, device: str = "meta") -> BoundedAttentionState: + if kind == "swa": + return FP8E8M0SlidingWindowState( + config=FP8E8M0SlidingWindowConfig(logical_layer_ids=(2, 5), window_size=128, kv_dim=512, rope_dim=64), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + window, compress_stride, candidate_dim, target = { + "kv_w8": (8, 4, 1024, "kv"), + "indexer_w8": (8, 4, 256, "indexer"), + "kv_w128": (128, 128, 512, "kv"), + "generic_w12": (12, 4, 24, "kv"), + }[kind] + return WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(2, 5), + window_size=window, + compress_stride=compress_stride, + candidate_dim=candidate_dim, + target=target, + ), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + + +def runtime_buffer(state: BoundedAttentionState) -> torch.Tensor: + return state.swa_buffer if isinstance(state, FP8E8M0SlidingWindowState) else state.buffer + + +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +def test_atomic_window_definitions(kind: str) -> None: + source, target = make_atomic_state(kind), make_atomic_state(kind, mtp_step=5) + assert source.get_pd_transfer_regions() == target.get_pd_transfer_regions() + cpu_regions = source.get_cpu_transfer_regions() + copied_regions = source.get_cpu_transfer_regions() + assert copied_regions == cpu_regions + if cpu_regions: + assert copied_regions is not cpu_regions and copied_regions[0] is not cpu_regions[0] + assert source.get_layer_view(2).shape[0] == 3 + assert not source.contains_logical_layer(8) + assert not hasattr(source.config, "compress_ratios") + if isinstance(source, WindowCompressorState): + assert source.config.overlap_tokens == source.config.window_size - source.config.compress_stride + assert (source.get_cpu_transfer_regions() == ()) == (source.config.overlap_tokens == 0) + source.initialize_page_pools(big_page_capacity=3, device="meta", pin_memory=False) + assert (source.big_pages_pool is None) == (source.config.overlap_tokens == 0) + else: + assert source.swa_capacity_tokens == 128 and target.swa_capacity_tokens == 256 + for state in (source, target): + attention_view = state.config.get_attention_view(packed_kv=state.get_layer_view(2)) + assert attention_view.shape == (3 * state.swa_capacity_tokens // 128, 128, 1, 584) + assert attention_view.stride(0) == state.config.get_packed_kv_bytes(128) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +@pytest.mark.parametrize("length", [0, 1, 5, 127, 128, 129, 255, 260]) +@pytest.mark.parametrize("source_step,target_step", [(0, 5), (5, 0)]) +def test_atomic_pd_windows_ignore_mtp_capacity(kind: str, length: int, source_step: int, target_step: int) -> None: + source = make_atomic_state(kind, mtp_step=source_step, device="cuda") + target = make_atomic_state(kind, mtp_step=target_step, device="cuda") + runtime_buffer(source).random_(0, 127) + target.reset_bounded_att_state([1]) + layout = pd_layout(source) + pages = torch.full((2, layout.byte_end + 32), 53, dtype=torch.uint8, device="cuda") + source.write_to_pd_page(req_idx=0, pd_page_buffers=pages, pd_page_index=1, state_token_len=length, layout=layout) + target.read_from_pd_page(pd_page_buffers=pages, req_idx=1, pd_page_index=1, state_token_len=length, layout=layout) + torch.cuda.synchronize() + if isinstance(source, FP8E8M0SlidingWindowState): + assert torch.equal(gather_swa_reference(source, 0, length), gather_swa_reference(target, 1, length)) + else: + src, dst = runtime_buffer(source), runtime_buffer(target) + for position in range(max(0, length - source.config.window_size), length): + assert torch.equal(src[:, 0, position % src.shape[2]], dst[:, 1, position % dst.shape[2]]) + assert (pages[0] == 53).all() and (pages[1, layout.byte_end :] == 53).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +def test_atomic_page_cpu_and_dp_roundtrip(kind: str) -> None: + source = make_atomic_state(kind, mtp_step=3, device="cuda") + target = make_atomic_state(kind, mtp_step=3, device="cuda") + source.initialize_page_pools(big_page_capacity=3) + target.initialize_page_pools(big_page_capacity=3) + runtime_buffer(source).random_(0, 127) + target.reset_bounded_att_state([0, 1]) + source.save_bounded_att_runtime_to_pages( + req_indices=[0], state_token_lens=[512], page_kind="big", page_slot_ids=[0] + ) + target.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=source) + regions = tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + for region in source.get_cpu_transfer_regions() + ) + if not regions: + assert source.big_pages_pool is None + with pytest.raises(AssertionError, match="stride"): + source.restore_bounded_att_runtime_from_pages([0], [513], "big", [0]) + return + layout = StateTransferLayout( + state_name=None, tokens_per_page=256, byte_start=64, byte_end=regions[-1].byte_end, regions=regions + ) + pages = torch.full((3, layout.byte_end + 32), 53, dtype=torch.uint8, pin_memory=True) + source.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0, 0], + cpu_page_buffers=pages, + cpu_page_indices=[1, 2], + layout=layout, + cpu_page_ready_flags=[True, False], + ) + target.read_state_pages_from_cpu( + cpu_page_buffers=pages, cpu_page_indices=[2], big_page_slot_ids=[0], layout=layout + ) + target.copy_big_pages_from_dp(src_page_ids=[0], dst_page_ids=[2], remote_dp_bounded_state=source) + target.restore_bounded_att_runtime_from_pages( + req_indices=[0], state_token_lens=[512], page_kind="big", page_slot_ids=[0] + ) + torch.cuda.synchronize() + assert torch.equal(source.big_pages_pool.buffer[0], target.big_pages_pool.buffer[0]) + assert torch.equal(source.big_pages_pool.buffer[0], target.big_pages_pool.buffer[2]) + assert torch.equal(runtime_buffer(source)[:, 0], runtime_buffer(target)[:, 1]) + assert (pages[:2] == 53).all() and (pages[2, :64] == 53).all() and (pages[2, layout.byte_end :] == 53).all() + if isinstance(source, FP8E8M0SlidingWindowState): + assert torch.equal(gather_swa_reference(source, 0, 512), gather_swa_reference(target, 0, 512)) + else: + for position in range(512 - source.config.overlap_tokens, 512): + ring = source.buffer.shape[2] + assert torch.equal(source.buffer[:, 0, position % ring], target.buffer[:, 0, position % ring]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_manager_joint_cpu_page_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state import AttentionStateManager + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + monkeypatch.setattr(args, "cpu_cache_token_page_size", 256) + monkeypatch.setattr(args, "pd_kv_page_size", 256) + manager = AttentionStateManager(request_capacity=3) + full_states = [make_full(ratio=ratio) for ratio in (4, 128)] + bounded = [ + make_atomic_state(kind, mtp_step=2, device="cuda") for kind in ("swa", "kv_w8", "indexer_w8", "kv_w128") + ] + for state in (*full_states, *bounded): + manager.register_state(model_role="main", state=state) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=512, page_size=256, big_page_tokens=256, device="cuda") + for full in full_states: + full.storage[:, :256].random_(0, 256) + for state in bounded: + runtime_buffer(state).random_(0, 127) + manager.save_bounded_att_runtime_to_pages([0], [256], "big", [0]) + expected = gather_swa_reference(bounded[0], 0, 256) + layouts = manager.get_cpu_page_layouts() + assert bounded[-1].state_name not in layouts + pages = torch.zeros((2, manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + manager.write_req_to_cpu( + token_indices=torch.arange(256, dtype=torch.int64, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + big_page_slot_ids=[0], + ) + manager.read_req_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.arange(256, 512, dtype=torch.int64, device="cuda"), + layouts=layouts, + big_page_slot_ids=[1], + req_idx=1, + state_token_len=256, + ) + torch.cuda.synchronize() + for full in full_states: + assert torch.equal(full.storage[:, :256], full.storage[:, 256:512]) + assert torch.equal(expected, gather_swa_reference(bounded[0], 1, 256)) + assert not pages[0].any() + ipc = manager.create_pd_ipc_copy() + assert all(state.big_pages_pool is None for state in ipc.get_bounded_attention_states()) + assert bounded[0].big_pages_pool is not None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_atomic_dp_ipc_preserves_shared_pools(monkeypatch: pytest.MonkeyPatch) -> None: + import json + from lightllm.common.attention_state import AttentionStateManager + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps({"enable_dp_prompt_cache_fetch": True})) + get_env_start_args.cache_clear() + try: + manager = AttentionStateManager(request_capacity=3) + for kind in ("swa", "kv_w8", "kv_w128"): + state = make_atomic_state(kind, device="cuda") + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + exported = manager.create_dp_ipc_copy() + for state in exported.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + assert pool.buffer.is_shared() + exported.register_page_pools_pin_memory() + finally: + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py new file mode 100644 index 0000000000..8327c36bf1 --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -0,0 +1,244 @@ +"""独立 FP8/E8M0 MLA：显式选择、固定编码、完整 payload 搬运和 padding 保留。""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + FP8E8M0PackedMLAConfig, + FP8E8M0PackedMLAState, + FullAttentionConfig, + FullAttentionState, + PackedMLAFullAttentionState, + create_mla_state, +) +from lightllm.common.attention_state.types import StateTransferLayout + + +if TYPE_CHECKING: + from easydict import EasyDict + + +def make_state(device: str = "meta", capacity: int = 16, tp_rank: int = 0) -> FP8E8M0PackedMLAState: + state = create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=((0, 4), (8,)), + global_head_num=1, + head_dim=512, + dtype=torch.bfloat16, + tp_rank=tp_rank, + tp_world_size=2, + ), + kv_type="fp8kv_dsa", + packed=True, + page_size=4, + device=device, + fp8_format="group64_e8m0", + ) + state.allocate_storage(token_capacity=capacity) + return state + + +def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDict) -> None: + pd_args.cpu_cache_token_page_size = 4 + state = make_state() + assert type(state) is FP8E8M0PackedMLAState + assert FP8E8M0PackedMLAState.__bases__ == (FullAttentionState,) + assert not isinstance(state, PackedMLAFullAttentionState) + assert isinstance(state.config, FP8E8M0PackedMLAConfig) + assert not isinstance(state.config, FullAttentionConfig) + assert not hasattr(state.config, "head_dim") + assert state.get_state_type() == "mla_packed_fp8_group64_e8m0" + assert state.storage.shape == (2, 20, 1, 720) and state.storage.is_meta + assert state.config.payload_bytes == 716 and state.get_token_bytes() == 1440 + assert state.get_storage_layer_index(logical_layer_id=0) == state.get_storage_layer_index(logical_layer_id=4) + assert state.get_layer_view(logical_layer_id=8).kv.shape == (20, 1, 584) + assert state.get_layer_view(logical_layer_id=8).indexer_k.shape == (20, 1, 132) + (region,) = state.get_cpu_transfer_regions() + assert region.shape == (2, 4, 1, 716) and region.byte_end == 2 * 4 * 716 + + # 默认仍是原来的 group128/FP32 编码，不能再通过 512 维自动猜测 GLM 格式。 + with pytest.raises(AssertionError, match="FP8 MLA kernel requires"): + create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=((0,),), global_head_num=1, head_dim=512, dtype=torch.bfloat16 + ), + kv_type="fp8kv_dsa", + packed=True, + device="meta", + ) + with pytest.raises(AssertionError, match="512-dimensional"): + create_mla_state( + config=FullAttentionConfig( + logical_layer_groups=((0,),), global_head_num=1, head_dim=576, dtype=torch.bfloat16 + ), + kv_type="fp8kv_dsa", + packed=True, + fp8_format="group64_e8m0", + device="meta", + ) + + +@pytest.mark.parametrize("draft", [False, True]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_glm_main_and_draft_explicitly_select_encoding(draft: bool, kv_type: str) -> None: + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + + config = dict( + num_hidden_layers=4, + num_attention_heads=4, + hidden_size=2048, + kv_lora_rank=512, + index_kpool=4, + index_head_dim=128, + layer_types=["linear_attention"] * 3 + ["deepseek_sparse_attention"], + linear_attn_config=dict(num_heads=4, head_dim=128, short_conv_kernel_size=4), + ) + args = SimpleNamespace( + tp=2, dp=1, llm_kv_type=kv_type, page_size=4, data_type="bfloat16", linear_att_ssm_data_type="float32" + ) + model_class = Glm5NextMTPModel if draft else Glm5NextTpPartModel + logical_ids = (14,) if draft else (10, 11, 12, 13) + (state,) = model_class.create_full_att_states( + model_config=config, + args=args, + model_logical_layer_ids=logical_ids, + tp_rank=1, + tp_world_size=2, + ) + assert type(state) is (FP8E8M0PackedMLAState if kv_type == "fp8kv_dsa" else PackedMLAFullAttentionState) + assert state.config.logical_layer_groups == ((logical_ids[-1],),) + assert state.storage is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fixed_quantizer_preserves_indexer_and_other_rows() -> None: + from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv + + state = make_state(device="cuda") + state.storage.fill_(173) + indices = torch.tensor([1, 7, 12], dtype=torch.int32, device="cuda") + kv = torch.randn((3, 1, 512), dtype=torch.bfloat16, device="cuda") + expected = torch.empty((3, 1, 584), dtype=torch.uint8, device="cuda") + write_fp8_e8m0_packed_kv( + kv=kv, token_indices=torch.arange(3, dtype=torch.int32, device="cuda"), packed_kv=expected + ) + state.write_full_att_kv(logical_layer_id=4, token_indices=indices, kv=kv) + view = state.get_layer_view(logical_layer_id=0) + torch.testing.assert_close(view.kv[indices], expected, rtol=0, atol=0) + assert view.indexer_k.eq(173).all() + assert state.storage[..., state.config.payload_bytes :].eq(173).all() + assert state.storage[1].eq(173).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("reorder_axes", [False, True]) +def test_full_payload_roundtrip_skips_padding( + transfer: str, + tp_rank: int, + reorder_axes: bool, + pd_args: EasyDict, +) -> None: + pd_args.cpu_cache_token_page_size, pd_args.pd_kv_page_size = 4, 16 + source = make_state(device="cuda") + target = make_state(device="cuda", tp_rank=tp_rank) + source.storage.random_(0, 256) + target.storage.fill_(173) + regions = source.get_cpu_transfer_regions() if transfer == "cpu" else source.get_pd_transfer_regions() + if reorder_axes: + regions = tuple( + replace( + region, + axes=("token", "layer", "head", "dim"), + shape=(region.shape[1], region.shape[0], *region.shape[2:]), + ) + for region in regions + ) + page_tokens = 4 if transfer == "cpu" else 16 + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=page_tokens, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.full( + (5, layout.byte_end + 32), + 99, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + page_view = pages[:, 16:-16] + src_ids = torch.tensor([11, 0, 4, 3, 13, 6, 5, 8, 2], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([15, 2, 3, 7, 12, 9, 1, 8, 0], dtype=torch.int64, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + if transfer == "cpu": + source.write_to_cpu( + token_indices=src_ids, cpu_page_buffers=page_view, cpu_page_indices=[3, 0, 4], layout=layout + ) + target.read_from_cpu( + cpu_page_buffers=page_view, cpu_page_indices=[3, 0, 4], token_indices=dst_ids, layout=layout + ) + else: + source.write_to_pd_page(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout) + target.read_from_pd_page(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, layout=layout) + stream.synchronize() + torch.testing.assert_close( + target.storage[:, dst_ids, :, :716], source.storage[:, src_ids, :, :716], rtol=0, atol=0 + ) + assert target.storage[..., 716:].eq(173).all() + assert target.storage[:, 16:].eq(173).all() + assert pages[:, :16].eq(99).all() and pages[:, -16:].eq(99).all() + assert pages[1].eq(99).all() and pages[2].eq(99).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_cpu_ready_pages_remain_unchanged(pd_args: EasyDict) -> None: + pd_args.cpu_cache_token_page_size = 4 + source = make_state(device="cuda") + source.storage.fill_(7) + regions = source.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, tokens_per_page=4, byte_start=0, byte_end=regions[0].byte_end, regions=regions + ) + pages = torch.full((3, layout.byte_end), 99, dtype=torch.uint8, pin_memory=True) + source.write_to_cpu( + token_indices=torch.arange(9, dtype=torch.int32, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[2, 0, 1], + layout=layout, + cpu_page_ready_flags=[False, True, False], + ) + torch.cuda.synchronize() + assert pages[0].eq(99).all() + assert pages[2].eq(7).all() + partial_page = pages[1].view(2, 4, 1, 716) + assert partial_page[:, 0].eq(7).all() + assert partial_page[:, 1:].eq(99).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_local_and_dp_copy_include_entire_physical_row() -> None: + source = make_state(device="cuda") + target = make_state(device="cuda", capacity=32) + source.storage.random_(0, 256) + target.storage.fill_(173) + src_ids = torch.tensor([0, 3, 11], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([5, 20, 31], dtype=torch.int64, device="cuda") + target.copy_tokens_from_dp(src_token_indices=src_ids, dst_token_indices=dst_ids, remote_dp_full_att_state=source) + torch.testing.assert_close(target.storage[:, dst_ids], source.storage[:, src_ids], rtol=0, atol=0) + source.copy_full_att_tokens( + src_indices=src_ids, dst_indices=torch.tensor([4, 6, 8], dtype=torch.int64, device="cuda") + ) + torch.testing.assert_close(source.storage[:, [4, 6, 8]], source.storage[:, src_ids], rtol=0, atol=0) diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py index 6567ec1211..059f67fdda 100644 --- a/unit_tests/common/att_resources/test_full_att_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -25,7 +25,7 @@ "int8": "group_quantized", "static_head": "static_fp8", "static_tensor": "static_fp8", - "mla": "mla", + "mla": "packed", "packed": "packed_mla", "packed_fp8": "packed_mla", "packed_glm": "packed_mla", diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 83753c87d7..51a56fd05a 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -23,6 +23,14 @@ from easydict import EasyDict +@pytest.fixture(autouse=True) +def manager_cpu_page_args(allocator_environment: None, monkeypatch: pytest.MonkeyPatch) -> None: + """本模块手工构造的 CPU 联合页使用 4-token 规格，启动参数须与布局一致。""" + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 4) + + def composition(device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): return MainDraftTestResourceManager(device, with_tail=with_tail, tail_pool_size=tail_pool_size) diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 2af879e03f..bc03b18597 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -46,7 +46,10 @@ def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_arg assert state.capacity == 8 and read_kv_storage(state).shape[1] == 8 + state.page_size assert not read_kv_storage(state).is_meta assert state.get_token_bytes() == token_bytes - assert state.get_cpu_transfer_regions() == cpu_regions + current_cpu_regions = state.get_cpu_transfer_regions() + assert current_cpu_regions == cpu_regions + assert current_cpu_regions is not cpu_regions + assert current_cpu_regions[0] is not cpu_regions[0] assert state.get_pd_transfer_regions() == pd_regions actual_bytes = sum(buffer.numel() * buffer.element_size() for buffer in full_state_test_buffers(state).values()) assert actual_bytes == token_bytes * (8 + state.page_size) diff --git a/unit_tests/common/att_resources/test_transfer_views.py b/unit_tests/common/att_resources/test_transfer_views.py index 0952c13503..7626abd0da 100644 --- a/unit_tests/common/att_resources/test_transfer_views.py +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -201,3 +201,26 @@ def test_max_copy_dtype_keeps_strided_last_axis() -> None: destination = torch.empty(size=(16,), dtype=torch.float32) copies = [(source, destination)] assert view_max_copy_dtype(copies=copies) is copies +def test_deepcopy_result_protects_cached_values_and_keeps_cache_controls() -> None: + from functools import lru_cache + from lightllm.common.attention_state.utils import deepcopy_result + + calls = 0 + + @deepcopy_result + @lru_cache(maxsize=None) + def describe() -> dict[str, list[int]]: + nonlocal calls + calls += 1 + return {"shape": [2, 128]} + + first = describe() + first["shape"].append(4) + second = describe() + assert second == {"shape": [2, 128]} + assert first is not second and first["shape"] is not second["shape"] + assert calls == 1 and describe.cache_info().hits == 1 + assert describe.cache_parameters() == {"maxsize": None, "typed": False} + describe.cache_clear() + assert describe.cache_info().currsize == 0 + assert describe() == second and calls == 2 diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index 72fb73a9e5..464f07c558 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -63,21 +63,9 @@ def _case(heads, ratio): tokens = 5 model = type("TestModel", (), {})() model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} - from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig - - model.req_manager = SimpleNamespace( - cache_config=DeepseekV4CacheConfig.from_model_config( - dict( - num_hidden_layers=3, - compress_ratios=[0, 4, 128], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - sliding_window=128, - ), - SimpleNamespace(mtp_mode=None), - ) - ) + model.config.update(head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128) + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] + model.tp_rank_ = 0 model.tp_world_size_ = 1 backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend(model=model) q = torch.randn((tokens, backend.padded_q_head_num, 512), dtype=torch.bfloat16, device="cuda") * 0.1 diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py new file mode 100644 index 0000000000..f1a5b9a4bd --- /dev/null +++ b/unit_tests/common/dsv4_test_utils.py @@ -0,0 +1,473 @@ +"""旧 V4 行为用例的测试夹具；数据操作委托新的 State，不向生产代码添加兼容接口。""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch +import triton + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state import AttentionStateManager +from lightllm.models.deepseek_v4.config import DSV4AttentionConfig +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState +from lightllm.common.attention_state.bounded_att.window_pages_pool import WindowPagesPool +from lightllm.common.attention_state.bounded_att.utils import create_page_pool_tensor +from lightllm.common.attention_state.full_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, +) +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.req_manager import ReqManager + +DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT = 256 + + +@dataclass +class DeepseekV4CacheConfig: + """V4 KV 布局和大小页状态缓存配置。""" + + compress_ratios: tuple[int, ...] + head_dim: int + rope_head_dim: int + indexer_head_dim: int + sliding_window: int + + # Packed FP8 KV format shared by Flash and Pro. + quant_group_dim = 64 + scale_alignment_bytes = 8 + # 每个 SWA 物理页的 token 数，用于 KV/scale 打包和 ring 容量对齐。 + swa_page_token_num = 128 + compressed_kv_token_alignment = DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT + short_compress_ratio = 4 + long_compress_ratio = 128 + + def __post_init__(self): + self.compress_ratios = tuple(self.compress_ratios) + self.layer_num = len(self.compress_ratios) + c4_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate == self.short_compress_ratio] + c128_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate == self.long_compress_ratio] + self.layer_to_c4_idx = {layer: slot for slot, layer in enumerate(c4_layers)} + self.layer_to_c128_idx = {layer: slot for slot, layer in enumerate(c128_layers)} + self.n_c4, self.n_c128 = len(self.layer_to_c4_idx), len(self.layer_to_c128_idx) + compressed_kv_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate] + self.layer_to_compressed_kv_idx = {layer: slot for slot, layer in enumerate(compressed_kv_layers)} + self.kv_data_bytes = self.head_dim + self.rope_head_dim + self.kv_scale_bytes = ( + triton.cdiv((self.head_dim - self.rope_head_dim) // self.quant_group_dim, self.scale_alignment_bytes) + * self.scale_alignment_bytes + ) + self.kv_bytes = self.kv_data_bytes + self.kv_scale_bytes + self.indexer_bytes = self.indexer_head_dim + torch.float32.itemsize + self.swa_slot_bytes = self.get_packed_kv_bytes(self.swa_page_token_num) // self.swa_page_token_num + self.compressed_kv_slot_bytes = self._get_compressed_kv_slot_bytes() + + @classmethod + def from_model_config(cls, config: Any, args: Any): + from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num + + layer_num = config["num_hidden_layers"] + compress_ratios = list(config["compress_ratios"]) + if args.mtp_mode is not None: + layer_num += get_added_mtp_kv_layer_num() + # Runtime-added draft layers use sliding-window attention. + compress_ratios.extend([0] * (layer_num - len(compress_ratios))) + return cls( + compress_ratios=tuple(compress_ratios[:layer_num]), + head_dim=config["head_dim"], + rope_head_dim=config["qk_rope_head_dim"], + indexer_head_dim=config["index_head_dim"], + sliding_window=config["sliding_window"], + ) + + def get_packed_kv_bytes(self, token_count: int): + return triton.cdiv(token_count * self.kv_bytes, self.kv_data_bytes) * self.kv_data_bytes + + def get_swa_state_shape(self): + return (self.layer_num, self.sliding_window, self.kv_bytes) + + def get_c4_state_shape(self): + return (self.n_c4, self.short_compress_ratio, 4 * self.head_dim) + + def get_indexer_state_shape(self): + return (self.n_c4, self.short_compress_ratio, 4 * self.indexer_head_dim) + + def get_runtime_state_shapes(self, mtp_step: int): + """每个请求的完整 ring；P/D 使用相同布局，可直接复制槽位。""" + swa_tokens = triton.cdiv(self.sliding_window + 2 * mtp_step, self.swa_page_token_num) * self.swa_page_token_num + c4_tokens = 2 * self.short_compress_ratio + mtp_step + c128_tokens = self.long_compress_ratio + mtp_step + return ( + (self.layer_num, swa_tokens, self.swa_slot_bytes), + (self.n_c4, c4_tokens, 4 * self.head_dim), + (self.n_c4, c4_tokens, 4 * self.indexer_head_dim), + (self.n_c128, c128_tokens, 2 * self.head_dim), + ) + + def get_cpu_cache_full_att_bytes(self): + args = get_env_start_args() + return args.cpu_cache_token_page_size * sum(self.compressed_kv_slot_bytes) + + def get_cpu_cache_big_page_bytes(self): + # MLA KV and compressor state are replicated across TP ranks. + swa_bytes = math.prod(self.get_swa_state_shape()) + c4_bytes = math.prod(self.get_c4_state_shape()) * torch.float32.itemsize + indexer_bytes = math.prod(self.get_indexer_state_shape()) * torch.float32.itemsize + return self.get_cpu_cache_full_att_bytes() + swa_bytes + c4_bytes + indexer_bytes + + def _get_compressed_kv_slot_bytes(self): + block_token_num = self.compressed_kv_token_alignment + # 整块按 KV 数据行对齐，同时能均分到 token 槽位。 + block_bytes_alignment = math.lcm(block_token_num, self.kv_data_bytes) + slot_bytes_list = [] + for ratio in self.compress_ratios: + if ratio == 0: + continue + compressed_token_num = block_token_num // ratio + block_bytes = self.get_packed_kv_bytes(compressed_token_num) + if ratio == self.short_compress_ratio: + block_bytes += compressed_token_num * self.indexer_bytes + block_bytes = triton.cdiv(block_bytes, block_bytes_alignment) * block_bytes_alignment + slot_bytes_list.append(block_bytes // block_token_num) + return tuple(slot_bytes_list) + + +def state_config(config: Any): + return DSV4AttentionConfig( + logical_layer_ids=tuple(range(config.layer_num)), + compress_ratios=config.compress_ratios, + head_dim=config.head_dim, + rope_head_dim=config.rope_head_dim, + indexer_head_dim=config.indexer_head_dim, + sliding_window=config.sliding_window, + ) + + +class DeepseekV4StateCacheManager: + def __init__(self, size: int, cache_config: DeepseekV4CacheConfig, keep_num: int = 0): + # 仅供旧用例保留原联合页视图；生产原子 State 分别拥有自己的页池。 + shapes = ( + cache_config.get_swa_state_shape(), + cache_config.get_c4_state_shape(), + cache_config.get_indexer_state_shape(), + ) + lengths = [ + math.prod(shape) * dtype.itemsize + for shape, dtype in zip(shapes, (torch.uint8, torch.float32, torch.float32)) + ] + self.buffer = create_page_pool_tensor(size=(size, sum(lengths)), dtype=torch.uint8) + offset = 0 + # manager 按规范 State 名排序，联合页依次为 indexer、KV compressor、SWA。 + for name, index, dtype in ( + ("indexer_state", 2, torch.float32), + ("c4_state", 1, torch.float32), + ("swa_state", 0, torch.uint8), + ): + length = lengths[index] + setattr(self, name, self.buffer[:, offset : offset + length].view(dtype).view(size, *shapes[index])) + offset += length + self.size = size + self.allocator = StatePageAllocator(size=size, reserved_slots=keep_num) + self.keep_num = keep_num + + def alloc_one_state_cache(self): + slots = self.allocator.alloc(1) + return slots[0] if slots is not None else None + + def alloc_state_cache(self, count: int): + return self.allocator.alloc(count) + + def free_state_cache(self, slots: list[int]): + self.allocator.free(slots) + + def get_free_cache_num(self): + return self.allocator.free_count + + def get_state_cache(self, slot: int): + return self.swa_state[slot], self.c4_state[slot], self.indexer_state[slot] + + +class DeepseekV4MemoryManager: + def __init__( + self, + size: int, + dtype: torch.dtype, + head_num: int, + head_dim: int, + layer_num: int, + cache_config: DeepseekV4CacheConfig, + ): + self.cache_config = cache_config + self.page_size = get_env_start_args().page_size + self.size = size = size // self.page_size * self.page_size + self.allocator = FullAttKVAllocator(size=size, page_size=self.page_size) + self.manager = AttentionStateManager(request_capacity=3) + self.full_states = {} + self.kv_buffer = [] + config = state_config(cache_config) + for layer, ratio in enumerate(config.compress_ratios): + if not ratio: + continue + state = CompressedAttentionState( + layout_config=CompressedAttentionConfig( + logical_layer_groups=((layer,),), + compress_ratio=ratio, + kv_dim=config.head_dim, + rope_dim=config.rope_head_dim, + block_token_size=config.compressed_kv_token_alignment, + indexer_dim=config.indexer_head_dim if ratio == 4 else 0, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + ), + page_size=self.page_size, + ) + state.allocate_storage(size) + self.manager.register_state("main", state) + self.full_states[layer] = state + self.kv_buffer.append(state.storage[0]) + big_tokens = get_env_start_args().linear_att_hash_page_size * get_env_start_args().linear_att_page_block_num + count = triton.cdiv(size, big_tokens) + self.big_page_buffers = DeepseekV4StateCacheManager(count + 2, cache_config, keep_num=2) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = count + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = count + 1 + self.operator = SimpleNamespace(copy_mem_to_mem=self.copy_mem_to_mem) + + def alloc(self, count: int): + return self.allocator.alloc(count) + + def free(self, indices: torch.Tensor): + self.allocator.free(indices) + + def get_cell_size(self): + return sum(state.get_token_bytes() for state in self.full_states.values()) + + def get_att_input_params(self, layer: int): + index = self.cache_config.layer_to_compressed_kv_idx.get(layer) + return self.kv_buffer[index] if index is not None else None + + def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): + for state in self.full_states.values(): + state.copy_full_att_tokens(src_mem_index, dst_mem_index) + + def alloc_paged_kv_move_buffer(self, count: int, page_size: int): + from copy import copy + + self.manager = copy(self.manager) + for original in self.manager.get_bounded_attention_states(): + state = copy(original) + if isinstance(state, FP8E8M0SlidingWindowState): + if hasattr(self, "req_to_swa"): + state.swa_buffer = self.req_to_swa + else: + field = ( + "req_to_c4_indexer_state" + if state.config.target == "indexer" + else ("req_to_c4_state" if state.config.window_size == 8 else "req_to_c128_state") + ) + if hasattr(self, field): + state.buffer = getattr(self, field) + setattr(self.manager, state.state_name, state) + self.manager._get_states.cache_clear() + get_env_start_args().pd_kv_page_size = page_size + self.manager.get_pd_token_page_layouts.cache_clear() + self.manager.get_pd_runtime_page_layouts.cache_clear() + self.manager.get_pd_page_bytes.cache_clear() + self.kv_move_buffer = torch.empty((count, self.manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + return self.kv_move_buffer + + def get_pd_kv_move_buffer_size(self): + return self.manager.get_pd_page_bytes() * get_env_start_args().pd_kv_page_num + + def get_paged_kv_move_buffer_shape(self, count: int, page_size: int): + return (count, self.manager.get_pd_page_bytes()) + + def write_mem_to_page_kv_move_buffer( + self, + indices: torch.Tensor, + page: int, + dp_index: int, + managers: list[DeepseekV4MemoryManager], + tp_size: int, + kind: str = "kv", + req_idx: int | None = None, + ): + for manager in managers[dp_index * tp_size : dp_index * tp_size + 1]: + manager.manager.write_to_pd_page( + pd_page_buffers=self.kv_move_buffer, + pd_page_index=page, + kind=kind, + layouts=manager.manager.get_pd_token_page_layouts() + if kind == "kv" + else manager.manager.get_pd_runtime_page_layouts(), + token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), + req_idx=req_idx, + state_token_len=getattr(manager, "state_token_len", 0), + ) + + def read_page_kv_move_buffer_to_mem( + self, + indices: torch.Tensor, + page: int, + dp_index: int, + managers: list[DeepseekV4MemoryManager], + tp_size: int, + kind: str = "kv", + req_idx: int | None = None, + ): + for manager in managers[dp_index * tp_size : (dp_index + 1) * tp_size]: + manager.manager.read_from_pd_page( + pd_page_buffers=self.kv_move_buffer, + pd_page_index=page, + kind=kind, + layouts=manager.manager.get_pd_token_page_layouts() + if kind == "kv" + else manager.manager.get_pd_runtime_page_layouts(), + token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), + req_idx=req_idx, + state_token_len=getattr(manager, "state_token_len", 0), + ) + + +class DeepseekV4ReqManager(ReqManager): + def __init__( + self, count: int, length: int, memory: DeepseekV4MemoryManager | None, config: Any, mtp_step: int = 0 + ): + super().__init__(count, length) + self.cache_config = config + self.states = DeepseekV4TpPartModel.create_bounded_att_states( + model_config=dict( + compress_ratios=config.compress_ratios, + head_dim=config.head_dim, + qk_rope_head_dim=config.rope_head_dim, + index_head_dim=config.indexer_head_dim, + sliding_window=config.sliding_window, + ), + args=SimpleNamespace(mtp_step=mtp_step), + model_logical_layer_ids=tuple(range(config.layer_num)), + request_capacity=count + 1, + device="cuda", + ) + self.swa_state = next(state for state in self.states if isinstance(state, FP8E8M0SlidingWindowState)) + self.config = self.swa_state.config + self.swa_buffer = self.req_to_swa = self.swa_state.swa_buffer + shapes = config.get_runtime_state_shapes(mtp_step) + buffers = [] + for window_size, target, shape in ((8, "kv", shapes[1]), (8, "indexer", shapes[2]), (128, "kv", shapes[3])): + state = next( + ( + item + for item in self.states + if isinstance(item, WindowCompressorState) + and item.config.window_size == window_size + and item.config.target == target + ), + None, + ) + buffers.append( + state.buffer if state is not None else torch.zeros((shape[0], count + 1, *shape[1:]), device="cuda") + ) + self.req_to_c4_state, self.req_to_c4_indexer_state, self.req_to_c128_state = buffers + self.window = SimpleNamespace( + swa_buffer=self.swa_buffer, + c4_buffer=buffers[0], + indexer_buffer=buffers[1], + c128_buffer=buffers[2], + request_capacity=count + 1, + swa_capacity_tokens=self.swa_state.swa_capacity_tokens, + HOLD_SWA_SLOT=self.swa_state.HOLD_SWA_SLOT, + ) + self.c4_state_tokens = self.req_to_c4_state.shape[2] + self.c128_state_tokens = self.req_to_c128_state.shape[2] + self.sliding_window = config.sliding_window + self.swa_capacity_tokens, self.HOLD_SWA_SLOT = self.window.swa_capacity_tokens, self.window.HOLD_SWA_SLOT + if memory is not None: + self.bind_mem_manager(memory) + + def bind_mem_manager(self, memory: DeepseekV4MemoryManager | None): + self.memory = memory + self.bind_hold_page(memory.allocator.hold_token_indices) + if not memory.manager._state_registration_finished: + for state in self.states: + memory.manager.register_state("main", state) + memory.manager.finish_state_registration() + memory.manager.full_att_allocator = memory.allocator + memory.manager.bounded_att_big_page_allocator = memory.big_page_buffers.allocator + memory.manager.bounded_att_cpu_load_big_page_slot = memory.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID + self._bind_page_views(memory.big_page_buffers, "big") + + def _bind_page_views(self, pool, kind): + for state in self.states: + if isinstance(state, WindowCompressorState) and not state.config.overlap_tokens: + continue + view = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + atomic_pool = object.__new__(WindowPagesPool) + atomic_pool.buffer = view + setattr(state, f"{kind}_pages_pool", atomic_pool) + + def get_compressor_state(self, logical_layer_id, target="kv"): + return next( + state + for state in self.states + if isinstance(state, WindowCompressorState) + and state.config.target == target + and state.contains_logical_layer(logical_layer_id) + ) + + def get_compressed_kv_state(self, logical_layer_id): + return self.memory.full_states[logical_layer_id] + + def get_compressed_kv_view(self, logical_layer_id: int): + return self.memory.full_states[logical_layer_id].get_layer_view(logical_layer_id) + + def get_att_input_params(self, layer: int): + compressed = self.get_compressed_kv_view(layer).kv if layer in self.memory.full_states else None + return ( + self.req_to_swa[layer].view(-1, self.config.get_packed_kv_bytes(self.config.page_token_size)), + compressed, + ) + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = DeepseekV4StateCacheManager(size, self.cache_config) + self._bind_page_views(self.small_page_buffers, "small") + self.memory.manager.bounded_att_small_page_allocator = self.small_page_buffers.allocator + return self.small_page_buffers + + def save_state(self, req: Any, slot: int, pool: DeepseekV4StateCacheManager): + for state in self.states: + if state.big_pages_pool is not None: + external = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + state._copy_window( + req_idx=req.req_idx, + external=external[slot], + state_token_len=req.get_chuncked_input_token_len(), + write=True, + ) + + def restore_state(self, req: Any, pool: DeepseekV4StateCacheManager, slot: int): + for state in self.states: + if state.big_pages_pool is not None: + external = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + state._copy_window( + req_idx=req.req_idx, external=external[slot], state_token_len=req.cur_kv_len, write=False + ) + + def free(self, reqs: list[int], indices: torch.Tensor = None): + super().free(reqs) + if indices is not None: + self.memory.free(indices) diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index 9b03f2e192..71a348012e 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -4,12 +4,9 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig -from lightllm.common.req_manager import DeepseekV4ReqManager +from unit_tests.common.dsv4_test_utils import DeepseekV4MemoryManager, DeepseekV4CacheConfig, DeepseekV4ReqManager from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices from lightllm.utils.envs_utils import get_env_start_args -from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer, get_indexer_k_buffer from lightllm.models.deepseek_v4.triton_kernel.swa_cache import get_swa_token_indices as _get_swa_token_indices @@ -17,9 +14,9 @@ def _compressed_kv_pools(manager): c4 = next(iter(manager.cache_config.layer_to_c4_idx)) c128 = next(iter(manager.cache_config.layer_to_c128_idx)) return ( - get_compressed_kv_buffer(manager, c4).unsqueeze(0), - get_indexer_k_buffer(manager, c4).unsqueeze(0), - get_compressed_kv_buffer(manager, c128).unsqueeze(0), + manager.full_states[c4].get_layer_view(c4).kv.unsqueeze(0), + manager.full_states[c4].get_layer_view(c4).indexer_k.unsqueeze(0), + manager.full_states[c128].get_layer_view(c128).kv.unsqueeze(0), ) @@ -57,6 +54,7 @@ def cache(monkeypatch, tmp_path, request): "mtp_step": mtp_step, "mtp_dynamic_verify": False, "enable_ep_moe": False, + "enable_dp_prompt_cache_fetch": False, } ), ) @@ -116,68 +114,62 @@ def test_big_page_pool_reserves_cpu_cache_staging_slots(cache): assert pool.get_free_cache_num() == capacity -def test_copy_mem_to_mem_preserves_token_slot_indices(cache): +def test_copy_mem_to_mem_preserves_complete_block_indices(cache): manager, _ = cache - source = torch.tensor([3, 257, 509], device="cuda") - destination = torch.tensor([1025, 1028, 1407], device="cuda") + # 压缩块内的 KV/scale/indexer 不可拆分；允许按非连续、乱序的完整块复制。 + source = torch.cat([torch.arange(start, start + 256, device="cuda") for start in (0, 512, 256)]) + destination = torch.cat([torch.arange(start, start + 256, device="cuda") for start in (1024, 1536, 1280)]) for buffer in manager.kv_buffer: buffer.zero_() - buffer[source] = torch.arange(1, 4, device="cuda", dtype=torch.uint8)[:, None, None] + values = torch.arange(1, 4, device="cuda", dtype=torch.uint8).repeat_interleave(256) + buffer[source] = values[:, None, None] manager.operator.copy_mem_to_mem(source, destination) for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination], buffer[source], rtol=0, atol=0) def test_cache_io_copies_individual_token_slots(cache): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_compressed_kv_pages - manager, _ = cache source = torch.tensor([3, 17, 257, 509], device="cuda") destination = torch.tensor([1025, 1028, 1407, 2047], device="cuda") for buffer in manager.kv_buffer: buffer.zero_() buffer[source] = torch.arange(1, 5, device="cuda", dtype=torch.uint8)[:, None, None] - pages = torch.empty((1, 4 * manager.get_cell_size()), dtype=torch.uint8, pin_memory=True) - page_ids = torch.zeros(1, dtype=torch.int32, device="cuda") - copy_compressed_kv_pages(manager, source, page_ids, pages, to_pages=True, page_size=4) - copy_compressed_kv_pages(manager, destination, page_ids, pages, to_pages=False, page_size=4) + layouts = manager.manager.get_cpu_page_layouts() + pages = torch.empty((1, manager.manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + for state in manager.full_states.values(): + state.write_to_cpu(source, pages, [0], layout=layouts[state.state_name]) + state.read_from_cpu(pages, [0], destination, layout=layouts[state.state_name]) torch.cuda.synchronize() for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination], buffer[source], rtol=0, atol=0) assert buffer[destination - 1].eq(0).all() untouched = torch.tensor([1024, 1026, 1406, 1408], device="cuda") assert buffer[untouched].eq(0).all() - snapshot = manager.get_index_kv_buffer(source) - assert all(isinstance(value, torch.Tensor) for value in snapshot.values()) - # The shared preload path truncates dictionary values along their token axis. - snapshot = {name: value[:, :2].cpu() for name, value in snapshot.items()} for buffer in manager.kv_buffer: buffer[destination] = 0 - manager.load_index_kv_buffer(destination[:2], snapshot) + # 同一 CPU 页只读前两个 token，不能覆盖目标中剩余槽位。 + for state in manager.full_states.values(): + state.read_from_cpu(pages, [0], destination[:2], layout=layouts[state.state_name]) for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination[:2]], buffer[source[:2]], rtol=0, atol=0) assert buffer[destination[2:]].eq(0).all() def test_request_runtime_is_independent_of_compressed_kv_capacity(cache, monkeypatch): - from lightllm.common.kv_cache_mem_manager import MemoryManager - original, requests = cache - observed = [] - - def profile(manager, mem_fraction): - observed.append(requests.req_to_swa.flatten(1, 2).nbytes) - assert _state_buffer(requests, "c4").is_cuda and _state_buffer(requests, "c128").is_cuda - assert not hasattr(manager, "swa_buffer") - manager.size = 8192 - - monkeypatch.setattr(MemoryManager, "profile_size", profile) - manager = DeepseekV4MemoryManager(None, torch.bfloat16, 1, 512, 3, original.cache_config) - assert observed == [requests.req_to_swa.flatten(1, 2).nbytes] - assert manager.size == 8192 - for buffer in manager.kv_buffer: - assert buffer.shape[0] == manager.size + manager.page_size - assert sum(buffer.nbytes for buffer in manager.kv_buffer) == original.get_cell_size() * (8192 + manager.page_size) + runtime = requests.window + before = tuple(tensor.data_ptr() for tensor in ( + runtime.swa_buffer, runtime.c4_buffer, runtime.indexer_buffer, runtime.c128_buffer, + )) + # 扩展压缩 KV 容量不能重新创建或移动按请求数分配的运行态。 + for state in original.full_states.values(): + state.allocate_storage(token_capacity=16384) + assert state.storage.shape[1] == 16384 + state.page_size + assert before == tuple(tensor.data_ptr() for tensor in ( + runtime.swa_buffer, runtime.c4_buffer, runtime.indexer_buffer, runtime.c128_buffer, + )) + assert runtime.request_capacity == requests.max_request_num + 1 @pytest.mark.parametrize("page_size", [256, 1024, 4096]) @@ -241,10 +233,16 @@ def _test_backend(requests, token_capacity=8192): model = type("TestModel", (), {})() model.req_manager = requests + model.att_state_manager = requests.memory.manager + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] + model.tp_rank_ = 0 + model.mtp_previous_draft_models = [] + model.__dict__["att_state_role"] = "main" model.config = { "num_attention_heads": 64, - "compress_ratios": requests.cache_config.compress_rates, + "compress_ratios": requests.cache_config.compress_ratios, "index_topk": 512, + "head_dim": 512, "qk_rope_head_dim": 64, "index_head_dim": 128, "sliding_window": requests.sliding_window, } model.tp_world_size_ = 1 model.max_req_num = requests.max_request_num @@ -270,6 +268,8 @@ def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): state.prepare_swa_workspace( SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, b_req_idx=req_ids, b_ready_cache_len=starts, b_seq_len=ends, @@ -384,7 +384,7 @@ def test_indexer_writer_only_writes_closed_groups(cache): _compressed_kv_pools(manager)[1].zero_() from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k import destindex_copy_indexer_k - destindex_copy_indexer_k(keys, slots[:259], positions, get_indexer_k_buffer(manager, 0), 64) + destindex_copy_indexer_k(keys, slots[:259], positions, manager.full_states[0].get_layer_view(0).indexer_k, 64) closed = torch.arange(3, 259, 4, device="cuda") actual = _read_packed_rows(_compressed_kv_pools(manager)[1], 0, slots[closed] // 4) # Reference quantization is independent of the cache writer and stays in the test. @@ -428,6 +428,11 @@ def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): ] model = model_cls.__new__(model_cls) model.req_manager = requests + model.att_state_manager = requests.memory.manager + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] + model.tp_rank_ = 0 + model.mtp_previous_draft_models = [] + model.__dict__["att_state_role"] = "main" model_input = ModelInput( batch_size=1, total_token_num=257, @@ -640,6 +645,7 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): "req_to_c128_state", ): setattr(memory, name, getattr(runtime, name)) + memory.state_token_len = length memory.alloc_paged_kv_move_buffer(1, manager.page_size) for start in range(0, length, manager.page_size): end = min(length, start + manager.page_size) @@ -654,26 +660,27 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): atol=0, ) - # Copy all ring slots, including the SWA-only layer and partially filled compressor pools. + # PD 传递与 MTP 容量无关的逻辑窗口，不再传递无效候选槽位或 padding。 source_memory.write_mem_to_page_kv_move_buffer([], 0, 0, [source_memory], 1, "att_state", source) manager.kv_move_buffer.copy_(source_memory.kv_move_buffer) manager.read_page_kv_move_buffer_to_mem([], 0, 0, [manager], 1, "att_state", destination) + positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda") + for layer in range(requests.cache_config.layer_num): + torch.testing.assert_close( + _read_swa(requests, layer, destination, positions, prefill=False), + _read_swa(prefill, layer, source, positions, prefill=False), rtol=0, atol=0, + ) ring_size = requests.swa_capacity_tokens - torch.testing.assert_close( - requests.req_to_swa[:, destination], - prefill.req_to_swa[:, source], - rtol=0, - atol=0, - ) - # Neither another request's ring nor the prefill workspace belongs to this transfer. assert requests.req_to_swa.flatten(1, 2)[:, : destination * ring_size].eq(0).all() assert requests.req_to_swa.flatten(1, 2)[:, (destination + 1) * ring_size :].eq(0).all() - for source_buffer, target_buffer in ( - (prefill.req_to_c4_state, requests.req_to_c4_state), - (prefill.req_to_c4_indexer_state, requests.req_to_c4_indexer_state), - (prefill.req_to_c128_state, requests.req_to_c128_state), + for source_buffer, target_buffer, width in ( + (prefill.req_to_c4_state, requests.req_to_c4_state, 8), + (prefill.req_to_c4_indexer_state, requests.req_to_c4_indexer_state, 8), + (prefill.req_to_c128_state, requests.req_to_c128_state, 128), ): - torch.testing.assert_close(source_buffer[:, source], target_buffer[:, destination], rtol=0, atol=0) + positions = torch.arange(max(0, length - width), length, device="cuda") + torch.testing.assert_close(source_buffer[:, source, positions % source_buffer.shape[2]], + target_buffer[:, destination, positions % target_buffer.shape[2]], rtol=0, atol=0) prefill.free([source], src) requests.free([destination], dst) assert manager.allocator.can_use_mem_size == manager.size @@ -715,6 +722,7 @@ def test_pd_transport_completion_covers_peer_gpu_state(cache, length): cudart = ctypes.CDLL(f"libcudart.so.{torch.version.cuda.split('.')[0]}") assert cudart.cudaDeviceEnablePeerAccess(1, 0) in (0, 704) cudart.cudaGetLastError() + transport.state_token_len = peer.state_token_len = length transport.alloc_paged_kv_move_buffer(1, transport.page_size) mem_managers = [transport, peer] for start in range(0, length, transport.page_size): @@ -734,16 +742,15 @@ def test_pd_transport_completion_covers_peer_gpu_state(cache, length): completed = torch.cuda.Event() completed.record(copy_stream) completed.synchronize() - # 完成事件只需保证数据可读，无需目标 stream 已空闲；用独立 stream 校验整块 state。 + # 在独立流验证逻辑窗口；无效 MTP 槽位不属于 PD 数据。 with torch.cuda.stream(consumer_stream): - torch.testing.assert_close( - requests.req_to_swa[:, destination], - requests.req_to_swa[:, source], - rtol=0, - atol=0, - ) - for buffer in (requests.req_to_c4_state, requests.req_to_c4_indexer_state, requests.req_to_c128_state): - torch.testing.assert_close(buffer[:, destination], buffer[:, source], rtol=0, atol=0) + positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda:1") + for layer in range(requests.cache_config.layer_num): + torch.testing.assert_close(_read_swa(requests, layer, destination, positions, prefill=False), + _read_swa(requests, layer, source, positions, prefill=False), rtol=0, atol=0) + for buffer, width in ((requests.req_to_c4_state, 8), (requests.req_to_c4_indexer_state, 8), (requests.req_to_c128_state, 128)): + positions = torch.arange(max(0, length - width), length, device="cuda:1") % buffer.shape[2] + torch.testing.assert_close(buffer[:, destination, positions], buffer[:, source, positions], rtol=0, atol=0) for buffer in peer.kv_buffer: torch.testing.assert_close(buffer[src.long()], buffer[dst.long()], rtol=0, atol=0) requests.free([source, destination], torch.cat([src, dst])) @@ -771,6 +778,8 @@ def compress(req_idx, slots, start, end, decode=False): state = SimpleNamespace( mem_manager=manager, req_manager=requests, + swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, mem_index=slots[start:end], position_ids=positions, is_prefill=not decode, @@ -921,16 +930,13 @@ def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, def build(): if dspark: build_dspark_swa_indices( - req_idx, - query_positions, - requests, - indices, - lengths, - write_slots, - sliding_window=128, + token_req_idx=req_idx, + position_ids=query_positions, + swa_state=requests.swa_state, + swa_indices=indices, + swa_lengths=lengths, + swa_write_slots=write_slots, block_size=token_num, - hold_req_id=requests.HOLD_REQUEST_ID, - hold_swa_slot=requests.HOLD_SWA_SLOT, ) else: build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) @@ -940,7 +946,10 @@ def build(): "Model", (), { - "config": {"num_attention_heads": 64, "compress_ratios": [0]}, + "config": dict(num_attention_heads=64, compress_ratios=[0], head_dim=512, + qk_rope_head_dim=64, index_head_dim=128, sliding_window=128), + "layers_infer": [SimpleNamespace(layer_num_=0)], + "tp_rank_": 0, "tp_world_size_": 1, "req_manager": requests, }, @@ -979,6 +988,8 @@ def forward(): reference_kv = backend.get_att_input_params( SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[source], ), @@ -1040,7 +1051,7 @@ def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash small = requests.create_small_page_cache_manager(2) block_num = 2048 // hash_page_size big_page_tokens = block_num * hash_page_size - radix = HybridAttPagedRadixCache(manager.size, 0, hash_page_size, block_num, manager, small) + radix = HybridAttPagedRadixCache(manager.size, 0, hash_page_size, block_num, manager.manager) args = get_env_start_args() args.linear_att_hash_page_size = hash_page_size args.linear_att_page_block_num = block_num @@ -1048,6 +1059,7 @@ def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash hit_len = tail_len if hit_extra_page else big_page_tokens for name, value in { "req_manager": requests, + "att_state_manager": manager.manager, "radix_cache": radix, "args": args, "is_hybrid_att_model": True, @@ -1197,6 +1209,10 @@ def test_c4_indexer_metadata_scoring_and_graph_replay(cache, monkeypatch, mode): indexer = indexer_module.DeepseekV4IndexInfer(0, config, 1) state = infer_struct_module.DeepseekV4InferStateInfo() state.mem_manager, state.req_manager = manager, requests + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.is_prefill = mode != "decode" state.microbatch_index, state.max_kv_seq_len = 0, 4096 state.b_req_idx = torch.tensor([req_idx] if state.is_prefill else [req_idx] * tokens, device="cuda") @@ -1253,7 +1269,8 @@ def select(): @pytest.mark.parametrize("is_prefill", [False, True]) -def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill): +@pytest.mark.parametrize("owned_layer_ids", [(0, 1, 2), (0, 2), (2,)]) +def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill, owned_layer_ids): from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace @@ -1262,6 +1279,8 @@ def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill second, second_slots = _reserve(manager, requests, 512) backend = _test_backend(requests) model = backend.model + model.layers_infer = [layer for layer in model.layers_infer if layer.layer_num_ in owned_layer_ids] + has_c128_layer = 1 in owned_layer_ids model.graph = SimpleNamespace(can_run=lambda *args: True, graph_max_len_in_batch=16384) model.dsv4_workspace = DeepseekV4Workspace(model) rope = torch.zeros((1024, 32), device="cuda") @@ -1272,6 +1291,10 @@ def make_state(lengths, microbatch_index): state = DeepseekV4InferStateInfo() state.is_prefill, state.microbatch_index = is_prefill, microbatch_index state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.batch_size = 2 state.max_kv_seq_len = max(lengths) state.input_ids = torch.zeros(2, dtype=torch.int64, device="cuda") @@ -1290,7 +1313,11 @@ def make_state(lengths, microbatch_index): swa_before = captured_state.dsv4_swa_indices.clone() other_state.dsv4_swa_indices.fill_(-7) torch.testing.assert_close(captured_state.dsv4_swa_indices, swa_before) - assert captured_state.dsv4_c128_indices.shape == (2, 64 if is_prefill else 128) + if has_c128_layer: + assert captured_state.dsv4_c128_indices.shape == (2, 64 if is_prefill else 128) + else: + assert captured_state.dsv4_c128_indices is None + assert captured_state.dsv4_c128_lengths is None captured_state.init_c4_indices(512) other_state.init_c4_indices(512) captured_state.dsv4_c4_indices.fill_(11) @@ -1302,7 +1329,8 @@ def make_state(lengths, microbatch_index): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): swa_output = captured_state.dsv4_swa_indices.clone() - c128_output = captured_state.dsv4_c128_indices.clone() + if has_c128_layer: + c128_output = captured_state.dsv4_c128_indices.clone() updated = make_state([258, 386], 0) assert updated.dsv4_c4_indices is None if is_prefill: @@ -1311,7 +1339,8 @@ def make_state(lengths, microbatch_index): captured_state.copy_for_cuda_graph(updated) graph.replay() torch.testing.assert_close(swa_output, updated.dsv4_swa_indices) - torch.testing.assert_close(c128_output, updated.dsv4_c128_indices) + if has_c128_layer: + torch.testing.assert_close(c128_output, updated.dsv4_c128_indices) assert captured_state.dsv4_c4_indices.eq(11).all() requests.free([first, second], torch.cat((first_slots, second_slots))) @@ -1333,7 +1362,8 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, monkeypatch.setitem(args, "mtp_mode", "dspark") monkeypatch.setitem(args, "enable_decode_microbatch_overlap", True) monkeypatch.setitem(args, "enable_prefill_microbatch_overlap", False) - config = dict(sliding_window=128, index_topk=512, num_attention_heads=64, compress_ratios=[4, 128, 0]) + config = dict(sliding_window=128, index_topk=512, num_attention_heads=64, compress_ratios=[4, 128, 0], + head_dim=512, qk_rope_head_dim=64, index_head_dim=128) model = type("TestModel", (), {})() vars(model).update( config=config, @@ -1344,7 +1374,10 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, has_vision=False, args=args, req_manager=requests, - tp_world_size_=1, + swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, + tp_rank_=0, tp_world_size_=1, att_state_role="main", + layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], ) backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) model.block_size = 3 @@ -1356,6 +1389,10 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.is_prefill, state.microbatch_index = mode == "prefill", microbatch_index state.is_draft_model = True state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 if state.is_prefill: state.batch_size = 1 state.input_ids = torch.ones(3, dtype=torch.int64, device="cuda") @@ -1420,7 +1457,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4StateCacheManager + from unit_tests.common.dsv4_test_utils import DeepseekV4StateCacheManager manager, requests = cache source, src = _reserve(manager, requests, length) @@ -1448,6 +1485,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( tail = slot else: source_checkpoints[end] = slot + offload_before = manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID].clone() # The live ring no longer contains the prompt's continuation checkpoints. requests.req_to_swa.flatten(1, 2).fill_(99) _state_buffer(requests, "c4").fill_(88) @@ -1491,6 +1529,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( producer = torch.cuda.current_stream() context = common.g_infer_context monkeypatch.setattr(context, "req_manager", requests) + monkeypatch.setattr(context, "att_state_manager", manager.manager) monkeypatch.setattr( context, "radix_cache", @@ -1512,7 +1551,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( disable_linear_att_small_page_cpu_cache=False, ) module.backend = SimpleNamespace( - is_master_in_dp=True, dp_world_size=1, model=SimpleNamespace(mem_manager=manager, req_manager=requests) + is_master_in_dp=True, dp_world_size=1, is_hybrid_att_model=True, model=SimpleNamespace(att_state_manager=manager.manager, req_manager=requests) ) module.cpu_cache_client = client module.cpu_cache_handle_queue = deque() @@ -1541,7 +1580,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( if tail is not None: torch.testing.assert_close( manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID], - expected_checkpoints[-1], + offload_before, rtol=0, atol=0, ) @@ -1597,16 +1636,8 @@ def allocate(req, count): rtol=0, atol=0, ) - if exhaust_small_pool: - assert restored.tail_small_page_buffer_id is None - else: - assert restored.tail_small_page_buffer_id is not None - torch.testing.assert_close( - requests.small_page_buffers.buffer[restored.tail_small_page_buffer_id], - expected_checkpoints[-1], - rtol=0, - atol=0, - ) + # CPU 碎页直接经 LOAD 大页恢复运行态，不重复占用普通小页。 + assert restored.tail_small_page_buffer_id is None dst = requests.req_to_token_indexs[destination, :length] for pool in _compressed_kv_pools(manager): torch.testing.assert_close( @@ -1668,6 +1699,8 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): slots = _swa_slots(requests, req, positions) state = SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[req], dsv4_swa_write_slots=slots, @@ -1700,27 +1733,33 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): def test_cpu_cache_preserves_already_published_pages(cache): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_kv_buffer_to_cpu_cache - manager, requests = cache req, slots = _reserve(manager, requests, 2048) _randomize_compressed_kv(manager) layout = manager.cache_config pages = torch.full((3, layout.get_cpu_cache_big_page_bytes()), 221, dtype=torch.uint8, pin_memory=True) - page_ids = torch.tensor([1], dtype=torch.int32, device="cuda") - readies = torch.ones(1, dtype=torch.bool, device="cuda") + page_ids = [1] + readies = [True] state_id = manager.big_page_buffers.alloc_one_state_cache() manager.big_page_buffers.buffer[state_id].fill_(7) - checkpoint_ids = torch.tensor([state_id], dtype=torch.int64, device="cuda") - copy_kv_buffer_to_cpu_cache(manager, slots, page_ids, readies, checkpoint_ids, pages) + checkpoint_ids = [state_id] + manager.manager.write_req_to_cpu( + token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_page_layouts(), cpu_page_ready_flags=readies, + big_page_slot_ids=checkpoint_ids, + ) torch.cuda.synchronize() assert pages.eq(221).all() - readies.zero_() - copy_kv_buffer_to_cpu_cache(manager, slots, page_ids, readies, checkpoint_ids, pages) + readies[0] = False + manager.manager.write_req_to_cpu( + token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_page_layouts(), cpu_page_ready_flags=readies, + big_page_slot_ids=checkpoint_ids, + ) torch.cuda.synchronize() offset = 0 - for pool in manager.kv_buffer: - expected = pool[slots.long()].cpu().reshape(-1) + for state in manager.manager.get_full_attention_states(): + expected = state.storage[0, slots.long()].cpu().reshape(-1) torch.testing.assert_close(pages[1, offset : offset + expected.numel()], expected, rtol=0, atol=0) offset += expected.numel() assert pages[1, layout.get_cpu_cache_full_att_bytes() :].eq(7).all() @@ -1767,6 +1806,10 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le state = DeepseekV4InferStateInfo() state.is_prefill = True state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.input_ids = torch.tensor(input_ids[ready:], dtype=torch.int64, device="cuda") state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor([len(input_ids)], dtype=torch.int32, device="cuda") @@ -1825,6 +1868,10 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): infer_state = DeepseekV4InferStateInfo() infer_state.is_prefill = True infer_state.req_manager = requests + infer_state.att_state_manager = requests.memory.manager + infer_state.att_state_role = "main" + infer_state.first_logical_layer_id = 0 + infer_state.first_c4_logical_layer_id = 0 infer_state.mem_manager = manager infer_state.input_ids = torch.ones(5, dtype=torch.int64, device="cuda") infer_state.b_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda") diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 8032134734..6e830f2a2b 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -61,7 +61,7 @@ def test_unsupported_parallel_layouts_fail_before_inference(monkeypatch, balance def test_memory_check_rejects_sequence_larger_than_full_token_pool(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.mem_manager = SimpleNamespace(size=60416) + model.att_state_manager = SimpleNamespace(full_att_allocator=SimpleNamespace(size=60416)) model.batch_max_tokens = 8192 model.max_seq_length = 262169 model.args = SimpleNamespace(performance_mode=None) @@ -72,7 +72,7 @@ def test_memory_check_rejects_sequence_larger_than_full_token_pool(): def test_memory_check_still_requires_one_prefill_batch(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.mem_manager = SimpleNamespace(size=8192) + model.att_state_manager = SimpleNamespace(full_att_allocator=SimpleNamespace(size=8192)) model.batch_max_tokens = 8192 model.max_seq_length = 262169 model.args = SimpleNamespace(performance_mode=None) diff --git a/unit_tests/models/deepseek_v4/test_vision_integration.py b/unit_tests/models/deepseek_v4/test_vision_integration.py index 76322216e6..9ffd9f3962 100644 --- a/unit_tests/models/deepseek_v4/test_vision_integration.py +++ b/unit_tests/models/deepseek_v4/test_vision_integration.py @@ -204,7 +204,7 @@ def test_swa_index_adds_bidirectional_image_visibility(): HOLD_SWA_SLOT=256, ) requests = SimpleNamespace( - cache_config=SimpleNamespace(swa_page_token_num=128), + config=SimpleNamespace(page_token_size=128), swa_capacity_tokens=128, HOLD_SWA_SLOT=128, HOLD_REQUEST_ID=1, @@ -276,11 +276,13 @@ def allocate(req, count): module.args = SimpleNamespace(page_size=256) module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=False, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), - mem_manager=SimpleNamespace( - operator=SimpleNamespace(load_cpu_cache_to_gpu=lambda **kwargs: loaded.append(req.cur_kv_len)) + att_state_manager=SimpleNamespace( + get_cpu_page_layouts=lambda: {}, + read_req_from_cpu=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), ), ), ) @@ -288,6 +290,7 @@ def allocate(req, count): lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), deref_pages=lambda page_list: released.extend(page_list), ) + module.cpu_cache_client.cpu_kv_cache_tensor = torch.empty((1, 1)) module.init_sync_group = None module.need_sync_compute_stream = lambda: False monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 8192) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index d9a2086ebf..f0490a6be4 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -21,7 +21,7 @@ fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import destindex_copy_kv_flashmla_fp8 +from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( _build_prefill_row_table_kernel, @@ -115,7 +115,7 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid source[1].mul_(1e-5) mem_index = torch.randperm(31, device="cuda")[:11].to(index_dtype) packed_kv = torch.full((31, 1, 720), 71, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) groups = source[..., :448].float().reshape(11, 1, 7, 64) scales = torch.exp2(torch.ceil(torch.log2((groups.abs().amax(-1) / 448).clamp_min(1e-4)))) values = (groups / scales.unsqueeze(-1)).to(torch.float8_e4m3fn) @@ -130,7 +130,7 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid assert (packed_kv[unused] == 71).all() # The packer is graph-safe for both source layouts and mutable destination mapping. - graph, _ = _capture_cuda_graph(lambda: destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv)) + graph, _ = _capture_cuda_graph(lambda: write_fp8_e8m0_packed_kv(source, mem_index, packed_kv)) graph.replay() assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) source.add_(0.125) @@ -173,7 +173,7 @@ def test_fp8_nope_decode_and_cuda_graph(num_heads): packed_kv = torch.zeros(4096, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) mem_index = torch.arange(4096, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) # Includes a HOLD row, short tails, and three consecutive MTP positions. lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda", dtype=torch.int32) @@ -201,7 +201,7 @@ def check(actual): check(out) q.normal_() source.normal_() - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) lengths.copy_(lengths.roll(1)) indices.random_(0, 4096) indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) @@ -226,7 +226,7 @@ def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chun source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) prefill_kv = source[prefill_rows].clone().add_(0.25) packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) @@ -274,7 +274,7 @@ def check(out): ragged_mem_index.copy_(ragged_mem_index.roll(1)) prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) topk_indices.copy_(topk_indices.roll(1, dims=1)) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) graph.replay() check(out) @@ -291,7 +291,7 @@ def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softm prefill_mem_index = ragged_mem_index[prefill_rows].long() prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) @@ -335,7 +335,7 @@ def check(actual): prefill_kv.normal_() ragged_mem_index.copy_(ragged_mem_index.roll(1)) prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) topk_indices.copy_(topk_indices.roll(1, dims=1)) graph.replay() check(output) @@ -369,7 +369,7 @@ def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q) ragged_n = 65536 packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) starts = torch.where( @@ -437,7 +437,7 @@ def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): num_q = 16 packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) @@ -477,7 +477,7 @@ def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) # Requests deliberately have different spans; starts are not fixed-stride. @@ -515,7 +515,7 @@ def check(actual): prefill_kv.normal_() ragged.copy_(ragged.roll(1)) prefill_mem.copy_(ragged[current_rows]) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) topk.copy_(topk.roll(1, dims=1)) graph.replay() check(output) @@ -531,7 +531,7 @@ def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() starts = torch.zeros(1, device="cuda", dtype=torch.int32) @@ -587,7 +587,7 @@ def run(active_state=state): q.normal_() source.normal_() prefill_kv.normal_() - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers @@ -600,7 +600,7 @@ def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 @@ -657,7 +657,7 @@ def run(active_state=state): assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() source.normal_() - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) prefill_kv.normal_() q.normal_() eager = run() @@ -683,7 +683,7 @@ def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (1, 2049), device="cuda", dtype=torch.int32) @@ -751,7 +751,7 @@ def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (batch, 2049), device="cuda", dtype=torch.int32) @@ -808,7 +808,7 @@ def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batc source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(8, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) counts = [2051] if batch == 1 else [0, 2048, 2049, 2050, 2051, 1985, 2051, 2051] @@ -896,7 +896,7 @@ def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_h source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) # BMM-style B/H stride: (512, batch * 512, 1). q = torch.randn(num_heads, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) assert q.stride() == (512, batch * 512, 1) @@ -979,7 +979,7 @@ def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypat source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(1, 8, 512, device="cuda", dtype=torch.bfloat16) late = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) late[:, 2175] = 17 @@ -1024,7 +1024,7 @@ def test_fp8_sparse_direct_decode_strided_inputs_and_packed_row_padding(monkeypa capacity = 4096 source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 1, 4352), -1, device="cuda", dtype=index_dtype)[..., ::2] indices[..., :2051] = torch.randint(0, capacity, (batch, 1, 2051), device="cuda", dtype=index_dtype) @@ -1103,7 +1103,7 @@ def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): topk[..., :7] = torch.randint(0, capacity, (batch, 1, 7), device="cuda", dtype=topk.dtype) packed = torch.empty(capacity, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) assert can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) actual = fp8_sparse_direct_decode(q, packed, topk, 0.0625) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) @@ -1119,7 +1119,7 @@ def test_fp8_sparse_direct_decode_large_workload_retains_native_graph(monkeypatc torch.manual_seed(53700 + num_heads + batch) packed = torch.empty(4096, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(4096, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(4096, device="cuda", dtype=torch.int32), packed) q = torch.randn(batch, num_heads, 512, device="cuda", dtype=torch.bfloat16) topk = torch.full((batch, 1, 2176), -1, device="cuda", dtype=torch.int32) topk[..., :2051] = torch.randint(0, 4096, (batch, 1, 2051), device="cuda", dtype=torch.int32) @@ -1155,7 +1155,7 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) @@ -1275,3 +1275,42 @@ def fake_context(*args): ) assert state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads assert not queries + + +def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + + args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) + config = FullAttentionConfig(logical_layer_groups=((3,),), global_head_num=1, head_dim=512, + dtype=torch.bfloat16, replicated=True) + state = create_mla_state( + config=config, kv_type="fp8kv_dsa", packed=True, fp8_format="group64_e8m0", page_size=4, + ) + state.allocate_storage(token_capacity=16) + state.storage.zero_() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + indices = torch.arange(4, dtype=torch.int64, device="cuda") + kv = torch.randn((4, 1, 512), dtype=torch.bfloat16, device="cuda") + state.write_full_att_kv(logical_layer_id=3, token_indices=indices, kv=kv) + expected = torch.zeros((4, 1, 584), dtype=torch.uint8, device="cuda") + write_fp8_e8m0_packed_kv(kv, indices, expected) + torch.testing.assert_close(state.get_layer_view(logical_layer_id=3).kv[:4], expected, rtol=0, atol=0) + state.get_layer_view(logical_layer_id=3).indexer_k[:4].random_(0, 256) + cpu_layout = manager.get_cpu_page_layouts()[state.state_name] + cpu_pages = torch.empty((1, manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + state.write_to_cpu(token_indices=indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], layout=cpu_layout) + state.read_from_cpu(cpu_page_buffers=cpu_pages, cpu_page_indices=[0], token_indices=indices + 4, layout=cpu_layout) + pd_layout = manager.get_pd_token_page_layouts()[state.state_name] + pd_pages = torch.empty((1, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(token_indices=indices, pd_page_buffers=pd_pages, pd_page_index=0, layout=pd_layout) + state.read_from_pd_page(pd_page_buffers=pd_pages, token_indices=indices + 8, pd_page_index=0, layout=pd_layout) + torch.cuda.synchronize() + for begin in (4, 8): + actual = state.get_layer_view(logical_layer_id=3) + torch.testing.assert_close(actual.kv[begin:begin + 4], actual.kv[:4], rtol=0, atol=0) + torch.testing.assert_close(actual.indexer_k[begin:begin + 4], actual.indexer_k[:4], rtol=0, atol=0) diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py index 9444485ab6..1e87b85a1b 100644 --- a/unit_tests/models/sensenova_loop/test_loop.py +++ b/unit_tests/models/sensenova_loop/test_loop.py @@ -8,7 +8,7 @@ from lightllm.models.registry import ModelRegistry from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types -from lightllm.common.state_cache_manager.linear_att import LinearAttCacheConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -35,10 +35,13 @@ def _config(layers=32, start=4, end=28, count=3): def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch, start, end, count): config = _config(start=start, end=end, count=count) args = StartArgs(model_dir="unused", tp=4, dp=1, data_type="bfloat16", mtp_mode="eagle_with_att", mtp_step=3) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) + monkeypatch.setattr("lightllm.models.qwen3next.config.get_env_start_args", lambda: args) monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) - cache = LinearAttCacheConfig.load_from_args() + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) + cache = Qwen3NextResourceConfig.from_model_config(normalized, args) + full_layers = cache.get_full_att_layer_indices() + linear_layers = tuple(i for i in range(cache.all_layer_num) if i not in full_layers) indices = get_loop_layer_indices(config["loop_settings"], 32) kinds = get_loop_layer_types(config["loop_settings"], config["text_config"]) assert len(kinds) == 32 + (end - start) * (count - 1) @@ -46,18 +49,18 @@ def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch slots = {} for execution, (physical, kind) in enumerate(zip(indices, kinds)): if kind == "full_attention": - slot = cache.get_full_att_kv_layer_index(execution) + slot = full_layers.index(execution) assert slot == full full += 1 else: - slot = cache.get_linear_state_layer_index(execution) + slot = linear_layers.index(execution) assert slot == linear linear += 1 assert slot not in slots.setdefault(physical, set()) slots[physical].add(slot) - assert cache.get_main_model_full_att_layer_num() == full - assert cache.linear_layer_num == linear - assert cache.get_full_att_kv_layer_index(len(indices)) == full # Independent MTP draft. + assert len(full_layers) == full + assert len(linear_layers) == linear + assert len(set(full_layers) | set(linear_layers)) == len(indices) if (start, end, count) == (4, 28, 3): assert (cache.all_layer_num, full, linear) == (80, 20, 60) if (start, end, count) == (4, 28, 2): @@ -122,7 +125,7 @@ def test_mhc_registry_and_embedded_mtp_reuse_existing_qwen_draft(tmp_path, strea draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 4 draft._init_config() assert draft.config["n_layer"] == draft.config["full_attention_interval"] == 1 - assert draft.config["layer_types"] == config["text_config"]["layer_types"] + assert draft.config["layer_types"] == ["full_attention"] def test_loop_tokenizer_loads_without_hf_loop_autoconfig(tmp_path): @@ -158,21 +161,23 @@ def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_lay if custom_layout: config["text_config"]["layer_types"] = ["full_attention", "linear_attention"] * 16 original = deepcopy(config) - args = StartArgs(tp=4, dp=1, data_type="bfloat16") - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) + args = StartArgs(tp=4, dp=1, data_type="bfloat16", linear_att_hash_page_size=256) + monkeypatch.setattr("lightllm.models.qwen3next.config.get_env_start_args", lambda: args) monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) - monkeypatch.setattr( - "lightllm.models.qwen3next.model.ReqManagerForMamba", - lambda *args, linear_config: SimpleNamespace(linear_config=linear_config), + config["text_config"]["num_attention_heads"] = 16 + original = deepcopy(config) + normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) + raw_cache = Qwen3NextResourceConfig.from_model_config(normalized, args) + logical_ids = tuple(range(normalized["num_hidden_layers"])) + full, = SenseNovaLoopTpPartModel.create_full_att_states( + model_config=normalized, args=args, model_logical_layer_ids=logical_ids, ) - raw_cache = LinearAttCacheConfig.load_from_args() - model = object.__new__(SenseNovaLoopTpPartModel) - model.args = args - model.max_req_num, model.batch_max_tokens, model.max_seq_length = 2, 8, 16 - model._init_req_manager() - assert dataclasses.asdict(raw_cache) == dataclasses.asdict(model.linear_config) - assert model.req_manager.linear_config is model.linear_config + bounded, = SenseNovaLoopTpPartModel.create_bounded_att_states( + model_config=normalized, args=args, model_logical_layer_ids=logical_ids, + ) + full_layers = raw_cache.get_full_att_layer_indices() + assert full.config.logical_layer_groups == tuple((i,) for i in full_layers) + assert bounded.config.logical_layer_ids == tuple(i for i in logical_ids if i not in full_layers) assert config == original diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py index 5f0c3e450f..5efcfeba3e 100644 --- a/unit_tests/server/test_pd_start_args.py +++ b/unit_tests/server/test_pd_start_args.py @@ -91,28 +91,18 @@ def test_disabled_cpu_cache_keeps_page_size_unset(startup_args): assert startup_args.cpu_cache_token_page_size is None -def test_other_memory_manager_selection_uses_default(monkeypatch): - from types import SimpleNamespace - import lightllm.models as models - from lightllm.common.kv_cache_mem_manager import MemoryManager - from lightllm.common.kv_cache_mem_manager import mem_utils - - for name in ("DeepseekV4TpPartModel", "Deepseek3_2TpPartModel", "Deepseek2TpPartModel"): - monkeypatch.setitem(models.__dict__, name, type(name, (), {})) - monkeypatch.setattr(mem_utils, "get_llm_model_class", lambda: type("Llama", (), {})) - monkeypatch.setattr( - mem_utils, "get_env_start_args", lambda: SimpleNamespace(model_dir="unused", llm_kv_type="None") - ) - mem_utils.select_mem_manager_class.cache_clear() - try: - assert mem_utils.select_mem_manager_class() is MemoryManager - finally: - mem_utils.select_mem_manager_class.cache_clear() +def test_default_attention_factory_uses_gqa(): + import torch + from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState, create_gqa_state + + state = create_gqa_state(config=FullAttentionConfig( + logical_layer_groups=((0,),), global_head_num=4, head_dim=128, dtype=torch.bfloat16, + ), device="meta") + assert isinstance(state, GQAFullAttentionState) def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_type(startup_args, monkeypatch, tmp_path): - from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager - from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args, get_added_mtp_kv_layer_num, get_llm_data_type from lightllm.utils.llm_utils import get_llm_model_class from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta @@ -143,18 +133,21 @@ def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_typ get_added_mtp_kv_layer_num, get_llm_data_type, get_llm_model_class, - select_mem_manager_class, calcu_cpu_cache_meta, ) for function in cached_functions: function.cache_clear() try: - assert select_mem_manager_class() is DeepseekV4MemoryManager + layouts = build_cpu_cache_layouts() + assert any("compressed_c4" in name for name in layouts) + assert any("compressed_c128" in name for name in layouts) + assert any("swa_w128_fp8" in name for name in layouts) + assert any("compressor_kv_w8_s4" in name for name in layouts) + assert not any("compressor_kv_w128" in name for name in layouts) meta = calcu_cpu_cache_meta() assert meta.page_num > 0 - assert meta.token_page_size == meta.layer_num == meta.num_heads == 1 assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - assert meta.calcu_one_page_size() == meta.head_dim + assert meta.page_bytes == max(layout.byte_end for layout in layouts.values()) assert get_env_start_args().llm_kv_type == "None" finally: for function in cached_functions: From 5ceae75ec084572cf148a7c1e0eb31fb7f108091 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 15:20:45 +0000 Subject: [PATCH 29/56] refactor: isolate attention configs and clarify packed FP8 encoding Define independent full-attention configs, keep packed buffer dtype internal, and update model factories and cache documentation. Handle variable-length CUDA IPC exports and refresh regression fixtures for independent main/draft states. --- docs/cache_manager_design.md | 30 +-- .../attention_state/full_att/__init__.py | 9 +- .../attention_state/full_att/base_state.py | 60 ------ .../full_att/create_gqa_state.py | 8 +- .../full_att/create_mla_state.py | 25 +-- .../attention_state/full_att/gqa_state.py | 46 ++++- .../full_att/group_quantized_state.py | 60 ++++-- .../attention_state/full_att/mla_state.py | 55 +++-- .../full_att/packed_mla_state.py | 195 +++++++++--------- .../full_att/static_fp8_state.py | 42 +++- lightllm/models/bloom/model.py | 4 +- lightllm/models/deepseek2/model.py | 7 +- lightllm/models/deepseek3_2/model.py | 7 +- lightllm/models/gemma3/model.py | 4 +- lightllm/models/gemma4/model.py | 4 +- lightllm/models/gemma_2b/model.py | 4 +- lightllm/models/glm5_next/model.py | 7 +- lightllm/models/glm5_next_mtp/model.py | 7 +- lightllm/models/llama/model.py | 4 +- lightllm/models/mistral/model.py | 4 +- lightllm/models/mixtral/model.py | 4 +- lightllm/models/qwen3next/model.py | 4 +- lightllm/models/starcoder/model.py | 4 +- lightllm/models/starcoder2/model.py | 4 +- .../model_infer/shared_att_resources_utils.py | 9 +- test/cpu_cache_kernel/test_speed.py | 4 +- .../common/att_resources/manager_fixtures.py | 8 +- .../common/att_resources/model_helpers.py | 16 +- .../common/att_resources/test_capacity.py | 92 ++++----- .../att_resources/test_compute_binding.py | 2 +- .../common/att_resources/test_dsv4_state.py | 4 +- .../att_resources/test_fp8_compute_binding.py | 1 - .../test_fp8_e8m0_packed_mla_state.py | 17 +- .../att_resources/test_full_att_cpu_kernel.py | 4 +- .../common/att_resources/test_full_state.py | 20 +- .../test_full_storage_initialization.py | 7 +- .../att_resources/test_gqa_pd_kernel.py | 8 +- .../common/att_resources/test_gqa_transfer.py | 6 +- .../test_hybrid_layer_mapping.py | 112 ++++++++++ .../common/att_resources/test_layer_groups.py | 11 +- .../test_layer_state_selection.py | 6 +- .../att_resources/test_model_calibration.py | 8 +- .../att_resources/test_model_resources.py | 8 +- .../test_quantized_full_state.py | 12 +- .../test_quantized_paged_storage.py | 2 +- .../att_resources/test_runtime_transport.py | 47 +++++ .../att_resources/test_state_creation.py | 32 ++- .../common/att_resources/test_tp_config.py | 24 ++- .../state_cache_manager/test_linear_att.py | 64 ------ unit_tests/models/glm5_next/test_cache.py | 4 +- unit_tests/models/glm5_next/test_fp8_kv.py | 7 +- unit_tests/models/glm5_next/test_pd_cache.py | 4 +- .../test_paged_side_allocations.py | 1 - .../model_infer/test_cpu_resource_layout.py | 3 + .../test_hybrid_cpu_state_cache.py | 6 +- unit_tests/server/test_pd_start_args.py | 4 +- 56 files changed, 656 insertions(+), 495 deletions(-) create mode 100644 unit_tests/common/att_resources/test_hybrid_layer_mapping.py delete mode 100644 unit_tests/common/state_cache_manager/test_linear_att.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index dce91a7fe3..94d9721689 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -132,7 +132,7 @@ AttentionStateManager # 公共接口与 allocator 入口 ```text FullAttentionState ├── capacity / page_size / device # 普通 token 容量、HOLD 页大小及分配设备 -├── storage_dtype # 实际 KV 编码类型；普通方案使用 config.dtype，INT4/8 为 int8，静态 FP8 为 uint8 +├── config # 每种具体方案的独立配置；packed 的物理 dtype/行宽由编码推导 ├── k_storage / v_storage # GQA 及其量化方案分别分配 K/V；MLA 使用单份 latent storage └── layer_mapping # 本模型的层映射 @@ -243,7 +243,7 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 ### 3.5 State 成员命名、资源身份与唯一性 -`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。`FullAttentionConfig`（含量化与 MLA 子类）、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 +`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。各具体 full-attention 配置、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 管理器成员名与 `state.state_name` 完全一致，统一使用 `{model_role}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState`、`AttentionTailState` 类型表达，名称中不再重复 `full_att`、`bounded_att`、`att_tail`。 @@ -311,7 +311,7 @@ lightllm/common/ │ │ └── state_page_allocator.py # StatePageAllocator：bounded_att 大小页各一个实例 │ ├── full_att/ │ │ ├── __init__.py -│ │ ├── base_state.py # FullAttentionConfig / FullAttentionState：通用配置与抽象接口 +│ │ ├── base_state.py # FullAttentionState：抽象行为接口 │ │ ├── create_gqa_state.py # GQA 统一创建入口：选择普通精度、INT4/INT8 或静态 FP8 │ │ ├── gqa_state.py # GQAFullAttentionState：兼容 MHA / GQA / MQA │ │ ├── gqa_transfer.py # GQA CPU/PD 共用的 K/V 页传输入口与 kernel @@ -367,7 +367,7 @@ lightllm/utils/ `allocators/` 与三类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 -`full_att/base_state.py` 同时定义 `FullAttentionConfig` 和 `FullAttentionState`；`bounded_att/base_state.py`、`att_tail/base_state.py` 分别定义另外两类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 +`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口，各具体配置与对应 State 放在同一文件中；`bounded_att/base_state.py`、`att_tail/base_state.py` 分别定义另外两类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 `KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 @@ -420,16 +420,16 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器和 K-pool。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 -GQA/MLA 的 `config` 使用逻辑 `FullAttentionConfig`。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.head_dim` 始终传逻辑 latent 维度，`packed=True` 时由入口计算包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 +GQA/MLA 的创建入口分别接收 `GQAAttentionConfig`、`MLAAttentionConfig`，每种具体 State 使用独立配置，不继承通用配置。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.latent_dim` 始终传逻辑 latent 维度，`packed=True` 时由 packed 配置推导包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 MLA 的 packed FP8 编码由 `fp8_format` 显式指定：默认 `"group128_fp32"` 使用原有 `PackedMLAFullAttentionState`；`"group64_e8m0"` 使用独立的 `FP8E8M0PackedMLAState`， -其配置为 `FP8E8M0PackedMLAConfig`。创建入口不通过 `head_dim == 512` 等维度条件推断格式。 +其配置为 `FP8E8M0PackedMLAConfig`。创建入口不通过 `latent_dim == 512` 等维度条件推断格式。 GLM main/draft 在各自的模型工厂中传入 `fp8_format="group64_e8m0"`；未量化时仍返回普通 packed State。 `FP8E8M0PackedMLAState` 固定使用 448 维 FP8、64 维 BF16 和七个 E8M0 scale（补齐为八字节）， KV 共 584 字节，随后保存 indexer key 及一个 FP32 scale；最后 64 维保留 BF16 不代表使用 RoPE。 -配置独立定义，不继承 `FullAttentionConfig`；只记录逻辑层分组、indexer 维度和 TP 坐标。 +配置独立定义，不继承其他方案的配置；只记录逻辑层分组、indexer 维度和 TP 坐标。 固定 latent 维度、单 head、副本规则和 uint8 编码由实现确定，不作为重复输入。 `token_bytes` 是含 16 字节对齐 padding 的物理行宽，`payload_bytes` 是不含行尾 padding 的有效字节数。 例如 indexer 维度为 128 时，GPU 行宽为 720 字节，CPU/PD 每个 token 只传 716 字节。 @@ -509,7 +509,7 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 #### 逻辑层与物理存储层 -具体 State 配置中的 `logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。GQA/MLA 使用 `FullAttentionConfig` 及其具体配置，序列压缩历史使用独立的 `CompressedAttentionConfig`，不要求所有配置继承同一基类。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 +具体 State 配置中的 `logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。GQA/MLA 分别使用各自的独立配置，序列压缩历史使用独立的 `CompressedAttentionConfig`，不要求所有配置继承同一基类。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 ```python logical_layer_groups = ( @@ -552,7 +552,7 @@ FullAttentionState └── CompressedAttentionState # V4 的 C4/C128 压缩块 ``` -所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。`FullAttentionConfig` 与抽象接口共同定义在 `full_att/base_state.py`，量化和 packed 专属配置可扩展该配置；配置结构的继承不代表 State 实现的继承。 +所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。配置也不建立继承关系：普通 GQA 使用 `GQAAttentionConfig`；分组量化使用 `GroupQuantizedAttentionConfig`，以 `scale_dtype` 明确 scale 精度；静态 FP8 使用 `StaticFP8AttentionConfig`，不要求调用方传入固定的物理 dtype。普通 MLA 使用 `MLAAttentionConfig`；packed MLA 使用 `PackedMLAAttentionConfig`，以 `latent_dim` 表达逻辑宽度，`dtype` 直接记录 FP16、BF16 或 `torch.float8_e4m3fn` 编码。FP8 行的 uint8 承载类型只在存储分配函数中确定，不提供 `storage_dtype` 配置或属性；物理行宽及 indexer 偏移由编码推导。MLA 固定单份 latent 并在 TP rank 间复制，不提供 GQA 的 head 数或 `replicated` 开关。 `gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 @@ -1790,7 +1790,7 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_attention_states()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -`create_full_att_states()` 由具体模型直接构造 `FullAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 +`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig` 或 `MLAAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 `AttentionStateManager.initialize_storage()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `allocate_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 @@ -1821,10 +1821,14 @@ State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_si | 配置 | 保存的信息 | 计算属性 | | --- | --- | --- | -| `FullAttentionConfig` | `global_head_num`、`replicated`、TP 坐标及原有层/维度/dtype 信息 | `head_num` 返回当前 rank 的 KV head 数 | +| `GQAAttentionConfig` | 逻辑层分组、全局 KV head 数、`head_dim`、KV 浮点 `dtype`、TP 坐标与 `replicated` | `head_num` 返回当前 rank 的单份 K 或 V 的 head 数 | +| `GroupQuantizedAttentionConfig` | GQA 几何与 TP 信息、`bits`、`group_size`、动态 `scale_dtype` | `head_num`；KV 编码固定为 INT4/INT8，scale 精度独立描述 | +| `StaticFP8AttentionConfig` | GQA 几何与 TP 信息、`scale_mode`，不传物理 dtype 或标定 Tensor | `head_num`；FP8 编码固定，标定通过独立流程绑定 | +| `MLAAttentionConfig` | 逻辑层分组、`latent_dim`、浮点 `dtype` 和 TP 坐标 | 固定复制单份 latent，无 head 数配置 | +| `PackedMLAAttentionConfig` | 逻辑层分组、`latent_dim`、KV 编码 `dtype`（FP16/BF16/FP8）、indexer 与编码参数、TP 坐标 | `kv_bytes`、`token_bytes`、`indexer_offset` 从编码推导；FP8 行使用内部 uint8 buffer 承载混合字段，不另设 `storage_dtype` 属性 | | `ConvSSMStateConfig` | 层、`global_k_heads`、`global_v_heads`、TP 坐标及 Conv/SSM 维度/dtype 信息；不保存资源身份 | `num_k_heads`、`num_v_heads` 分别返回本地 head 数，`conv_channels` 返回本地卷积通道数 | -full-attention 的 `replicated=True` 表示每个 rank 保存全部 head，例如复制的 MLA latent；否则按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 +GQA 的 `replicated=True` 表示每个 rank 保存全部 head；MLA 固定复制单份 latent，不设置该参数。GQA 的 `replicated=False` 时按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 模型类统一暴露 `normalize_model_config(model_config)`：不创建实例、不加载权重或初始化 CUDA，返回新的有效配置，重复调用保持结果不变。模型启动读取配置后、CPU 服务创建 meta State 前，以及各模型 State 工厂，都使用这个入口。DFlash 的嵌套覆盖、原生 MTP 的自有层数等差异由子类实现，不再保留 `_resource_config` 这条单独的转换路径。`basemodel_utils.py` 只处理通用 text_config 展开和字段别名，不反射调用模型方法，也不选择 attention 算法。 @@ -1961,7 +1965,7 @@ SWA 为 `main_swa_w128_fp8_group64_e8m0_state`。同一个压缩器实现也可 | `ReqManager` | 请求管理 | 只保存请求槽位和逻辑 token 到物理 token 的映射 | 压缩历史通过独立的 `create_compressed_att_state()` 创建，不经过 MLA 创建入口。 -`CompressedAttentionConfig` 是独立的原子布局定义，不继承 `FullAttentionConfig`，也不持有 +`CompressedAttentionConfig` 是独立的原子布局定义，不继承其他 full-attention 配置，也不持有 `DSV4AttentionConfig`。它直接保存逻辑层分组 `logical_layer_groups`、压缩比例 `compress_ratio`、 KV/RoPE 维度 `kv_dim` / `rope_dim`、每块覆盖的原始 token 数 `block_token_size`、可选 `indexer_dim` 和当前 TP 信息。`indexer_dim=0` 表示不保存 indexer，不根据压缩比例推断。 diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index b65340cfb3..d3e1f56a9d 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -1,15 +1,16 @@ from .create_mla_state import create_mla_state from .create_gqa_state import create_gqa_state -from .base_state import FullAttentionConfig, FullAttentionState -from .gqa_state import GQAFullAttentionState -from .mla_state import MLAFullAttentionState +from .base_state import FullAttentionState +from .gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .mla_state import MLAAttentionConfig, MLAFullAttentionState __all__ = [ "create_gqa_state", "create_mla_state", "FullAttentionState", "GQAFullAttentionState", - "FullAttentionConfig", + "GQAAttentionConfig", + "MLAAttentionConfig", "MLAFullAttentionState", ] diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index 83fe4efd8b..1a44a487ef 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -1,11 +1,8 @@ from __future__ import annotations from functools import lru_cache - - from typing import List from abc import ABC, abstractmethod -from dataclasses import dataclass from typing import Any import torch @@ -17,63 +14,6 @@ ) -@dataclass(frozen=True) -class FullAttentionConfig: - # 物理存储层与逻辑推理层的对应关系，外层顺序就是所有 Tensor 的 layer 轴顺序。 - # 每个非空内层元组列出共享同一份存储的逻辑层号；层号可以不连续，但不能重复出现。 - # 例如 ((0, 4), (8,)) 只分配两层：逻辑层 0/4 使用 storage[0]，逻辑层 8 使用 storage[1]。 - # 不共享时写成 ((0,), (4,), (8,))，此时分配三层；容量和传输按 len(logical_layer_groups) 计算。 - # 组内共享该方案的全部字段（KV、scale、indexer 等），因此布局及配套参数必须兼容。 - # 分组不决定执行顺序或写入者：模型负责让共享消费者跳过 KV 写入，避免覆盖生产层数据。 - logical_layer_groups: tuple[tuple[int, ...], ...] - # TP 切分前的 KV head 数：GQA 的 K/V 各有这么多个，不把二者相加，也不计 TP 副本。 - # 不是 query head 数；例如全局 8 个 KV head，实际有 8 个 K head 和 8 个 V head。 - # MLA 的压缩 latent 作为一个 KV 字段存储，通常设为 1，并在各 TP rank 复制。 - global_head_num: int - # GQA 中每个 K/V head 的原始向量维度；INT4 的实际编码宽度由具体 State 换算。 - # 普通 MLA 中为 latent 宽度；packed MLA 中为整行元素数，包含 indexer、scale 和 padding。 - head_dim: int - # 具体方案使用的数据类型：普通 GQA/MLA 为 KV 精度，分组 INT4/INT8 为 scale 精度。 - # 静态 FP8 和 FP8 packed MLA 使用 uint8 承载编码字节，实际字段类型由 State 解释。 - dtype: torch.dtype - # 当前 State 所属模型的 TP 组内 rank，不是全局进程 rank 或 DP rank。 - tp_rank: int = 0 - # 当前模型 TP 组大小，用于计算本地 head 数以及传输时的全局 head 范围。 - tp_world_size: int = 1 - # True：每个 rank 都保存全部全局 heads，例如 MLA latent。 - # False：heads 足够时均匀切分；heads 少于 TP rank 数时，每组相邻 ranks 复制一个 head。 - # 例如 global_head_num=2、TP=8 时，rank 0..3 保存 head 0，rank 4..7 保存 head 1。 - replicated: bool = False - - @property - def head_num(self) -> int: - """当前 TP rank 实际存储的 KV head 数。""" - # 仍是单份 K 或 V 的 head 数；GQA 的 K/V 独立分配，各自的 head 轴大小均为 head_num。 - return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) - - def __post_init__(self) -> None: - # 只保留这一份分组配置；层号映射在 State 构造时展开，不另存逻辑/物理层数。 - assert isinstance(self.logical_layer_groups, tuple) and all( - isinstance(logical_layer_group, tuple) and logical_layer_group - for logical_layer_group in self.logical_layer_groups - ), "layer groups must be a tuple of nonempty tuples" - logical_layer_ids = [ - logical_layer_id - for logical_layer_group in self.logical_layer_groups - for logical_layer_id in logical_layer_group - ] - assert all(isinstance(logical_layer_id, int) and logical_layer_id >= 0 for logical_layer_id in logical_layer_ids), ( - "logical layer IDs must be nonnegative integers" - ) - assert len(logical_layer_ids) == len(set(logical_layer_ids)), "a logical layer must belong to exactly one storage group" - assert self.global_head_num > 0 and 0 <= self.tp_rank < self.tp_world_size - assert ( - self.replicated - or self.global_head_num % self.tp_world_size == 0 - or self.tp_world_size % self.global_head_num == 0 - ), "KV heads require complete TP shards or replicas" - - class FullAttentionState(ABC): """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 diff --git a/lightllm/common/attention_state/full_att/create_gqa_state.py b/lightllm/common/attention_state/full_att/create_gqa_state.py index 042e9042e2..d25a7c37fc 100644 --- a/lightllm/common/attention_state/full_att/create_gqa_state.py +++ b/lightllm/common/attention_state/full_att/create_gqa_state.py @@ -6,14 +6,13 @@ import torch -from .base_state import FullAttentionConfig -from .gqa_state import GQAFullAttentionState +from .gqa_state import GQAAttentionConfig, GQAFullAttentionState from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState def create_gqa_state( - config: FullAttentionConfig, + config: GQAAttentionConfig, kv_type: Literal["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"] = "None", page_size: int = 1, device: str | torch.device = "cuda", @@ -32,7 +31,7 @@ def create_gqa_state( logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, head_dim=config.head_dim, - dtype=config.dtype, + scale_dtype=config.dtype, tp_rank=config.tp_rank, tp_world_size=config.tp_world_size, replicated=config.replicated, @@ -45,7 +44,6 @@ def create_gqa_state( logical_layer_groups=config.logical_layer_groups, global_head_num=config.global_head_num, head_dim=config.head_dim, - dtype=torch.uint8, tp_rank=config.tp_rank, tp_world_size=config.tp_world_size, replicated=config.replicated, diff --git a/lightllm/common/attention_state/full_att/create_mla_state.py b/lightllm/common/attention_state/full_att/create_mla_state.py index 05e3459f4b..bd2e5a6ecb 100644 --- a/lightllm/common/attention_state/full_att/create_mla_state.py +++ b/lightllm/common/attention_state/full_att/create_mla_state.py @@ -2,19 +2,18 @@ from __future__ import annotations -from dataclasses import replace from typing import Literal import torch -from .base_state import FullAttentionConfig, FullAttentionState -from .mla_state import MLAFullAttentionState +from .base_state import FullAttentionState +from .mla_state import MLAAttentionConfig, MLAFullAttentionState from .packed_mla_state import PackedMLAAttentionConfig, PackedMLAFullAttentionState from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState def create_mla_state( - config: FullAttentionConfig, + config: MLAAttentionConfig, kv_type: Literal["None", "fp8kv_dsa"] = "None", packed: bool = False, indexer_dim: int = 128, @@ -24,7 +23,7 @@ def create_mla_state( ) -> FullAttentionState: """创建普通或 packed MLA State，统一处理量化编码、TP 副本和物理行宽。 - config.head_dim 始终是逻辑 KV latent 维度，尚未包含 indexer、scale 和 padding。 + config.latent_dim 始终是逻辑 KV latent 维度，尚未包含 indexer、scale 和 padding。 packed=True 时把 indexer 字段存入同一行；fp8kv_dsa 仅适用于 packed MLA。 fp8_format 由模型显式选择，仅在 packed FP8 时生效；不根据 latent 维度推断编码。 group128_fp32 与 group64_e8m0 使用独立 State。 @@ -32,7 +31,7 @@ def create_mla_state( if not packed: assert kv_type == "None", f"unsupported MLA encoding: {kv_type}" return MLAFullAttentionState( - layout_config=replace(config, replicated=True), + layout_config=config, page_size=page_size, device=device, ) @@ -40,8 +39,7 @@ def create_mla_state( fp8 = kv_type == "fp8kv_dsa" assert not fp8 or fp8_format in ("group128_fp32", "group64_e8m0"), f"unknown packed MLA FP8 format: {fp8_format}" if fp8 and fp8_format == "group64_e8m0": - assert config.head_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" - assert config.global_head_num == 1, "FP8/E8M0 packed MLA requires one replicated latent" + assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" return FP8E8M0PackedMLAState( layout_config=FP8E8M0PackedMLAConfig( logical_layer_groups=config.logical_layer_groups, @@ -52,19 +50,12 @@ def create_mla_state( page_size=page_size, device=device, ) - storage_dtype = torch.uint8 if fp8 else config.dtype - # 原有 DSA 格式：512 FP8 + 16 字节 FP32 scales + 64 BF16，另附 indexer key/scale。 - kv_bytes = 656 if fp8 else config.head_dim * config.dtype.itemsize - row_bytes = (kv_bytes + indexer_dim + 4 + 15) // 16 * 16 packed_config = PackedMLAAttentionConfig( logical_layer_groups=config.logical_layer_groups, - global_head_num=config.global_head_num, - head_dim=row_bytes // storage_dtype.itemsize, - dtype=storage_dtype, + latent_dim=config.latent_dim, + dtype=torch.float8_e4m3fn if fp8 else config.dtype, tp_rank=config.tp_rank, tp_world_size=config.tp_world_size, - replicated=True, - kv_dim=config.head_dim, indexer_dim=indexer_dim, fp8_nope_dim=512 if fp8 else 0, ) diff --git a/lightllm/common/attention_state/full_att/gqa_state.py b/lightllm/common/attention_state/full_att/gqa_state.py index 0086963514..ffcbed33db 100644 --- a/lightllm/common/attention_state/full_att/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa_state.py @@ -3,6 +3,7 @@ from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod +from dataclasses import dataclass import torch @@ -11,7 +12,7 @@ from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionConfig, FullAttentionState +from .base_state import FullAttentionState from .gqa_transfer import transfer_gqa_pages from .dp_transfer import copy_gqa_tokens @@ -22,6 +23,47 @@ logger = init_logger(name=__name__) +@dataclass(frozen=True) +class GQAAttentionConfig: + """普通 GQA（含 MHA/MQA）的逻辑 KV 布局；不描述量化编码或物理行 padding。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # K/V 浮点存储精度。 + dtype: torch.dtype + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" + + class GQAFullAttentionState(FullAttentionState): """GQA 的 full-attention KV 存储，MHA 和 MQA 均为其特例。 @@ -42,7 +84,7 @@ class GQAFullAttentionState(FullAttentionState): def __init__( self, - layout_config: FullAttentionConfig, + layout_config: GQAAttentionConfig, page_size: int = 1, device: str | torch.device = "cuda", ) -> None: diff --git a/lightllm/common/attention_state/full_att/group_quantized_state.py b/lightllm/common/attention_state/full_att/group_quantized_state.py index 172af40d61..9029ca1a0c 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/group_quantized_state.py @@ -12,7 +12,7 @@ from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionConfig, FullAttentionState +from .base_state import FullAttentionState from .group_quantized_transfer import transfer_group_quantized_pages from .dp_transfer import copy_group_quantized_tokens @@ -24,20 +24,52 @@ @dataclass(frozen=True) -class GroupQuantizedAttentionConfig(FullAttentionConfig): - # 每个 KV 元素的量化位数，仅支持 4/8；INT4 将两个编码值打包到一个字节。 +class GroupQuantizedAttentionConfig: + """INT4/INT8 GQA 的独立布局；KV 编码为整数，动态 scale 使用单独精度。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # 每个 token、每组 K/V 量化 scale 的浮点精度，不是 KV 编码类型。 + scale_dtype: torch.dtype + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + # KV 量化位数；INT4 将两个编码值打包进一个字节。 bits: int = 8 - # 沿每个 head 的 dim 轴分组，每个 token、每组共用一个动态 scale；须整除 head_dim。 - # K/V 分别保存 scale，精度由继承的 dtype 指定。 + # 沿 head_dim 分组，每组共用一个动态 scale；必须为二次幂且整除 head_dim。 group_size: int = 8 + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + def __post_init__(self) -> None: - """创建配置时校验通用 TP 约束及量化分组要求。""" - super().__post_init__() + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" assert self.bits in (4, 8) and self.group_size >= 2 and not ( self.group_size & (self.group_size - 1) ), "INT4/INT8 quantization requires a power-of-two group size" - assert self.head_dim % self.group_size == 0 and self.dtype in ( + assert self.head_dim % self.group_size == 0 and self.scale_dtype in ( torch.float16, torch.bfloat16, torch.float32 ), "group scales require a floating dtype and complete groups" @@ -76,7 +108,7 @@ def __init__( def get_state_type(self) -> str: """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" - return f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.dtype)}_scale" + return f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.scale_dtype)}_scale" def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" @@ -88,14 +120,14 @@ def allocate_storage(self, token_capacity: int) -> None: len(self.config.logical_layer_groups), token_capacity + self.page_size, self.config.head_num, self.config.head_dim // self.config.group_size, ) - self.k_scales = torch.empty(size=shape, dtype=self.config.dtype, device=self.device) - self.v_scales = torch.empty(size=shape, dtype=self.config.dtype, device=self.device) + self.k_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) + self.v_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) def get_token_bytes(self) -> int: """累计量化 KV 编码与逐 token 分组 scale 的开销，无需分配存储。""" config = self.config head_bytes = config.head_dim * config.bits // 8 - head_bytes += config.head_dim // config.group_size * config.dtype.itemsize + head_bytes += config.head_dim // config.group_size * config.scale_dtype.itemsize return len(config.logical_layer_groups) * 2 * config.head_num * head_bytes # ======================================================================== @@ -354,8 +386,8 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: field_specs = ( ("k", torch.int8, config.head_dim * config.bits // 8, "dim"), ("v", torch.int8, config.head_dim * config.bits // 8, "dim"), - ("k_scale", config.dtype, config.head_dim // config.group_size, "group"), - ("v_scale", config.dtype, config.head_dim // config.group_size, "group"), + ("k_scale", config.scale_dtype, config.head_dim // config.group_size, "group"), + ("v_scale", config.scale_dtype, config.head_dim // config.group_size, "group"), ) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) diff --git a/lightllm/common/attention_state/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla_state.py index fed5d918b8..71b34a05a6 100644 --- a/lightllm/common/attention_state/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla_state.py @@ -3,6 +3,7 @@ from functools import lru_cache from typing import TYPE_CHECKING, List from math import prod +from dataclasses import dataclass import torch @@ -10,7 +11,7 @@ from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionConfig, FullAttentionState +from .base_state import FullAttentionState from .packed_transfer import transfer_packed_pages from .dp_transfer import copy_packed_tokens @@ -18,6 +19,33 @@ from lightllm.common.attention_state.types import StateTransferLayout +@dataclass(frozen=True) +class MLAAttentionConfig: + """普通 MLA 的独立 latent 布局；单份 latent 固定在各 TP rank 完整复制。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 编码前的完整 latent 宽度，包含模型需要的 RoPE 部分，不含 indexer 或 padding。 + latent_dim: int + # 未量化 latent 的浮点精度，也作为 MLA 工厂选择 packed 编码前的原始精度。 + dtype: torch.dtype + # 所属模型 TP 组内的 rank；各 rank 保存相同 latent，仅 rank 0 写出共享外部页。 + tp_rank: int = 0 + # 保存完整 latent 副本的 TP rank 数，不改变 latent 的维度或物理行宽。 + tp_world_size: int = 1 + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.latent_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + + class MLAFullAttentionState(FullAttentionState): """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" @@ -27,7 +55,7 @@ class MLAFullAttentionState(FullAttentionState): def __init__( self, - layout_config: FullAttentionConfig, + layout_config: MLAAttentionConfig, page_size: int = 1, device: str | torch.device = "cuda", ) -> None: @@ -56,8 +84,8 @@ def allocate_storage(self, token_capacity: int) -> None: size=( len(self.config.logical_layer_groups), token_capacity + self.page_size, - self.config.head_num, - self.config.head_dim, + 1, + self.config.latent_dim, ), dtype=self.config.dtype, device=self.device, @@ -66,10 +94,7 @@ def allocate_storage(self, token_capacity: int) -> None: def get_token_bytes(self) -> int: """由配置计算当前 TP rank 的单 token latent 字节数，无需分配存储。""" config = self.config - return ( - len(config.logical_layer_groups) * config.head_num * config.head_dim - * config.dtype.itemsize - ) + return len(config.logical_layer_groups) * config.latent_dim * config.dtype.itemsize # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 @@ -246,13 +271,9 @@ def read_from_pd_page( # ======================================================================== def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: - """描述唯一的 MLA latent 字段，外部页保存全局 heads，副本只由首 rank 写出。""" + """各 rank 复制同一 latent，仅 rank 0 写出共享外部页，所有 rank 均可读回。""" config = self.config - shape = (len(config.logical_layer_groups), page_tokens, config.global_head_num, config.head_dim) - # 本地分片可在多个 rank 复制；先去掉副本编号，再定位外部页的 head 起点。 - local_head_count = config.head_num - replica_count = config.tp_world_size * local_head_count // config.global_head_num - head_start = config.tp_rank // replica_count * local_head_count + shape = (len(config.logical_layer_groups), page_tokens, 1, config.latent_dim) return ( TransferRegion( field="kv", @@ -262,9 +283,9 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: axes=("layer", "token", "head", "dim"), shape=shape, axis_mappings=( - AxisMapping(axis="head", local_start=0, external_start=head_start, length=local_head_count), - AxisMapping(axis="dim", local_start=0, external_start=0, length=config.head_dim), + AxisMapping(axis="head", local_start=0, external_start=0, length=1), + AxisMapping(axis="dim", local_start=0, external_start=0, length=config.latent_dim), ), - write_enabled=config.tp_rank % replica_count == 0, + write_enabled=config.tp_rank == 0, ), ) diff --git a/lightllm/common/attention_state/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla_state.py index 6e1e21fe8f..52edb0e1dd 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla_state.py @@ -12,7 +12,7 @@ from ..utils import deepcopy_result, dtype_name from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionConfig, FullAttentionState +from .base_state import FullAttentionState from .packed_mla_transfer import transfer_packed_mla_pages from .dp_transfer import copy_packed_tokens @@ -24,37 +24,60 @@ @dataclass(frozen=True) -class PackedMLAAttentionConfig(FullAttentionConfig): - # MLA KV latent 的有效维度，包含需要保存的 RoPE 部分，不含 indexer、scale 或行 padding。 - kv_dim: int = 576 - # 每个 token 的 indexer key 维度，与 MLA latent 一起存放在同一物理行。 +class PackedMLAAttentionConfig: + """group128/FP32 scale 或浮点 packed MLA，固定单份 latent 并在各 TP rank 复制。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 完整逻辑 latent 宽度，不含 indexer、scale 和物理 padding；FP8 为 576。 + latent_dim: int + # KV 编码精度：FP16、BF16 或 float8_e4m3fn；FP8 仅量化非 RoPE 部分。 + # packed FP8 的底层 buffer 使用 uint8 承载 KV、FP32 scale 和 BF16 RoPE，不作为配置输入。 + dtype: torch.dtype + # indexer key 的 FP8 维度，另保存一个 FP32 scale，不计入此维度。 indexer_dim: int = 128 - # 0 表示不使用 FP8 latent 编码；非零表示需要 FP8 量化的非 RoPE 部分维度。 - # 当前 FP8 实现要求 512 维非 RoPE latent，剩余 64 维 RoPE 以 BF16 保存。 + # 0 为未量化；512 表示前 512 维按 group128 编码，其余 64 维以 BF16 保存。 fp8_nope_dim: int = 0 - # FP8 非 RoPE latent 的分组量化宽度，当前 kernel 要求每 128 个元素对应一个 scale。 + # FP8 非 RoPE 特征的量化分组宽度，当前编码固定为 128。 quant_group_size: int = 128 + # 所属模型 TP 组内的 rank；各 rank 保存相同 latent，仅 rank 0 写出共享外部页。 + tp_rank: int = 0 + # 保存完整 latent 副本的 TP rank 数，不改变 latent 的维度或物理行宽。 + tp_world_size: int = 1 + + @property + def kv_bytes(self) -> int: + """单行 KV 区域字节数，FP8 包含 512 字节数据、16 字节 scales 和 64 维 BF16。""" + return 656 if self.dtype == torch.float8_e4m3fn else self.latent_dim * self.dtype.itemsize + + @property + def token_bytes(self) -> int: + """KV、indexer key/scale 及整行 16 字节对齐所需的物理字节数。""" + return (self.kv_bytes + self.indexer_dim + 4 + 15) // 16 * 16 + + @property + def indexer_offset(self) -> int: + """保持既有格式：FP8 紧接 KV，浮点格式将对齐间隙放在 KV 与 indexer 之间。""" + return self.kv_bytes if self.dtype == torch.float8_e4m3fn else self.token_bytes - self.indexer_dim - 4 def __post_init__(self) -> None: - """创建配置时校验通用 TP 约束、MLA 编码及 packed 行宽。""" - super().__post_init__() - assert self.head_num == 1 and self.indexer_dim > 0, ( - "packed MLA requires one replicated latent and a positive indexer dimension" - ) - if self.fp8_nope_dim: - assert (self.fp8_nope_dim, self.kv_dim, self.quant_group_size, self.dtype) == ( - 512, - 576, - 128, - torch.uint8, - ), "FP8 MLA kernel requires 512 latent dimensions, 64 RoPE dimensions and groups of 128" - payload_bytes = 656 + self.indexer_dim + 4 + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.latent_dim > 0 and self.indexer_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + if self.dtype == torch.float8_e4m3fn: + assert (self.fp8_nope_dim, self.latent_dim, self.quant_group_size) == (512, 576, 128), ( + "FP8 MLA kernel requires 512 latent dimensions, 64 RoPE dimensions and groups of 128" + ) else: - assert self.dtype in (torch.float16, torch.bfloat16), "unquantized packed MLA requires FP16 or BF16" - payload_bytes = self.kv_dim * self.dtype.itemsize + self.indexer_dim + 4 - assert self.head_dim * self.dtype.itemsize == (payload_bytes + 15) // 16 * 16, ( - "packed token width must include aligned KV and indexer storage" - ) + assert self.dtype in (torch.float16, torch.bfloat16) and self.fp8_nope_dim == 0, ( + "unquantized packed MLA requires FP16 or BF16 without FP8 dimensions" + ) @dataclass(frozen=True) @@ -79,7 +102,6 @@ def __init__( ) -> None: """按 packed MLA 配置准备 KV 与 indexer 存储。""" self.config = layout_config - self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化子类显式设置，区别于 scale 精度。 self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { @@ -94,34 +116,31 @@ def __init__( def get_state_type(self) -> str: """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" - encoding = f"fp8_group{self.config.quant_group_size}" if self.config.fp8_nope_dim else dtype_name(dtype=self.config.dtype) + encoding = ( + f"fp8_group{self.config.quant_group_size}" + if self.config.dtype == torch.float8_e4m3fn else dtype_name(dtype=self.config.dtype) + ) return "mla_packed_" + encoding def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity - shape = self._storage_shape() - self.storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) - - def get_token_bytes(self) -> int: - """由配置计算当前 TP rank 的单 token 开销，包含 packed 行 padding，无需分配存储。""" - config = self.config - return ( - len(config.logical_layer_groups) * config.head_num * config.head_dim - * self.storage_dtype.itemsize + # FP8 行混合保存多种字段，用字节 buffer 承载，字段视图再解释各自的 dtype。 + storage_dtype = torch.uint8 if self.config.dtype == torch.float8_e4m3fn else self.config.dtype + self.storage = torch.empty( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + 1, + self.config.token_bytes // storage_dtype.itemsize, + ), + dtype=storage_dtype, + device=self.device, ) - @property - def kv_nbytes(self) -> int: - """返回一行中 KV 区域的字节数，包含该编码需要的 scale 与 RoPE 数据。""" - return 656 if self.config.fp8_nope_dim else self.config.kv_dim * self.config.dtype.itemsize - - @property - def indexer_offset(self) -> int: - """返回一行中 indexer key 区域的字节偏移。""" - if self.config.fp8_nope_dim: - return self.kv_nbytes - return self.config.head_dim * self.config.dtype.itemsize - self.config.indexer_dim - 4 + def get_token_bytes(self) -> int: + """返回全部物理层的单 token 存储字节数，包含编码、indexer 和 padding。""" + return len(self.config.logical_layer_groups) * self.config.token_bytes # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 @@ -139,8 +158,11 @@ def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] raw = row.view(torch.uint8) - kv = raw[..., : self.kv_nbytes] if self.config.fp8_nope_dim else row[..., : self.config.kv_dim] - indexer = raw[..., self.indexer_offset : self.indexer_offset + self.config.indexer_dim + 4] + kv = ( + raw[..., : self.config.kv_bytes] + if self.config.dtype == torch.float8_e4m3fn else row[..., : self.config.latent_dim] + ) + indexer = raw[..., self.config.indexer_offset : self.config.indexer_offset + self.config.indexer_dim + 4] return PackedMLALayerView(kv=kv, indexer_k=indexer) def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: @@ -150,7 +172,7 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" - if self.config.fp8_nope_dim: + if self.config.dtype == torch.float8_e4m3fn: from lightllm.common.basemodel.triton_kernel.kv_copy.destindex_copy_kv_flashmla_fp8 import ( destindex_copy_kv_flashmla_fp8, ) @@ -167,7 +189,7 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, else: from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.kv_dim]) + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.latent_dim]) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" @@ -222,8 +244,8 @@ def write_to_cpu( page_buffers=cpu_page_buffers, page_indices=cpu_page_indices, layout=layout, - kv_dim=self.config.kv_dim, - indexer_offset=self.indexer_offset, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, indexer_dim=self.config.indexer_dim, write=True, page_readies=cpu_page_ready_flags, @@ -244,8 +266,8 @@ def read_from_cpu( page_buffers=cpu_page_buffers, page_indices=cpu_page_indices, layout=layout, - kv_dim=self.config.kv_dim, - indexer_offset=self.indexer_offset, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, indexer_dim=self.config.indexer_dim, write=False, ) @@ -263,11 +285,11 @@ def copy_tokens_from_dp( """校验 MLA 编码与区域尺寸，再复制完整 KV 和 indexer 数据。""" assert ( remote_dp_full_att_state.config.fp8_nope_dim, - remote_dp_full_att_state.config.kv_dim, + remote_dp_full_att_state.config.latent_dim, remote_dp_full_att_state.config.indexer_dim, ) == ( self.config.fp8_nope_dim, - self.config.kv_dim, + self.config.latent_dim, self.config.indexer_dim, ), "incompatible packed MLA format" assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( @@ -306,8 +328,8 @@ def write_to_pd_page( page_buffers=pd_page_buffers, page_indices=[pd_page_index], layout=layout, - kv_dim=self.config.kv_dim, - indexer_offset=self.indexer_offset, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, indexer_dim=self.config.indexer_dim, write=True, ) @@ -327,8 +349,8 @@ def read_from_pd_page( page_buffers=pd_page_buffers, page_indices=[pd_page_index], layout=layout, - kv_dim=self.config.kv_dim, - indexer_offset=self.indexer_offset, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, indexer_dim=self.config.indexer_dim, write=False, ) @@ -350,50 +372,36 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 - if config.fp8_nope_dim: + if config.dtype == torch.float8_e4m3fn: field_specs = [ ("kv_nope", torch.uint8, 512, "dim"), ("kv_scale", torch.float32, 4, "group"), ("kv_rope", torch.bfloat16, 64, "dim"), ] else: - field_specs = [("kv", config.dtype, config.kv_dim, "dim")] + field_specs = [("kv", config.dtype, config.latent_dim, "dim")] field_specs.extend(( ("indexer_k", torch.uint8, config.indexer_dim, "dim"), ("indexer_scale", torch.float32, 1, "scale"), )) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) - local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) - global_shape = list(local_shape) - global_shape[axes.index("token")] = page_tokens - head_axis = axes.index("head") - # field 已按 K/V 拆开：这里读取的是单个字段的本地 head 数，不是合并 buffer 的 2H。 - local_head_count = global_shape[head_axis] - # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 - # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 - # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 - global_shape[head_axis] = config.global_head_num - - # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 - # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 - replica_count = config.tp_world_size * local_head_count // config.global_head_num - # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 - head_start = config.tp_rank // replica_count * local_head_count - # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 - write_enabled = config.tp_rank % replica_count == 0 + local_shape = (len(config.logical_layer_groups), page_tokens, 1, field_width) + global_shape = local_shape + # 全部 TP ranks 保存同一 latent，只有 rank 0 写出共享外部页。 + write_enabled = config.tp_rank == 0 axis_mappings = tuple( AxisMapping( axis=axis, local_start=0, - external_start=head_start if axis == "head" else 0, + external_start=0, length=dimension, ) for axis, dimension in zip(axes, local_shape) if axis not in ("layer", "token") ) - # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + # 阶段二：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 element_bytes = field_dtype.itemsize byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes field_region = TransferRegion( @@ -408,17 +416,6 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ) regions.append(field_region) byte_offset = field_region.byte_end - logger.debug( - "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", - self.state_name, - field_name, - head_start, - local_head_count, - replica_count, - write_enabled, - field_region.byte_start, - field_region.byte_end, - ) logger.debug( "%s transfer regions built: fields=%d, total_bytes=%d", self.state_name, @@ -426,11 +423,3 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: byte_offset, ) return tuple(regions) - - def _storage_shape(self) -> tuple[int, ...]: - return ( - len(self.config.logical_layer_groups), - self.capacity + self.page_size, - self.config.head_num, - self.config.head_dim, - ) diff --git a/lightllm/common/attention_state/full_att/static_fp8_state.py b/lightllm/common/attention_state/full_att/static_fp8_state.py index 458ebae0c0..b874cc421d 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/static_fp8_state.py @@ -14,7 +14,7 @@ from ..utils import deepcopy_result from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionConfig, FullAttentionState +from .base_state import FullAttentionState from .static_fp8_transfer import transfer_static_fp8_pages from .dp_transfer import copy_gqa_tokens @@ -26,14 +26,44 @@ @dataclass(frozen=True) -class StaticFP8AttentionConfig(FullAttentionConfig): - # per_tensor：每层 K/V 各一个固定 scale；per_head：每层、每个本地 K/V head 各一个。 - # 标定参数不随 token 增长，也不作为动态 KV 字段写入 CPU/PD 页。 +class StaticFP8AttentionConfig: + """静态 FP8 GQA 的逻辑布局；KV 固定为 FP8 编码，标定 Tensor 由后续流程绑定。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + # 每层 K/V 各一个 scale，或每个本地 K/V head 各一个；不占用逐 token 页空间。 scale_mode: str = "per_tensor" + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + def __post_init__(self) -> None: - """创建配置时校验通用 TP 约束及标定方式。""" - super().__post_init__() + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" assert self.scale_mode in ("per_tensor", "per_head"), "invalid static FP8 scale mode" diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 0767c9f652..5cf512e10c 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -37,7 +37,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -45,7 +45,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 21593fe97d..047c619dfa 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -69,7 +69,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -77,10 +77,9 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_mla_state( - config=FullAttentionConfig( + config=MLAAttentionConfig( logical_layer_groups=logical_layer_groups, - global_head_num=1, - head_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, ), diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index e899db92f4..105d735816 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -25,7 +25,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -33,10 +33,9 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_mla_state( - config=FullAttentionConfig( + config=MLAAttentionConfig( logical_layer_groups=logical_layer_groups, - global_head_num=1, - head_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, ), diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index 17f6246153..15dc1f37f0 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -149,14 +149,14 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=256, diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index e586d33fad..dd87a1a505 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -96,7 +96,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -124,7 +124,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=per_token_k_width // head_dim, head_dim=head_dim, diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index 2e95f10d14..3abaf00466 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -47,7 +47,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -55,7 +55,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config["hidden_size"] // config["num_attention_heads"], diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 3146ba3b30..61aa9602c4 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -116,7 +116,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -125,10 +125,9 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_mla_state( - config=FullAttentionConfig( + config=MLAAttentionConfig( logical_layer_groups=logical_layer_groups, - global_head_num=1, - head_dim=cfg.full_att_head_dim, + latent_dim=cfg.full_att_head_dim, dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 57df11fc88..1d2ac2d764 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -81,7 +81,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -89,10 +89,9 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_mla_state( - config=FullAttentionConfig( + config=MLAAttentionConfig( logical_layer_groups=logical_layer_groups, - global_head_num=1, - head_dim=config["kv_lora_rank"], + latent_dim=config["kv_lora_rank"], dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 2a33eda402..c89567b3f1 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -45,7 +45,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -53,7 +53,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index b4dea62592..26e8fd1ffe 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -39,7 +39,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -47,7 +47,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index b90658e289..e2c298967c 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -47,7 +47,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -55,7 +55,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 36460f6d28..0604d063e5 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -107,7 +107,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -117,7 +117,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=cfg.full_att_all_num_kv_heads, head_dim=cfg.full_att_head_dim, diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index 49d295e298..5ca086b03e 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -35,7 +35,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -43,7 +43,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=1, head_dim=config["hidden_size"] // config["num_attention_heads"], diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 35ed6d35e2..66d5ebb4a5 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -38,7 +38,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -46,7 +46,7 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=logical_layer_groups, global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py index dc03199294..312d524d50 100644 --- a/lightllm/server/router/model_infer/shared_att_resources_utils.py +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -39,18 +39,17 @@ def publish_att_resources( else: ipc_resources = resources.create_dp_ipc_copy() with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): - # Each importer consumes its own CUDA IPC reference counters. + # 每个读者独立序列化 CUDA IPC 引用；计数偏移变化可能使各份 pickle 长度不同。 payloads = [ForkingPickler.dumps((ipc_resources, request_mapping)).tobytes() for _ in range(count)] - size = len(payloads[0]) - if any(len(payload) != size for payload in payloads): - raise RuntimeError("rank storage exports have inconsistent sizes") + size = max(len(payload) for payload in payloads) shm = create_or_link_shm( name=f"rank_storage_{kind}_{get_current_rank_in_node()}", expected_size=8 + count * size, force_mode="create" ) shm.buf[:4] = count.to_bytes(4, "little") shm.buf[4:8] = size.to_bytes(4, "little") for index, payload in enumerate(payloads): - shm.buf[8 + index * size : 8 + (index + 1) * size] = payload + # 每个读者占用等长槽位；pickle.loads 会忽略完整对象之后的补零。 + shm.buf[8 + index * size : 8 + (index + 1) * size] = payload.ljust(size, b"\0") shm.close() diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index db67d9aebb..8236c5ff36 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -14,7 +14,7 @@ from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, - FullAttentionConfig, + GQAAttentionConfig, ) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.common.attention_state import AttentionStateManager @@ -50,7 +50,7 @@ def main(): ) tokens = args.tokens_per_page * args.pages full = GQAFullAttentionState( - FullAttentionConfig( + GQAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in tuple(range(12))), global_head_num=1, head_dim=256, diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index a252410165..950f7cde10 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -7,7 +7,7 @@ from functools import lru_cache from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.full_att import MLAFullAttentionState +from lightllm.common.attention_state.full_att import MLAAttentionConfig, MLAFullAttentionState from lightllm.common.attention_state.bounded_att import GDNState from lightllm.common.attention_state.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state @@ -32,7 +32,11 @@ def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = T self.register_state("draft0", draft_state) self.register_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( - self.main_gqa_fp32_state.config, + layout_config=MLAAttentionConfig( + logical_layer_groups=self.main_gqa_fp32_state.config.logical_layer_groups, + latent_dim=self.main_gqa_fp32_state.config.head_dim, + dtype=self.main_gqa_fp32_state.config.dtype, + ), page_size=self.main_gqa_fp32_state.page_size, device=device, ) prepared_full_state.allocate_storage(self.main_gqa_fp32_state.capacity) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 54c19df896..78ae9342d5 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -10,7 +10,9 @@ import torch from lightllm.common.attention_state.types import TensorField from lightllm.common.attention_state.full_att import ( - FullAttentionState, FullAttentionConfig, StaticFP8FullAttentionState, + FullAttentionState, GQAAttentionConfig, MLAAttentionConfig, + GroupQuantizedAttentionConfig, StaticFP8AttentionConfig, PackedMLAAttentionConfig, + CompressedAttentionConfig, FP8E8M0PackedMLAConfig, StaticFP8FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, ) @@ -52,7 +54,11 @@ def bind_test_fp8_calibration( def rebuild_full_state( template: FullAttentionState, - config: FullAttentionConfig | None = None, + config: ( + GQAAttentionConfig | MLAAttentionConfig | GroupQuantizedAttentionConfig + | StaticFP8AttentionConfig | PackedMLAAttentionConfig + | CompressedAttentionConfig | FP8E8M0PackedMLAConfig | None + ) = None, token_capacity: int | None = None, page_size: int | None = None, device: str | torch.device | None = None, @@ -167,15 +173,15 @@ def full_state_test_fields( if isinstance(state, PackedMLAFullAttentionState): storage = buffers["kv"] raw = storage.view(torch.uint8) - if state.config.fp8_nope_dim: + if state.config.dtype == torch.float8_e4m3fn: fields = { "kv_nope": TensorField(tensor=raw[..., :512], axes=axes), "kv_scale": TensorField(tensor=raw[..., 512:528].view(torch.float32), axes=("layer", "token", "head", "group")), "kv_rope": TensorField(tensor=raw[..., 528:656].view(torch.bfloat16), axes=axes), } else: - fields = {"kv": TensorField(tensor=storage[..., :state.config.kv_dim], axes=axes)} - start = state.indexer_offset + fields = {"kv": TensorField(tensor=storage[..., :state.config.latent_dim], axes=axes)} + start = state.config.indexer_offset end = start + state.config.indexer_dim fields["indexer_k"] = TensorField(tensor=raw[..., start:end], axes=axes) fields["indexer_scale"] = TensorField(tensor=raw[..., end:end + 4].view(torch.float32), axes=("layer", "token", "head", "scale")) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index b6679a9835..fe3d1c8037 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -27,13 +27,14 @@ from lightllm.common.attention_state.full_att import ( FullAttentionState, GQAFullAttentionState, - FullAttentionConfig, + GQAAttentionConfig, + MLAAttentionConfig, + create_gqa_state, + create_mla_state, GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, MLAFullAttentionState, - PackedMLAAttentionConfig, PackedMLAFullAttentionState, - StaticFP8AttentionConfig, StaticFP8FullAttentionState, ) @@ -42,45 +43,40 @@ def definition(kind: str): - config = FullAttentionConfig( - logical_layer_groups=((0,), (2,)), - global_head_num=2, - head_dim=16, - dtype=torch.bfloat16, - ) - if kind in ("int4", "int8"): - return GroupQuantizedFullAttentionState( - GroupQuantizedAttentionConfig(**vars(config), bits=4 if kind == "int4" else 8), - page_size=4, device="meta", + if kind == "mla" or kind.startswith("packed"): + return create_mla_state( + config=MLAAttentionConfig( + logical_layer_groups=((0,), (2,)), + latent_dim=512 if kind == "packed_glm" else 576, + dtype=torch.bfloat16, + ), + kv_type="fp8kv_dsa" if kind == "packed_fp8" else "None", + packed=kind.startswith("packed"), + page_size=4, + device="meta", ) + kv_type = { + "int4": "int4kv", "int8": "int8kv", + "static_head": "fp8kv_sph", "static_tensor": "fp8kv_spt", + }.get(kind, "None") + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + ), + kv_type=kv_type, + page_size=4, + device="meta", + ) if kind.startswith("static"): heads = 2 if kind == "static_head" else 1 - state = StaticFP8FullAttentionState( - StaticFP8AttentionConfig( - **vars(config), scale_mode="per_head" if heads == 2 else "per_tensor" - ), - page_size=4, device="meta", - ) bind_test_fp8_calibration( state=state, scales=torch.ones((2, heads * 2), dtype=torch.float32), q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, ) - return state - if kind.startswith("packed"): - fp8 = kind == "packed_fp8" - kv_dim = 512 if kind == "packed_glm" else 576 - config = replace( - config, - global_head_num=1, - head_dim=800 if fp8 else kv_dim + 72, - dtype=torch.uint8 if fp8 else config.dtype, - ) - return PackedMLAFullAttentionState( - PackedMLAAttentionConfig(**vars(config), kv_dim=kv_dim, fp8_nope_dim=512 if fp8 else 0), page_size=4, device="meta" - ) - if kind == "mla": - return MLAFullAttentionState(replace(config, global_head_num=1, head_dim=576), page_size=4, device="meta") - return GQAFullAttentionState(config, page_size=4, device="meta") + return state def allocated_bytes(state: FullAttentionState, token_only: bool = False): @@ -294,12 +290,14 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( (GQAFullAttentionState, "main", 3, torch.bfloat16), (GQAFullAttentionState, "main", 1, torch.float32), ): - config = FullAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=dim, - dtype=dtype, - ) + if state_cls is MLAFullAttentionState: + config = MLAAttentionConfig( + logical_layer_groups=((0,),), latent_dim=dim, dtype=dtype, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((0,),), global_head_num=1, head_dim=dim, dtype=dtype, + ) manager.register_state(role, state_cls(config, page_size=1, device="meta")) manager.finish_state_registration() # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 @@ -369,7 +367,7 @@ def _get_transfer_regions(self, page_tokens: int): logical_layer_groups=((0,),), global_head_num=1, head_dim=128, - dtype=torch.float32, + scale_dtype=torch.float32, bits=8, group_size=128, ) @@ -464,7 +462,7 @@ def get_pd_transfer_regions(self): manager = AttentionStateManager(request_capacity=1) if full_page_tokens is not None: pd_args.pd_kv_page_size = full_page_tokens - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=1, @@ -551,7 +549,7 @@ def _get_transfer_regions(self, page_tokens: int): pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 manager = AttentionStateManager(request_capacity=1) - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=3, @@ -583,7 +581,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( pd_args: EasyDict, ): manager = AttentionStateManager(request_capacity=1) - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=1, @@ -866,7 +864,7 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation pd_args.pd_kv_page_size = 3 pd_args.cpu_cache_token_page_size = 3 manager = AttentionStateManager(request_capacity=1) - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=3, @@ -935,7 +933,7 @@ def get_cpu_transfer_regions(self): ) pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=2, diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index 3a70b0fd58..031bc33f40 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -49,7 +49,7 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, infer.att_state_role = role infer.mem_index = torch.tensor([1, 5], dtype=torch.int32, device="cuda") heads = 1 if kind == "packed_glm" else 2 * state.config.head_num - dim = state.config.kv_dim if kind == "packed_glm" else state.config.head_dim + dim = state.config.latent_dim if kind == "packed_glm" else state.config.head_dim kv = torch.full( size=(2, heads, dim), fill_value=0.5, diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py index 87e0133801..169ba44485 100644 --- a/unit_tests/common/att_resources/test_dsv4_state.py +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -98,10 +98,10 @@ def test_compressed_meta_geometry() -> None: def test_compressed_factory_uses_independent_config() -> None: - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_compressed_att_state + from lightllm.common.attention_state.full_att import create_compressed_att_state definition = make_full(ratio=4) - assert not isinstance(definition.config, FullAttentionConfig) + assert type(definition.config).__bases__ == (object,) state = create_compressed_att_state(config=definition.config, page_size=256, device="meta") assert isinstance(state, CompressedAttentionState) assert state.storage is None and state.get_token_bytes() == 180 diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 6f28089874..8c79ea9d36 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -151,7 +151,6 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu logical_layer_groups=((5,),), global_head_num=head_num, head_dim=128, - dtype=torch.bfloat16, scale_mode=scale_mode, ), page_size=4, device="cuda", diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py index 8327c36bf1..36210e397d 100644 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -12,7 +12,7 @@ from lightllm.common.attention_state.full_att import ( FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState, - FullAttentionConfig, + MLAAttentionConfig, FullAttentionState, PackedMLAFullAttentionState, create_mla_state, @@ -26,10 +26,9 @@ def make_state(device: str = "meta", capacity: int = 16, tp_rank: int = 0) -> FP8E8M0PackedMLAState: state = create_mla_state( - config=FullAttentionConfig( + config=MLAAttentionConfig( logical_layer_groups=((0, 4), (8,)), - global_head_num=1, - head_dim=512, + latent_dim=512, dtype=torch.bfloat16, tp_rank=tp_rank, tp_world_size=2, @@ -51,7 +50,7 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic assert FP8E8M0PackedMLAState.__bases__ == (FullAttentionState,) assert not isinstance(state, PackedMLAFullAttentionState) assert isinstance(state.config, FP8E8M0PackedMLAConfig) - assert not isinstance(state.config, FullAttentionConfig) + assert type(state.config).__bases__ == (object,) assert not hasattr(state.config, "head_dim") assert state.get_state_type() == "mla_packed_fp8_group64_e8m0" assert state.storage.shape == (2, 20, 1, 720) and state.storage.is_meta @@ -65,8 +64,8 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic # 默认仍是原来的 group128/FP32 编码，不能再通过 512 维自动猜测 GLM 格式。 with pytest.raises(AssertionError, match="FP8 MLA kernel requires"): create_mla_state( - config=FullAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=512, dtype=torch.bfloat16 + config=MLAAttentionConfig( + logical_layer_groups=((0,),), latent_dim=512, dtype=torch.bfloat16 ), kv_type="fp8kv_dsa", packed=True, @@ -74,8 +73,8 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic ) with pytest.raises(AssertionError, match="512-dimensional"): create_mla_state( - config=FullAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=576, dtype=torch.bfloat16 + config=MLAAttentionConfig( + logical_layer_groups=((0,),), latent_dim=576, dtype=torch.bfloat16 ), kv_type="fp8kv_dsa", packed=True, diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index c6f46ed9fe..ffddf2f375 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -11,7 +11,7 @@ from unit_tests.common.att_resources.test_state_identity import allocated_state from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout from lightllm.common.attention_state.types import AxisMapping -from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -19,7 +19,7 @@ def test_narrow_head_fields_merge_before_page_transfer(page_device: str, monkeypatch: pytest.MonkeyPatch) -> None: from lightllm.common.attention_state.full_att import gqa_transfer - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,), (1,)), global_head_num=4, head_dim=1, diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 787025b01e..5e5f768f7d 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -16,7 +16,8 @@ FullAttentionState, GQAFullAttentionState, MLAFullAttentionState, - FullAttentionConfig, + GQAAttentionConfig, + MLAAttentionConfig, ) from lightllm.common.attention_state.types import ( AxisMapping, StateTransferLayout, TransferRegion, @@ -24,12 +25,15 @@ def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): - prepared_full_state = cls(FullAttentionConfig( - logical_layer_groups=((2,), (5,)), - global_head_num=1, - head_dim=4, - dtype=torch.float32, - ), page_size=4, device=device) + if cls is MLAFullAttentionState: + config = MLAAttentionConfig( + logical_layer_groups=((2,), (5,)), latent_dim=4, dtype=torch.float32, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((2,), (5,)), global_head_num=1, head_dim=4, dtype=torch.float32, + ) + prepared_full_state = cls(config, page_size=4, device=device) prepared_full_state.allocate_storage(12) return prepared_full_state @@ -223,7 +227,7 @@ def test_invalid_late_region_shape_cannot_partially_write_pd_page(): def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((1,), (3,), (5,), (7,), (9,), (11,)), global_head_num=1, head_dim=4, diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index 32a1849db8..a239b823df 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -16,7 +16,7 @@ import pytest import torch -from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState +from lightllm.common.attention_state.full_att import MLAFullAttentionState, PackedMLAFullAttentionState from lightllm.utils.envs_utils import get_env_start_args from unit_tests.common.att_resources.test_capacity import allocated_bytes from unit_tests.server.router.model_infer.test_cpu_resource_layout import CASES, make_runtime, runtime @@ -51,9 +51,10 @@ def test_model_storage_write_and_copy_preserve_all_fields( for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") - dim = config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else config.head_dim + is_mla = isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) + dim = config.latent_dim if is_mla else config.head_dim kv = torch.randn( - size=(2, (1 if "kv" in full_state_test_buffers(state) else 2) * config.head_num, dim), + size=(2, 1 if is_mla else 2 * config.head_num, dim), device="cuda", dtype=torch.bfloat16, ) diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py index ffc4607ee0..1f64107157 100644 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -8,7 +8,7 @@ import pytest import torch -from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState from lightllm.common.attention_state.types import StateTransferLayout if TYPE_CHECKING: @@ -30,7 +30,7 @@ def test_gqa_pd_direct_copy_between_tp_layouts( pd_args.pd_kv_page_size = 4 # P 保存完整 heads，D 只恢复 rank 1 的一半 heads；两侧外部页格式相同。 source = GQAFullAttentionState( - layout_config=FullAttentionConfig( + layout_config=GQAAttentionConfig( logical_layer_groups=((2,), (5,)), global_head_num=4, head_dim=128, @@ -137,7 +137,7 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: pd_args.pd_kv_page_size = 4 # MQA 的 KV head 在 TP ranks 间复制，非首 rank 禁止写出但仍能读回。 state = GQAFullAttentionState( - layout_config=FullAttentionConfig( + layout_config=GQAAttentionConfig( logical_layer_groups=((0,), (2,)), global_head_num=1, head_dim=128, @@ -194,7 +194,7 @@ def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pyt pytest.skip("CUDA peer access required") pd_args.pd_kv_page_size = 4 state = GQAFullAttentionState( - layout_config=FullAttentionConfig( + layout_config=GQAAttentionConfig( logical_layer_groups=((0,), (2,)), global_head_num=2, head_dim=128, diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index 8c0441f70a..0df056ad4e 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -5,7 +5,7 @@ import pytest import torch -from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState from lightllm.common.attention_state.full_att.gqa_transfer import transfer_gqa_pages from lightllm.common.attention_state.types import StateTransferLayout @@ -18,7 +18,7 @@ @pytest.mark.parametrize("page_device", ["cpu", "cuda"]) def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_args: "EasyDict") -> None: pd_args.pd_kv_page_size = 32 - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,), (2,), (4,)), global_head_num=3, head_dim=128, @@ -110,7 +110,7 @@ def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( from lightllm.common.attention_state.full_att import gqa_transfer pd_args.pd_kv_page_size = 2 - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,), (2,)), global_head_num=2, head_dim=128, diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py new file mode 100644 index 0000000000..d9ba1052ce --- /dev/null +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -0,0 +1,112 @@ +import pytest + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("draft_layers", [0, 1, 3]) +@pytest.mark.parametrize( + "layer_types,expected_indices,linear_layers", + [ + ( + ("linear_attention", "linear_attention", "linear_attention", "full_attention") * 2, + (0, 1, 2, 0, 3, 4, 5, 1), + 6, + ), + ( + ( + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + ), + (0, 0, 1, 1, 2), + 3, + ), + ( + ("linear_attention", "full_attention", "linear_attention") * 2 + ("full_attention",), + (0, 0, 1, 2, 1, 3, 2), + 4, + ), + ], +) +def test_hybrid_layer_mapping_and_mtp_states( + layer_types: tuple[str, ...], + expected_indices: tuple[int, ...], + linear_layers: int, + draft_layers: int, +) -> None: + model_config = dict( + num_hidden_layers=len(layer_types), + hidden_size=128, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_num_key_heads=1, + linear_num_value_heads=1, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + full_attention_interval=4, + layer_types=layer_types, + ) + args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + config = Qwen3NextResourceConfig.from_model_config(model_config, args) + logical_layer_ids = tuple(range(len(layer_types))) + (full_state,) = Qwen3NextTpPartModel.create_full_att_states( + model_config=model_config, args=args, model_logical_layer_ids=logical_layer_ids, + ) + (bounded_state,) = Qwen3NextTpPartModel.create_bounded_att_states( + model_config=model_config, args=args, model_logical_layer_ids=logical_layer_ids, device="meta", + ) + manager = AttentionStateManager(request_capacity=1) + manager.register_state(model_role="main", state=full_state) + manager.register_state(model_role="main", state=bounded_state) + # draft 各自创建独立 State，不再把其层数和物理行追加到 main 的配置中。 + for draft_index in range(draft_layers): + (state,) = Qwen3_5MTPModel.create_full_att_states( + model_config=model_config, + args=args, + model_logical_layer_ids=(len(layer_types) + draft_index,), + ) + manager.register_state(model_role=f"draft{draft_index}", state=state) + manager.finish_state_registration() + + assert config.all_layer_num == len(layer_types) + full_layers = len(layer_types) - linear_layers + assert config.get_full_att_layer_indices() == tuple( + layer for layer, kind in enumerate(layer_types) if kind != "linear_attention" + ) + assert len(full_state.config.logical_layer_groups) == full_layers + assert bounded_state.conv_state.shape[0] == bounded_state.ssm_state.shape[0] == linear_layers + assert sum(len(state.config.logical_layer_groups) for state in manager.get_full_attention_states()) == ( + full_layers + draft_layers + ) + for layer, (kind, expected) in enumerate(zip(layer_types, expected_indices)): + is_linear = kind == "linear_attention" + state = manager.get_state( + model_role="main", + state_type_filter="", + logical_layer_id=layer, + state_class=BoundedAttentionState if is_linear else FullAttentionState, + ) + if is_linear: + conv, ssm = state.get_layer_view(logical_layer_id=layer) + assert conv.storage_offset() == expected * state.conv_state.stride(0) + assert ssm.storage_offset() == expected * state.ssm_state.stride(0) + else: + assert state.get_storage_layer_index(logical_layer_id=layer) == expected + for draft in range(draft_layers): + state = manager.get_state( + model_role=f"draft{draft}", + state_type_filter="gqa", + logical_layer_id=len(layer_types) + draft, + ) + assert state is not full_state + assert state.get_storage_layer_index(logical_layer_id=len(layer_types) + draft) == 0 diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index c7d98e1437..6046e1475f 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -8,7 +8,6 @@ from torch.utils._pytree import tree_leaves from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import FullAttentionConfig from lightllm.models.gemma4.model import Gemma4TpPartModel from unit_tests.common.att_resources.model_helpers import ( full_state_test_buffers, @@ -57,15 +56,11 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - assert [state.get_storage_layer_index(logical_layer_id=layer) for layer in (0, 4, 8, 11)] == [0, 0, 1, 1] +@pytest.mark.parametrize("kind", SCHEMES) @pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) -def test_invalid_layer_groups_are_rejected(groups: tuple) -> None: +def test_invalid_layer_groups_are_rejected(kind: str, groups: tuple) -> None: with pytest.raises(AssertionError): - FullAttentionConfig( - logical_layer_groups=groups, - global_head_num=1, - head_dim=128, - dtype=torch.bfloat16, - ) + replace(definition(kind).config, logical_layer_groups=groups) def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "EasyDict") -> None: diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index b3ea892073..f153ffd96e 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -7,7 +7,7 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import FullAttentionConfig, FullAttentionState, GQAFullAttentionState +from lightllm.common.attention_state.full_att import GQAAttentionConfig, FullAttentionState, GQAFullAttentionState from lightllm.common.attention_state.bounded_att import KDAState, GDNState from unit_tests.common.att_resources.test_bounded_state import make_state from unit_tests.common.att_resources.test_tail_state import make_tail @@ -25,7 +25,7 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: model.max_req_num = 0 model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) state = GQAFullAttentionState( - layout_config=FullAttentionConfig( + layout_config=GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=128, @@ -49,7 +49,7 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo for model_role in ("main", "draft0"): for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): manager.register_state(model_role, GQAFullAttentionState( - layout_config=FullAttentionConfig( + layout_config=GQAAttentionConfig( logical_layer_groups=groups, global_head_num=1, head_dim=128, diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index 5439dd3bb2..7866525860 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -134,7 +134,7 @@ def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, tp_world_size: int, ): from builtins import open as original_open - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state # 第 0 行属于普通精度的共享层；其余行分别属于 main、draft2、draft10 的 FP8 State。 scales = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) @@ -168,7 +168,7 @@ def count_open(file: str, *args: object, **kwargs: object): ("main", ((0, 1),), 0, "None"), ): state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=groups, global_head_num=4, head_dim=16, @@ -226,10 +226,10 @@ def test_cpu_layout_does_not_require_calibration_file( @pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state state = create_gqa_state( - config=FullAttentionConfig( + config=GQAAttentionConfig( logical_layer_groups=((0, 1), (2,)), global_head_num=2, head_dim=16, diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 57b521a440..a74b71fcc3 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -17,7 +17,7 @@ from unit_tests.common.att_resources.model_helpers import create_full_attention_roles from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, - FullAttentionConfig, + GQAAttentionConfig, ) from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig from lightllm.models.qwen3next.config import Qwen3NextResourceConfig @@ -59,19 +59,17 @@ def make_resources( full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), - global_head_num=1, - head_dim=584, dtype=torch.bfloat16, tp_rank=rank, tp_world_size=tp, - kv_dim=512, + latent_dim=512, ), page_size=4, ) full.allocate_storage(16) else: full = GQAFullAttentionState( - FullAttentionConfig( + GQAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in (1, 3, *range(4, 4 + draft_layers))), global_head_num=2, head_dim=16, diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 21b5773a47..9345913fc7 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -41,7 +41,7 @@ def make_quant_state(kind: str, device: str | torch.device | int): logical_layer_groups=((1,), (3,)), global_head_num=2, head_dim=16, - dtype=torch.float16, + scale_dtype=torch.float16, bits=int(kind[-1]), ) prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) @@ -53,7 +53,6 @@ def make_quant_state(kind: str, device: str | torch.device | int): logical_layer_groups=((1,), (3,)), global_head_num=2, head_dim=16, - dtype=torch.uint8, scale_mode="per_head" if per_head else "per_tensor", ) scales = torch.arange( @@ -71,10 +70,9 @@ def make_quant_state(kind: str, device: str | torch.device | int): return prepared_full_state fp8 = kind == "packed_fp8" config = PackedMLAAttentionConfig( + latent_dim=576, logical_layer_groups=((1,), (3,)), - global_head_num=1, - head_dim=800 if fp8 else 648, - dtype=torch.uint8 if fp8 else torch.bfloat16, + dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, fp8_nope_dim=512 if fp8 else 0, ) prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) @@ -261,7 +259,7 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): storage_buffer.zero_() if isinstance(state, PackedMLAFullAttentionState): state.get_layer_view(1).indexer_k.fill_(29) - dims = state.config.kv_dim if isinstance(state, PackedMLAFullAttentionState) else state.config.head_dim + dims = state.config.latent_dim if isinstance(state, PackedMLAFullAttentionState) else state.config.head_dim heads = 1 if isinstance(state, PackedMLAFullAttentionState) else 2 * state.config.head_num kv = torch.randn((2, heads, dims), dtype=torch.bfloat16, device="cuda") ids = torch.tensor([1, 6], dtype=torch.int32, device="cuda") @@ -289,7 +287,7 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): else: view = state.get_layer_view(1) assert (view.indexer_k == 29).all() - if not state.config.fp8_nope_dim: + if state.config.dtype != torch.float8_e4m3fn: assert torch.equal(view.kv[ids], kv) else: fields = full_state_test_fields(state) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py index 896a6073db..658c340281 100644 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -16,7 +16,7 @@ def test_quantized_kv_buffers_include_the_full_hold_page(): logical_layer_groups=((0,), (1,), (2,)), global_head_num=2, head_dim=16, - dtype=torch.float16, + scale_dtype=torch.float16, bits=bits, ), page_size=4, device="cpu", diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index de556140f8..9bbbb9fec0 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -9,6 +9,8 @@ from typing import TYPE_CHECKING from unit_tests.common.att_resources.model_helpers import main_linear_state from types import SimpleNamespace +from contextlib import nullcontext +import pickle import traceback import uuid @@ -455,6 +457,51 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe assert main_linear_state(manager).mtp_state_index[2] == 0 +@pytest.mark.parametrize("kind", ["pd", "dp"]) +def test_ipc_readers_accept_different_pickle_lengths(kind: str, monkeypatch: pytest.MonkeyPatch) -> None: + """各读者独立导出；模拟 IPC 计数变化引起的长度差异，并验证补齐后仍可逐份读取。""" + resources = SimpleNamespace( + create_pd_ipc_copy=lambda: {"kind": "pd"}, + create_dp_ipc_copy=lambda: {"kind": "dp"}, + ) + mapping = [[0, 1, 2]] + payloads = [] + storage = {} + + def serialize(value: tuple[object, object]) -> memoryview: + payload = pickle.dumps(value) + b"\0" * (3 * len(payloads)) + payloads.append(payload) + return memoryview(payload) + + def shared_memory(name: str, expected_size: int, force_mode: str) -> SimpleNamespace: + if force_mode == "create": + storage[name] = SimpleNamespace(buf=bytearray(expected_size), close=lambda: None) + return storage[name] + + monkeypatch.setattr(shared_att_resources_utils, "kv_trans_use_p2p", lambda: False) + monkeypatch.setattr(shared_att_resources_utils, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(shared_att_resources_utils, "get_service_shm_name", lambda name: name) + monkeypatch.setattr(shared_att_resources_utils, "FileLock", lambda path: nullcontext()) + monkeypatch.setattr(shared_att_resources_utils, "create_or_link_shm", shared_memory) + monkeypatch.setattr(shared_att_resources_utils.ForkingPickler, "dumps", serialize) + + shared_att_resources_utils.publish_att_resources( + resources=resources, request_mapping=mapping, kind=kind, reader_count=3, + ) + assert len(payloads) == 3 and len({len(payload) for payload in payloads}) == 3 + buffer = storage[f"rank_storage_{kind}_0"].buf + slot_bytes = int.from_bytes(buffer[4:8], "little") + assert slot_bytes == max(len(payload) for payload in payloads) + for index, payload in enumerate(payloads): + assert buffer[8 + index * slot_bytes : 8 + (index + 1) * slot_bytes] == payload.ljust(slot_bytes, b"\0") + for _ in payloads: + imported_resources, imported_mapping = shared_att_resources_utils.load_att_resources(rank_in_node=0, kind=kind) + assert imported_resources == {"kind": kind} + assert imported_mapping == mapping + with pytest.raises(RuntimeError, match="exhausted"): + shared_att_resources_utils.load_att_resources(rank_in_node=0, kind=kind) + + def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: try: torch.cuda.set_device(0) diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index acc8e837ac..f1eb62ec87 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -10,7 +10,7 @@ import torch from lightllm.common.attention_state.full_att import ( - FullAttentionConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, + GQAAttentionConfig, MLAAttentionConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, StaticFP8FullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, create_gqa_state, create_mla_state, ) @@ -36,7 +36,7 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( storage_dtype: torch.dtype, width: int, ): - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((3,), (7,)), global_head_num=8, head_dim=128, @@ -66,6 +66,11 @@ def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( assert torch.equal(state.scales, scales) if q_scales is not None: assert torch.equal(state.q_scales, q_scales) + assert type(state.config).__bases__ == (object,) + if kv_type in ("int4kv", "int8kv"): + assert state.config.scale_dtype == config.dtype + elif kv_type.startswith("fp8"): + assert not hasattr(state.config, "dtype") state.device = torch.device("cpu") state.allocate_storage(8) assert state.k_storage.shape == state.v_storage.shape == (2, 12, 2, width) @@ -92,10 +97,9 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( row_width: int, expected_type: str, ): - config = FullAttentionConfig( + config = MLAAttentionConfig( logical_layer_groups=((2,), (6,)), - global_head_num=1, - head_dim=kv_dim, + latent_dim=kv_dim, dtype=torch.bfloat16, tp_rank=1, tp_world_size=4, @@ -108,14 +112,21 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( device="meta", ) assert state.get_state_type() == expected_type - assert state.config.replicated and not config.replicated - assert state.config.head_dim == row_width and config.head_dim == kv_dim + assert state.config.latent_dim == config.latent_dim == kv_dim + assert state.config.dtype == (torch.float8_e4m3fn if kv_type == "fp8kv_dsa" else config.dtype) + assert not hasattr(state.config, "storage_dtype") + assert not hasattr(state.config, "replicated") assert state.capacity == 0 and read_kv_storage(state) is None + assert type(state.config).__bases__ == (object,) + assert not hasattr(state.config, "head_num") and not hasattr(state.config, "global_head_num") state.device = torch.device("cpu") state.allocate_storage(8) assert read_kv_storage(state).shape == (2, 12, 1, row_width) + assert read_kv_storage(state).dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else config.dtype) if packed: - assert state.config.kv_dim == kv_dim + # 保持旧算子读取的物理布局：FP8 indexer 紧接 KV，浮点格式保留中间的对齐间隙。 + expected_indexer_offset = 656 if kv_type == "fp8kv_dsa" else kv_dim * 2 + 12 + assert state.config.indexer_offset == expected_indexer_offset assert state.get_layer_view(2).indexer_k.shape[0] == 12 @@ -163,7 +174,7 @@ def test_kpool_creation_preserves_pool_and_mtp_dimensions(): def test_create_rejects_encoding_from_another_logical_family(): - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=1, head_dim=576, @@ -171,6 +182,9 @@ def test_create_rejects_encoding_from_another_logical_family(): ) with pytest.raises(AssertionError, match="unsupported GQA encoding"): create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") + config = MLAAttentionConfig( + logical_layer_groups=((0,),), latent_dim=576, dtype=torch.bfloat16, + ) with pytest.raises(AssertionError, match="unsupported packed MLA encoding"): create_mla_state( config=config, diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index 8919a81d1d..ebf7ef10c1 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -18,7 +18,8 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.attention_state.full_att import ( - FullAttentionConfig, + GQAAttentionConfig, + MLAAttentionConfig, GQAFullAttentionState, MLAFullAttentionState, ) @@ -57,13 +58,18 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( if not torch.cuda.is_available(): pytest.skip("CUDA required") cpu_args.cpu_cache_token_page_size = 4 - config = FullAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=heads, - head_dim=4, - dtype=torch.float32, - replicated=replicated, - ) + if state_cls is MLAFullAttentionState: + config = MLAAttentionConfig( + logical_layer_groups=((0,),), latent_dim=4, dtype=torch.float32, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=heads, + head_dim=4, + dtype=torch.float32, + replicated=replicated, + ) reference = state_cls(config, page_size=4, device="cuda") reference.allocate_storage(4) for storage_buffer in full_state_test_buffers(reference).values(): @@ -157,7 +163,7 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( source = GDNState(config, request_capacity=1, device="meta") other_config = replace(config, logical_layer_ids=(1,), global_k_heads=4) else: - config = FullAttentionConfig( + config = GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=2, head_dim=4, diff --git a/unit_tests/common/state_cache_manager/test_linear_att.py b/unit_tests/common/state_cache_manager/test_linear_att.py deleted file mode 100644 index bceb67f4f6..0000000000 --- a/unit_tests/common/state_cache_manager/test_linear_att.py +++ /dev/null @@ -1,64 +0,0 @@ -import pytest -import torch - -from lightllm.common.state_cache_manager import LinearAttCacheConfig - - -@pytest.mark.parametrize("draft_layers", [0, 1, 3]) -@pytest.mark.parametrize( - "layer_types,expected_indices,linear_layers", - [ - ( - ("linear_attention", "linear_attention", "linear_attention", "full_attention") * 2, - (0, 1, 2, 0, 3, 4, 5, 1), - 6, - ), - ( - ( - "deepseek_sparse_attention", - "linear_attention", - "linear_attention", - "deepseek_sparse_attention", - "linear_attention", - ), - (0, 0, 1, 1, 2), - 3, - ), - ( - ("linear_attention", "full_attention", "linear_attention") * 2 + ("full_attention",), - (0, 0, 1, 2, 1, 3, 2), - 4, - ), - ], -) -def test_hybrid_layer_mapping_and_mtp_tail(layer_types, expected_indices, linear_layers, draft_layers): - config = LinearAttCacheConfig( - tp_world_size=1, - full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=128, - global_linear_k_heads=1, - global_linear_v_heads=1, - num_linear_k_heads=1, - num_linear_v_heads=1, - head_linear_k_dim=128, - head_linear_v_dim=128, - conv_kernel_size=4, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - layer_types=layer_types, - draft_full_att_kv_layer_num=draft_layers, - ) - assert config.all_layer_num == len(layer_types) - assert config.linear_layer_num == linear_layers - full_layers = len(layer_types) - linear_layers - assert config.get_main_model_full_att_layer_num() == full_layers - assert config.get_full_att_kv_layer_num_with_draft_model() == full_layers + draft_layers - for layer, (kind, expected) in enumerate(zip(layer_types, expected_indices)): - get_index = ( - config.get_linear_state_layer_index if kind == "linear_attention" else config.get_full_att_kv_layer_index - ) - assert get_index(layer) == expected - for draft in range(draft_layers): - assert config.get_full_att_kv_layer_index(len(layer_types) + draft) == full_layers + draft diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 080e4ef99d..301e0ae90a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -88,11 +88,9 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in full_layer_ids), - global_head_num=1, - head_dim=584, dtype=torch.bfloat16, tp_world_size=tp_world_size, - kv_dim=512, + latent_dim=512, ), page_size=args.page_size, ) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index f0490a6be4..2d1c6fca03 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -1279,13 +1279,14 @@ def fake_context(*args): def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): from lightllm.common.attention_state import AttentionStateManager - from lightllm.common.attention_state.full_att import FullAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) - config = FullAttentionConfig(logical_layer_groups=((3,),), global_head_num=1, head_dim=512, - dtype=torch.bfloat16, replicated=True) + config = MLAAttentionConfig( + logical_layer_groups=((3,),), latent_dim=512, dtype=torch.bfloat16, + ) state = create_mla_state( config=config, kv_type="fp8kv_dsa", packed=True, fp8_format="group64_e8m0", page_size=4, ) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 9415496fef..bc971862a4 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -71,12 +71,10 @@ def create(tp, target_mtp_step=None): full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in full_layer_ids), - global_head_num=1, - head_dim=584, dtype=torch.bfloat16, tp_rank=rank, tp_world_size=tp, - kv_dim=512, + latent_dim=512, ), page_size=args.page_size, ) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py index a3c2bb5795..bef86b47b7 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -111,7 +111,6 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat ) source_table = torch.arange(32, dtype=torch.int32).reshape(2, 16) backend = SimpleNamespace( - is_hybrid_att_model=False, model=SimpleNamespace(mem_manager=SimpleNamespace()), dp_world_size=tp_size, rank_in_dp=tp_rank, diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index b04d9ddc54..ff1c29add4 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -123,6 +123,8 @@ def make_runtime( num_hidden_layers=1, n_layer=1, full_attention_interval=1, + # draft 只有一个自有层，不能沿用 main 的多层类型列表。 + layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], ) draft_dir = tmp_path / "draft" draft_dir.mkdir() @@ -179,6 +181,7 @@ def make_runtime( num_hidden_layers=1, n_layer=1, full_attention_interval=1, + layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], ) draft.max_total_token_num, draft.max_req_num = 16, 1 draft.tp_world_size_, draft.tp_rank_ = tp, rank diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 0fcff8d09f..291edfddc2 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -10,7 +10,7 @@ from lightllm.common.attention_state.full_att import ( GQAFullAttentionState, - FullAttentionConfig, + GQAAttentionConfig, ) from lightllm.common.req_manager import ReqManager from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule @@ -60,10 +60,12 @@ def linear_cache(monkeypatch, tmp_path, request): conv_kernel_size=3, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, + full_attention_interval=2, + all_layer_num=2, layer_types=("linear_attention", "full_attention"), ) full = GQAFullAttentionState( - FullAttentionConfig( + GQAAttentionConfig( logical_layer_groups=tuple((layer,) for layer in config.get_full_att_layer_indices()), global_head_num=1, head_dim=8, diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py index 5efcfeba3e..f41b5ad22a 100644 --- a/unit_tests/server/test_pd_start_args.py +++ b/unit_tests/server/test_pd_start_args.py @@ -93,9 +93,9 @@ def test_disabled_cpu_cache_keeps_page_size_unset(startup_args): def test_default_attention_factory_uses_gqa(): import torch - from lightllm.common.attention_state.full_att import FullAttentionConfig, GQAFullAttentionState, create_gqa_state + from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState, create_gqa_state - state = create_gqa_state(config=FullAttentionConfig( + state = create_gqa_state(config=GQAAttentionConfig( logical_layer_groups=((0,),), global_head_num=4, head_dim=128, dtype=torch.bfloat16, ), device="meta") assert isinstance(state, GQAFullAttentionState) From 5ab58c174dae76717cc5aa6bd13b5b4721a86613 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 31 Oct 2026 16:07:52 +0000 Subject: [PATCH 30/56] refactor: unify K-pool window states and cache transfer regions --- docs/CN/source/models/add_new_model.md | 2 +- docs/EN/source/models/add_new_model.md | 2 +- docs/cache_manager_design.md | 366 +++++++----------- .../attention_state/att_tail/__init__.py | 5 - .../attention_state/att_tail/base_state.py | 92 ----- .../att_tail/create_kpool_state.py | 22 -- .../attention_state/att_tail/kpool_state.py | 174 --------- .../att_tail/kpool_transfer.py | 186 --------- .../attention_state/bounded_att/base_state.py | 10 +- .../bounded_att/fp8_e8m0_swa_state.py | 2 + .../attention_state/bounded_att/gdn_state.py | 2 + .../attention_state/bounded_att/kda_state.py | 2 + .../bounded_att/window_compressor_state.py | 52 ++- .../bounded_att/window_compressor_transfer.py | 80 +++- .../attention_state/full_att/base_state.py | 11 +- .../full_att/compressed_att_state.py | 2 + .../full_att/fp8_e8m0_packed_mla_state.py | 2 + .../attention_state/full_att/gqa_state.py | 2 + .../full_att/group_quantized_state.py | 2 + .../attention_state/full_att/mla_state.py | 2 + .../full_att/packed_mla_state.py | 2 + .../full_att/static_fp8_state.py | 2 + lightllm/common/attention_state/manager.py | 100 ++--- lightllm/common/attention_state/types.py | 6 +- lightllm/common/basemodel/basemodel.py | 20 +- lightllm/models/glm5_next/indexer.py | 2 +- lightllm/models/glm5_next/model.py | 107 +++-- lightllm/models/glm5_next_mtp/model.py | 20 +- .../server/router/model_infer/infer_batch.py | 2 - lightllm/utils/cpu_cache_layout.py | 2 +- test/att_resources/validate_model_cache.py | 29 +- unit_tests/common/att_resources/conftest.py | 6 +- .../common/att_resources/manager_fixtures.py | 14 +- .../common/att_resources/model_helpers.py | 34 +- .../common/att_resources/test_capacity.py | 34 +- .../att_resources/test_kpool_transfer.py | 332 ---------------- .../test_layer_state_selection.py | 12 +- .../att_resources/test_manager_helpers.py | 317 +-------------- .../att_resources/test_model_resources.py | 7 +- .../att_resources/test_network_transport.py | 6 +- .../att_resources/test_role_transfers.py | 17 +- .../att_resources/test_runtime_transport.py | 30 +- .../att_resources/test_state_creation.py | 14 +- .../att_resources/test_state_identity.py | 36 +- .../att_resources/test_state_queries.py | 17 +- .../common/att_resources/test_tail_state.py | 124 ------ .../test_window_compressor_state.py | 192 +++++++++ unit_tests/models/glm5_next/test_cache.py | 25 +- unit_tests/models/glm5_next/test_kernels.py | 6 +- unit_tests/models/glm5_next/test_layers.py | 4 +- unit_tests/models/glm5_next/test_pd_cache.py | 38 +- .../model_infer/test_multimodal_chunking.py | 6 +- 52 files changed, 775 insertions(+), 1806 deletions(-) delete mode 100644 lightllm/common/attention_state/att_tail/__init__.py delete mode 100644 lightllm/common/attention_state/att_tail/base_state.py delete mode 100644 lightllm/common/attention_state/att_tail/create_kpool_state.py delete mode 100644 lightllm/common/attention_state/att_tail/kpool_state.py delete mode 100644 lightllm/common/attention_state/att_tail/kpool_transfer.py delete mode 100644 unit_tests/common/att_resources/test_kpool_transfer.py delete mode 100644 unit_tests/common/att_resources/test_tail_state.py create mode 100644 unit_tests/common/att_resources/test_window_compressor_state.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 64cdddd139..60f276ac94 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -124,7 +124,7 @@ class TpPartBaseModel: | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | 读取初始化模型的 config.json, 并进行一些 key 名的同名合法化操作 | | def _verify_params(self)： | 校验参数 | -| `create_full_att_states(...)` / `create_bounded_att_states(...)` / `create_att_tail_states(...)` | 创建并返回当前模型的 full / bounded / tail State 列表；默认返回空列表，公共流程统一注册 | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` | 创建并返回当前模型的 full / bounded State 列表；默认返回空列表，公共流程统一注册 | | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index 34838b80e1..6d814a8fe7 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -124,7 +124,7 @@ Commonly used interfaces that need to be inherited and implemented | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | Read the config.json for initializing the model and perform some key name legalization operations | | def _verify_params(self)： | Validate parameters | -| `create_full_att_states(...)` / `create_bounded_att_states(...)` / `create_att_tail_states(...)` | Return this model's full / bounded / tail State lists; defaults return empty lists and the shared startup stage registers the results | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` | Return this model's full / bounded State lists; defaults return empty lists and the shared startup stage registers the results | | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 94d9721689..b6eef47841 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -4,7 +4,7 @@ ## 1. 设计约定 -统一在 `lightllm/common/attention_state/` 下组织可复用的资源实现，由 `AttentionStateManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的三类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 +统一在 `lightllm/common/attention_state/` 下组织可复用的资源实现，由 `AttentionStateManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的两类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： @@ -17,10 +17,10 @@ class AttentionStateManager: 这三个入口表示资源管理器对分配域的管理责任，聚合操作与传输接口见第 6.2 节。具体初始化必须保证所需 allocator 已创建或绑定，无需的 allocator 为 `None`。例如纯 full-attention 模型只初始化 `full_att_allocator`；启用 bounded 状态页的组合再按配置初始化大小页 allocator。GPU 运行态的存在本身不要求启用页池。 -- 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att`、`main_att_tail`、`draft_att_tail` 成员。具体模型按需组织资源组合与实现类型，拥有的 State 成员统一遵守第 3.5 节的角色、类型和存储方案命名。 -- `full_att`、`bounded_att`、`att_tail` 作为可复用实现的目录分类。full attention 保存历史 KV；bounded attention 由 `BoundedAttentionState` 抽象基类约定操作，当前包含 `KDAState`、`GDNState`，以后可加入 sliding-window 等有界状态，各自定义运行态与可选页池；attention tail 保存索引、池化、压缩等模块在增量计算中需要跨步骤保留的尾部数据及控制信息；具体实现决定恢复边界和快照要求。 +- 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att` 成员。具体模型按需组织资源组合与实现类型，拥有的 State 成员统一遵守第 3.5 节的角色、类型和存储方案命名。 +- `full_att` 保存全局历史 KV；`bounded_att` 管理绑定请求的有界运行态，包括 KDA/GDN、SWA，以及用于 KV 压缩和 K-pool 索引池化的窗口候选/gate。各 State 决定自身的快照需求和恢复边界。 - main/draft 可以采用不同方案；同一共享分配域的参与方使用同一个 allocator 实例，遵守统一的索引和回收约定。资源对象持有 buffer 和布局信息；管理器创建或绑定共享 allocator，调用方直接使用对应 allocator 分配与归还索引。 -- `AttentionStateManager` 提供运行态清零、清零指定请求的 tail、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 +- `AttentionStateManager` 提供运行态清零、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 - `req_to_token_indexs` 和请求槽位仍由 `ReqManager` 持有。资源管理器接收显式传入的物理索引、请求槽位及恢复长度，提供资源操作接口。 - 内部调用契约使用 `assert`，例如页类型、物理索引范围、Tensor 形状及共享分配域的一致性；用户配置、外部页面内容及格式兼容性继续显式校验，资源不足等运行失败保留相应异常。GPU 上的 MTP 状态选择保留异步断言，避免引入主机同步。 - 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 @@ -46,7 +46,7 @@ class AttentionStateManager: 管理器的 `write_req_to_cpu` / `read_req_from_cpu` 保留原名；绝对状态长度仍统一使用 `state_token_len` / `state_token_lens`，不改为 prefix 命名。`AxisMapping` 的 `local_start` / `external_start` 和 `local_circular` / `external_circular` 始终表示本地与外部两侧， -读回时仅反转数据方向；映射集合统一命名为 `axis_mappings`，单次覆盖为 `axis_mapping_overrides`。 +读回时仅反转数据方向；映射集合统一命名为 `axis_mappings`，请求窗口起点由 State 根据本次绝对前缀长度确定。 新增 Python 文件、类和接口按 **公共函数在前、私有辅助函数在后** 排列。已有基础流程文件以保留功能、减少 review 差异为优先，保留原有函数相对顺序，不为满足分组规范整体重排。 @@ -64,12 +64,12 @@ class AttentionStateManager: State 的基类与具体实现同样使用中文大块分隔注释，依次组织基础定义、该类资源的数据操作、CPU Cache、DP、PD、私有辅助实现。KDA/GDN 直接继承 `BoundedAttentionState`；所有具体 full-attention State 直接继承 `FullAttentionState`。具体 State 之间不继承实现，分别在自己的文件中按上述分块组织完整操作。 -- 分类数据操作统一带 `full_att`、`bounded_att` 或 `att_tail` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_runtime_to_pages`、`reset_attention_tail_state`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 +- 分类数据操作统一带 `full_att` 或 `bounded_att` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_runtime_to_pages`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 - `get_layer_view` 用于计算访问，CPU/DP/PD 保持按传输方向命名。State 不增加仅用于统一名字的转发方法。 - 所有 attention 配置字段必须提供中文注释，说明业务含义、单位或对应 Tensor 轴、全局/TP 本地范围，以及适用方案和默认值的特殊语义；派生属性还需说明计算关系。资源身份只保存在 State，不放入 config。 - 已知的调用参数使用显式 `key=value`，如 `size=...`、`dtype=...`、`device=...`、`pin_memory=...`；不先组装 `options` 字典再通过 `**options` 传入。 - 资源身份、布局范围和分配域选择等必要检查直接在使用处展开，不再包装成只有断言、布尔判断或转发循环的小函数。同一操作已校验的槽位继续复用，避免逐页、逐 State 重复校验分配器。 -- 每个公共接口使用中文 docstring 说明操作对象、索引含义和数据范围；在有差异时明确量化参数、前缀长度、MTP 游标及完整 tail 的处理。私有实现使用独立分块，关键布局或状态转换另加说明。 +- 每个公共接口使用中文 docstring 说明操作对象、索引含义和数据范围；在有差异时明确量化参数、前缀长度、MTP 游标及完整逻辑窗口的处理。私有实现使用独立分块，关键布局或状态转换另加说明。 - bounded 运行态重置、大小页保存与恢复、MTP 提交和运行态复制集中排列。pool 只保存快照 buffer；前缀与槽位的关联、页面发布和引用生命周期由调用流程管理。 ## 2. 整体结构与职责边界 @@ -86,8 +86,7 @@ flowchart TD Concrete --> Alloc[按需初始化或绑定的共享 allocator] Concrete --> Internal[模型自行组织的内部资源] Internal --> Full[full attention 实现] - Internal --> Bounded[bounded attention：KDAState / GDNState 等] - Internal --> Indexer[attention tail 实现] + Internal --> Bounded[bounded attention：KDA/GDN、SWA、窗口压缩器] Concrete -->|绑定具体视图| Forward[模型与 attention backend] Transport[DP/PD 调度] -->|调用具体资源操作| Concrete ``` @@ -150,9 +149,8 @@ KDAState # 一种具体 attention 状态实现 GDNState # 独立定义 GDN 状态、布局和快照操作 └── 按其实现组织 Conv / SSM / MTP 状态及可选大小页 pool -AttentionTailState # 抽象操作契约，不固定内部存储 -└── KPoolAttentionTailState # GLM K-pool 的具体状态 - └── tail_buffer # 原始 key/gate ring，按层、请求组织 +WindowCompressorState # bounded 窗口候选与 gate，供 V4 压缩和 GLM K-pool 使用 +└── buffer # 同 dtype 的 candidate|gate 拼接 Tensor ``` 同一共享分配域中的 allocator 只初始化一次，由 main/draft 共用的唯一 `AttentionStateManager` 管理生命周期。调用方解除相关引用后，向对应 allocator 单次归还索引，不因模型数增加而重复分配或释放。 @@ -167,7 +165,7 @@ AttentionTailState # 抽象操作契约，不固定内部存储 | KV、量化 scale、packed indexer buffer 与层映射 | `AttentionStateManager` 组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | | KDA/GDN 的 Conv/SSM、MTP 游标 | `AttentionStateManager`持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | | 大小页 checkpoint buffer | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID 保存数据；不持有独立空闲列表或长度记录 | -| indexer 原始尾部窗口 | `AttentionStateManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | +| indexer 原始尾部窗口 | `AttentionStateManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需窗口布局 | | 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | | 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复关联资源数据，解除引用后直接归还 allocator | @@ -197,7 +195,7 @@ main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAStat 共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 -linear 运行态与两侧 indexer tail 均由 `ReqManager` 提供的 `req_idx` 定位。indexer tail 不单独分配 checkpoint 槽位；其状态布局和恢复规则由各自实现描述。 +linear 运行态与两侧索引池化窗口均由 `ReqManager` 提供的 `req_idx` 定位。索引池化窗口不单独分配 checkpoint 槽位；其状态布局和恢复规则由各自实现描述。 ### 3.3 分配状态与数据状态分离 @@ -205,7 +203,7 @@ allocator 是索引分配的唯一事实来源，资源对象持有存储和布 同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 -对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由`AttentionStateManager`维护。 +对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，索引池化窗口资源管理增量计算需要的原始短窗口；二者的组合关系由`AttentionStateManager`维护。 ### 3.4 按具体 attention 实现定义状态 @@ -243,11 +241,11 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 ### 3.5 State 成员命名、资源身份与唯一性 -`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。各具体 full-attention 配置、KDA/GDN 的 `ConvSSMStateConfig` 和 `KPoolStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 +`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。各具体 full-attention 配置、KDA/GDN 的 `ConvSSMStateConfig` 和 `WindowCompressorStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 -管理器成员名与 `state.state_name` 完全一致，统一使用 `{model_role}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState`、`AttentionTailState` 类型表达，名称中不再重复 `full_att`、`bounded_att`、`att_tail`。 +管理器成员名与 `state.state_name` 完全一致，统一使用 `{model_role}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState` 类型表达，名称中不再重复 `full_att`、`bounded_att`。 -三类 State 都必须实现 `get_state_type()`，只返回中间的算法及方案标记，使用小写 snake_case。方案标记区分量化位宽、group size、scale 方式/精度，KDA/GDN 分别标注 Conv 和 SSM 精度，K-pool 标注压缩比例与存储精度。层号、容量、设备、TP rank 和请求信息不进入名称。 +两类 State 都必须实现 `get_state_type()`，只返回中间的算法及方案标记，使用小写 snake_case。方案标记区分量化位宽、group size、scale 方式/精度，KDA/GDN 分别标注 Conv 和 SSM 精度，K-pool 标注压缩比例与存储精度。层号、容量、设备、TP rank 和请求信息不进入名称。 同一角色、同一标记只能注册一个 State，包含该方案的全部逻辑层；不同存储方案可以共存，通过算法字符串、层归属及可选的 State 类别筛选。名字用于资源定位，不能替代完整布局兼容性检查。标记注册后保持稳定，meta 定义、实际存储和 CPU/DP/PD 进程使用相同名称，不提供旧名称别名。 @@ -259,9 +257,9 @@ State 直接持有按层组织的 Tensor，不额外设置通用状态管理对 | MLA BF16 | `draft0_mla_bf16_state` | | 带 indexer 字段的 packed MLA FP8 | `main_mla_packed_fp8_group128_state` | | KDA / GDN | `main_kda_conv_bf16_ssm_fp32_state` / `draft0_gdn_conv_bf16_ssm_fp32_state` | -| K-pool C4 indexer tail | `main_kpool_c4_bf16_state` | +| K-pool W4/S4 索引池化窗口 | `main_compressor_indexer_w4_s4_bf16_state` | -普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；K-pool 使用 `kpool_cN_*` 标明压缩比例与存储精度。完整 shape、标定方式等仍保存在 State 配置中。 +普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；窗口压缩器使用 `compressor_{target}_wW_sS_*` 标明用途、窗口、步长与存储精度。完整 shape、标定方式等仍保存在 State 配置中。 以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： @@ -272,9 +270,9 @@ AttentionStateManager ├── bounded_att_small_page_allocator ├── main_mla_packed_bf16_state ├── main_kda_conv_bf16_ssm_fp32_state -├── main_kpool_c4_bf16_state +├── main_compressor_indexer_w4_s4_bf16_state ├── draft0_mla_packed_bf16_state -└── draft0_kpool_c4_bf16_state +└── draft0_compressor_indexer_w4_s4_bf16_state ``` 当前 GLM draft 不使用 KDA，因此不创建 `draft0_kda_conv_bf16_ssm_fp32_state`。Qwen 的线性状态为 `main_gdn_conv_bf16_ssm_fp32_state`；Full Attention 按实际配置创建 `main_gqa_bf16_state` 等成员。MLA 与普通 KV 都可使用同一个 `full_att_allocator` 分配域。 @@ -293,7 +291,7 @@ main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 ## 4. 目录设计 -目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器使用 `manager.py`，三类公共 State 契约使用各自目录下的 `base_state.py`；allocator 和 pool 文件按对应类名命名，具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 +目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器使用 `manager.py`，两类公共 State 契约使用各自目录下的 `base_state.py`；allocator 和 pool 文件按对应类名命名，具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 ```text lightllm/common/ @@ -340,7 +338,7 @@ lightllm/common/ │ │ ├── fp8_e8m0_swa_state.py # FP8E8M0SlidingWindowConfig / State：仅管理近期 KV │ │ ├── create_window_compressor_state.py # 窗口压缩器原子运行态创建入口 │ │ ├── window_compressor_state.py # WindowCompressorStateConfig / State：候选与 gate -│ │ ├── window_pages_pool.py # WindowPagesPool：单一窗口的时间有序快照 +│ │ ├── window_pages_pool.py # SWA 与压缩器的时间有序窗口快照池 │ │ ├── fp8_e8m0_swa_transfer.py # FP8/E8M0 packed SWA 的 GPU 窗口搬运 │ │ ├── window_compressor_transfer.py # 候选/gate ring 与时间有序快照的 GPU 搬运 │ │ ├── create_kda_state.py # KDA 统一创建入口 @@ -349,27 +347,21 @@ lightllm/common/ │ │ ├── create_gdn_state.py # GDN 统一创建入口 │ │ ├── gdn_state.py # GDNState：独立算法状态与页操作 │ │ └── gdn_pages_pool.py # GDNPagesPool：支持不同 K/V head 数和维度的快照 -│ └── att_tail/ -│ ├── __init__.py -│ ├── base_state.py # AttentionTailState：尾部状态抽象操作契约 -│ ├── kpool_transfer.py # transfer_kpool_pages：单次 kernel 双向搬运环形逻辑窗口 -│ ├── create_kpool_state.py # K-pool 统一创建入口 -│ └── kpool_state.py # KPoolAttentionTailState：原生 ring 与恢复规则 └── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 lightllm/models// -├── model.py # 三类 State 创建接口，决定本模型的资源组合 +├── model.py # 两类 State 创建接口，决定本模型的资源组合 └── config.py # 模型专有配置与派生参数 lightllm/utils/ └── cpu_cache_layout.py # build_cpu_cache_layouts：构建 CPU cache 联合页布局 ``` -`allocators/` 与三类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 +`allocators/` 与两类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 -`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口，各具体配置与对应 State 放在同一文件中；`bounded_att/base_state.py`、`att_tail/base_state.py` 分别定义另外两类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 +`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口，各具体配置与对应 State 放在同一文件中；`bounded_att/base_state.py` 定义有界运行态的公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 -`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。`KPoolAttentionTailState` 管理请求运行态，不创建独立大小页 pool。 +`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。K-pool 使用 W4/S4 的 `WindowCompressorState`，在整组边界无重叠历史，不创建大小页 pool。 `ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `state_name`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 @@ -393,15 +385,15 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；`AttentionTailState` 实现 PD 快照读写，CPU/DP 缓存前缀恢复时按对齐边界重置；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`att_tail/kpool_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运最近 `pool_size` 个逻辑位置；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 和 tail 的共用传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；K-pool tail 直接提交本方案的环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在三类函数中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。tail 使用固定 `pool_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；tail 检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；无重叠的窗口 State 在 CPU 对齐边界不需要快照；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`bounded_att/window_compressor_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运外部快照所需的逻辑窗口；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 的传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；窗口压缩器直接提交环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在各方案中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。压缩器使用固定 `window_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；窗口搬运检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 -无需模型特有初始化的场景直接使用 `AttentionStateManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention、bounded attention 和 indexer tail，不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `register_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 +无需模型特有初始化的场景直接使用 `AttentionStateManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention 与 bounded attention（包括压缩窗口），不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `register_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 -模型特有的 State 组合由 `model.py` 中三类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool tail；main/draft 分别返回各自的 State。所有模型共用 `AttentionStateManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 +模型特有的 State 组合由 `model.py` 中两类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool 窗口；main/draft 分别返回各自的 State。所有模型共用 `AttentionStateManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 -token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 tail 重置直接在 `AttentionStateManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttentionStateManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states`、`_tail_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 state_name 匹配时可构建临时映射，它不决定模型的资源组织方式。 +token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交直接在 `AttentionStateManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttentionStateManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 state_name 匹配时可构建临时映射，它不决定模型的资源组织方式。 #### 按逻辑算法创建 State @@ -415,10 +407,9 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交及 | KDA | `bounded_att.create_kda_state(config, request_capacity, mtp_step=0, *, device)` | KDA Conv/SSM 运行态 | | GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | | SWA | `bounded_att.create_sliding_window_state(config, request_capacity, mtp_step, device)` | 近期原始 KV 窗口，当前采用 FP8/E8M0 packed 编码 | -| 窗口压缩器 | `bounded_att.create_window_compressor_state(config, request_capacity, mtp_step, device)` | 按窗口长度和产出步长管理候选/gate，KV 和 indexer 分别创建实例 | -| K-pool | `att_tail.create_kpool_state(config, request_capacity, mtp_step=0, *, device)` | 带池化尾部窗口及 MTP 历史槽位的 K-pool State | +| 窗口压缩器 | `bounded_att.create_window_compressor_state(config, request_capacity, mtp_step, device)` | 按窗口长度和产出步长管理候选/gate，KV 与索引 key 池化分别创建实例 | -每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器和 K-pool。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 +每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器。K-pool 复用窗口压缩器入口。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 GQA/MLA 的创建入口分别接收 `GQAAttentionConfig`、`MLAAttentionConfig`，每种具体 State 使用独立配置，不继承通用配置。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.latent_dim` 始终传逻辑 latent 维度，`packed=True` 时由 packed 配置推导包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 @@ -437,7 +428,7 @@ State 将连续的 KV/scale/indexer 字节作为一个传输区域，复用 `tra 同卡 token 复制与 DP 迁移仍复制完整物理行。原 packed State 及其传输实现无需识别 E8M0 编码。 定制量化写入放在 `fp8_e8m0_packed_mla_write.py`，State 不依赖 GLM 模型目录中的实现。 -各模型的三类 State 工厂负责构造配置并调用对应算法入口；`lightllm/utils/cpu_cache_layout.py` 复用模型工厂创建 meta State，生成 CPU cache 联合页布局，不加载标定。full-attention 的实际存储由管理器后续统一分配；GPU 启动时 bounded 和 tail 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 +各模型的两类 State 工厂负责构造配置并调用对应算法入口；`lightllm/utils/cpu_cache_layout.py` 复用模型工厂创建 meta State，生成 CPU cache 联合页布局，不加载标定。full-attention 的实际存储由管理器后续统一分配；GPU 启动时 bounded 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 ```python state = create_gqa_state( @@ -448,11 +439,11 @@ manager.register_state("main", state) # 全部 State 准备完成后，由 manager 统一规划容量并分配存储。 ``` -根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名及传输视图的尾维合并、整批最大搬运 dtype 选择，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded/tail 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 +根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名及传输视图的尾维合并、整批最大搬运 dtype 选择，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 -管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后清零目标 tail；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttentionStateManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 +管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后恢复目标压缩窗口；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttentionStateManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 -`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisMapping`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 +`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisMapping`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求两类 State 继承同一个大接口。 ## 5. allocator 与状态页 pool @@ -470,7 +461,7 @@ allocator 入口的存在只能说明该分配域已启用，具体恢复边界 大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 -`FullAttKVAllocator` 的申请量和返回索引均以 token 表示，分配和回收须满足整页约束；`StatePageAllocator` 返回状态页槽位 ID。GPU 运行态和 indexer tail 使用请求槽位，不使用上述状态页分配器。 +`FullAttKVAllocator` 的申请量和返回索引均以 token 表示，分配和回收须满足整页约束；`StatePageAllocator` 返回状态页槽位 ID。GPU 运行态（包括压缩窗口）使用请求槽位，不使用上述状态页分配器。 HOLD 页位于普通 KV 分配域之后，由 full-attention 存储预留。初始化流程将 allocator 的 `hold_token_indices` 显式交给 `ReqManager.bind_hold_page` 建立 padding 请求映射；ReqManager 不持有 allocator，也不负责释放 KV。`FullAttKVAllocator` 构造时直接创建 `shared_can_use_token_num: SharedInt` 并发布初始容量，在 `alloc`、`free` 和 `free_all` 中直接更新。对应的 `KVCapacityReader` 定义在同一文件中，router 按本地 DP 分组读取共享计数；Reader 不重置计数。读写双方使用同一服务前缀与 `rank_in_node` 命名，因此构造 allocator 前须完成服务名和 rank 初始化，无需外部绑定或容量回调。 @@ -495,7 +486,7 @@ CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排 前缀长度与大小页槽位的关联由请求及缓存流程维护。`big_pages_pool` 和 `small_pages_pool` 仅保存算法快照，不重复持有长度数组或页有效性标记。 -大小页保存/恢复和 bounded 运行态复制只接收索引与布局，不接收前缀长度。`reset_attention_tail_state(req_indices)` 只清零指定请求的 tail，manager 的 CPU/DP 操作不传递专用于重置的长度。PD tail 读写单独接收 `state_token_len`，用于把 prefill 末尾的逻辑位置映射到两端各自的 ring；该值来自 PD 任务的 `end_kv_index`，不保存到 State 或传输页内。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 调用方只能发布已完成全部参与 State 保存的页面，并通过 stream/event 保证异步操作顺序。尚未捕获的中间边界不能发布为可复用快照；失败时不发布对应页面，等待已提交操作结束后再回收槽位。前缀身份、引用与就绪状态均由调用流程管理。 @@ -519,7 +510,7 @@ logical_layer_groups = ( # 不共享时使用 ((0,), (4,), (8,))，对应三份物理层存储。 ``` -层命名统一使用 `logical_` 和 `storage_` 区分含义：`logical_layer_groups` 描述共享分组，`logical_layer_id` 是推理层传入的逻辑层号，`storage_layer_index` 是 Tensor layer 轴上的物理下标。State 的 `_logical_to_storage_layer` 保存两者的对应关系。KDA/GDN 和 K-pool 当前仍是一对一存储，配置使用 `logical_layer_ids`；三类 State 的 `get_layer_view(logical_layer_id=...)` 使用相同参数语义。FP8 标定的 `storage_layer_count`、`storage_layer_start` 表示物理行数和偏移，不能传逻辑层号。Tensor 轴标签 `"layer"` 保留，表示存储层轴。 +层命名统一使用 `logical_` 和 `storage_` 区分含义：`logical_layer_groups` 描述共享分组，`logical_layer_id` 是推理层传入的逻辑层号，`storage_layer_index` 是 Tensor layer 轴上的物理下标。State 的 `_logical_to_storage_layer` 保存两者的对应关系。KDA/GDN 和 K-pool 当前仍是一对一存储，配置使用 `logical_layer_ids`；两类 State 的 `get_layer_view(logical_layer_id=...)` 使用相同参数语义。FP8 标定的 `storage_layer_count`、`storage_layer_start` 表示物理行数和偏移，不能传逻辑层号。Tensor 轴标签 `"layer"` 保留，表示存储层轴。 静态 FP8 的 FA3/FlashInfer 后端通过 `AttControl.logical_layer_id` 接收当前逻辑层号，再由 State 的 `get_storage_layer_index()` 查询标定行，不通过 KV 指针反查层号。该接口属于 `FullAttentionState` 的公共契约，所有具体实现按自己的层映射返回物理行；层视图和 KV 写入也使用此接口。 @@ -529,13 +520,13 @@ State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get #### State 层归属与模型装配 -三类 State 均提供 `contains_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 +两类 State 均提供 `contains_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 -`AttentionStateManager.get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_role` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`state_type_filter` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`kpool_c4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `state_type_filter=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 +`AttentionStateManager.get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_role` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`state_type_filter` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`compressor_indexer_w4_s4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `state_type_filter=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 -该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_state_registration_finished=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。三类模型工厂只返回 State 列表，`prepare_attention_states()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 +该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_state_registration_finished=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。两类模型工厂只返回 State 列表，`prepare_attention_states()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 -推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层 tail 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 +推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层窗口 State 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 #### 继承与内部 storage 的关系 @@ -657,7 +648,7 @@ class FullAttentionState: #### 存储字段与 State 元数据 -三类 State 不提供完整存储枚举接口，各实现显式确定参与传输的固定字段。Full attention 的布局直接根据配置计算形状、dtype 和字节区间；PD/DP 调用处显式列出实际字段：GQA 为 K/V，MLA 为 latent，分组量化增加动态 scale，packed MLA 包含对应的 KV、scale、RoPE 与 indexer 字段。静态标定参数独立保留，不加入 token 传输字段。KDA/GDN 直接从配置描述 Conv/SSM，K-pool 直接描述完整 tail ring。 +两类 State 不提供完整存储枚举接口，各实现显式确定参与传输的固定字段。Full attention 的布局直接根据配置计算形状、dtype 和字节区间；PD/DP 调用处显式列出实际字段：GQA 为 K/V，MLA 为 latent，分组量化增加动态 scale，packed MLA 包含对应的 KV、scale、RoPE 与 indexer 字段。静态标定参数独立保留，不加入 token 传输字段。KDA/GDN 直接从配置描述 Conv/SSM，K-pool 直接描述完整压缩器 ring。 `TensorField` 仅作为内部搬运数据描述，仅保留 `tensor`、`axes`，字段名称由所在字典的 key 提供；Tensor 自身提供 dtype、shape、stride、device 和底层存储视图。运行态、MTP 游标及大小页池由具体 State 直接持有，不额外导出一份完整字段列表。 @@ -667,15 +658,19 @@ CPU/PD 的字段选择由 State 显式定义，公共区域接口如下： | --- | --- | --- | | `FullAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | | `BoundedAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | -| `AttentionTailState` | 无；CPU 缓存恢复通过边界重置 | `get_pd_transfer_regions()` | -所有接口返回 `tuple[TransferRegion, ...]`，可以基于 meta State 调用，不分配传输 buffer。State 决定哪些字段参与传输、字段区域的 dtype/shape、TP 切片和共享页的副本写入者；每份描述以 `byte_start` / `byte_end` 记录左闭右开的字节区间，两个位置均从该 State 的区域起点计算。三类 State 的 CPU/PD 区域接口统一无参数。full-attention 从全局启动参数读取 `cpu_cache_token_page_size` 或 `pd_kv_page_size`，按固定页容量展开形状与字段偏移；全局 head 数与本地 TP 切片均由自身 config 推导。bounded/indexer 每次描述一份运行态。 +两类 State 的 CPU/PD 区域接口统一使用内层 `@lru_cache(maxsize=None)`、外层 +`@deepcopy_result`：缓存首次构建的字段描述，每次查询返回独立副本，包含区域内的轴映射。 +首次查询后配置与全局页规格保持不变；实际存储分配完成后，管理器清除各 State 的两类区域缓存， +再清除管理器的派生布局和页大小缓存。子类覆盖区域接口时也须保留这两个装饰器。 + +所有接口返回 `tuple[TransferRegion, ...]`，可以基于 meta State 调用，不分配传输 buffer。State 决定哪些字段参与传输、字段区域的 dtype/shape、TP 切片和共享页的副本写入者；每份描述以 `byte_start` / `byte_end` 记录左闭右开的字节区间，两个位置均从该 State 的区域起点计算。两类 State 的 CPU/PD 区域接口统一无参数。full-attention 从全局启动参数读取 `cpu_cache_token_page_size` 或 `pd_kv_page_size`，按固定页容量展开形状与字段偏移；全局 head 数与本地 TP 切片均由自身 config 推导。bounded（包括压缩窗口） 每次描述一份运行态。 本次有效 token 数属于读写流程，不参与布局生成。CPU cache 使用固定的全局 CPU 页布局；State 的传输接口在不足整页时仍使用相同的字段形状和偏移，只搬运有效 token。缓存页的保存边界和发布条件由外层流程控制。 GQA、MLA 及各量化方案只导出 KV、动态 scale 和随 token 保存的 indexer 字段，不传静态标定参数或 packed padding。KDA/GDN 只导出有效 Conv 窗口和一份已接受 SSM，不传 MTP 候选区、游标或整个大小页池。K-pool 导出完整原生 key/gate ring。CPU 与 PD 可以在具体子类中分别覆盖；当前字段布局相同的实现共用自己的私有构建方法。 -manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字段的实际最大结束位置安排后续资源，并通过布局终点计算整页大小；manager 不解析 Conv/SSM/K-pool 的内部结构。CPU 服务复用相同 State 接口；当前 CPU 联合页仍只包含 full/bounded，tail 按恢复边界重置。DP 映射只包含 full/bounded attention，调用它们无参数的 PD 区域接口，只取字段名和本地轴切片，不使用区域的页形状与字节偏移；本次复制的 token/request 数量由调用参数决定。 +manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字段的实际最大结束位置安排后续资源，并通过布局终点计算整页大小；manager 不解析 Conv/SSM/K-pool 的内部结构。CPU 服务复用相同 State 接口；当前 CPU 联合页包含 full/bounded 中需要保存的字段，无重叠窗口不占 CPU 页。DP 映射只包含 full/bounded attention，调用它们无参数的 PD 区域接口，只取字段名和本地轴切片，不使用区域的页形状与字节偏移；本次复制的 token/request 数量由调用参数决定。 资源身份使用 `state.state_name`；层映射、容量和 TP 分片事实继续由 State 的 config、capacity、page_size 及计算接口维护。CPU/DP/PD 由抽象接口约定并由具体 State 实现，管理器直接调用对应方法，不再维护额外的传输能力集合。信息保留在其所属对象，传输布局仅描述当前路径所需字段。 @@ -761,14 +756,13 @@ State 在后续 `allocate_storage()` 中自行分配数据 buffer：GQA 与静 | 接口 | 语义 | | --- | --- | | `AttentionStateManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | -| `register_state(model_role, state)` | 将已创建的 full/bounded/indexer State 原对象绑定为规范成员，同步 State 与 config 身份；只允许在创建完成前调用 | +| `register_state(model_role, state)` | 将已创建的 full/bounded State 原对象绑定为规范成员，同步 state_name 与管理器成员名；只允许在创建完成前调用 | | `finish_state_registration()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_state_registration_finished = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | | `load_full_att_calibration(args)` | 在 `finish_state_registration()` 后、存储分配前加载静态 FP8 标定；文件只读一次，逐 State 绑定物理层区间及 TP 参数；没有静态 FP8 State 时直接返回 | | `initialize_storage(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：在完成创建标记和标定绑定后，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | | `get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_state_registration_finished` 并使用 `lru_cache` 缓存 | | `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | | `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | -| `get_attention_tail_states(role=None)` | 返回指定角色或全部 `AttentionTailState` 的 tuple | | `get_full_att_token_bytes()` | 返回当前 TP rank 全部 full-attention State 的单 token 存储字节数，包含动态 scale 和行 padding | | `get_cpu_page_bytes()` | 依赖 `get_cpu_page_layouts()`，取全部 State 的最大 `byte_end` 并按 16 字节对齐；包含字段间及 State 间的必要 padding，页内不保存格式指纹、长度或版本 | | `get_cpu_page_layouts()` | 按全局 `cpu_cache_token_page_size` 生成 CPU 联合页中各 State 的字段偏移和 TP 映射，使用 `lru_cache` 缓存 | @@ -784,11 +778,11 @@ State 在后续 `allocate_storage()` 中自行分配数据 buffer：GQA 与静 State 装配期间由管理器内部遍历 `__dict__`，公共聚合查询及 `get_state` 均会断言失败，避免缓存不完整的 main/draft 集合；模型准备阶段只创建和注册 State，不查询推理资源。装配完成后先调用 `finish_state_registration()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 -容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `allocate_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded/indexer 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 +容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `allocate_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded（包括压缩窗口） 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 实际存储分配完成后执行 warmup。CPU 布局由第一次容量查询或读写时按需生成；PD 容量规划会通过 `get_pd_page_bytes()` 在 meta State 上提前生成 PD 布局，不分配实际传输 buffer。PD 页的 token 容量由全局 `get_env_start_args().pd_kv_page_size` 唯一确定，两个 PD 布局 getter 均不接收页 token 数，按 manager 实例缓存并直接返回 `Mapping[str, StateTransferLayout]`，没有对应资源时返回空字典。 -PD 的依赖方向为：`get_pd_page_bytes()` 分别查询两个独立的布局 getter。full-attention 和 bounded/indexer 各自生成偏移及本类布局所需大小，互不查询，也不调用 `get_pd_page_bytes()`。实际块大小取两类布局占用的最大值，统一在 `get_pd_page_bytes()` 完成。State 提供包含内部字段 padding 的相对偏移，manager 按字段的实际最大结束位置确定下一个 State 的起点，不再用裸存储字节数预留 PD 区域。PD 收发流程按 `kv` / `att_state` 任务选择布局。 +PD 的依赖方向为：`get_pd_page_bytes()` 分别查询两个独立的布局 getter。full-attention 和 bounded（包括压缩窗口） 各自生成偏移及本类布局所需大小，互不查询，也不调用 `get_pd_page_bytes()`。实际块大小取两类布局占用的最大值，统一在 `get_pd_page_bytes()` 完成。State 提供包含内部字段 padding 的相对偏移，manager 按字段的实际最大结束位置确定下一个 State 的起点，不再用裸存储字节数预留 PD 区域。PD 收发流程按 `kv` / `att_state` 任务选择布局。 CPU getter 从全局启动参数读取 `cpu_cache_token_page_size`，同样按 manager 实例缓存。不保存独立的 `cpu_page_bytes`、`pd_page_bytes`、`cpu_transfer_layouts`、`pd_transfer_layouts` 或 `cpu_page_tokens` 成员。存储分配可能改变 State 的物理页规格，因此分配完成后同时清除布局及派生 CPU/PD 页大小缓存。初始化时的异常直接向上传播，初始化失败即终止该流程；不增加属性拦截或回滚逻辑。 @@ -805,20 +799,19 @@ IPC 发布按用途调用 `create_pd_ipc_copy()` 或 `create_dp_ipc_copy()`， | 操作 | 接口 | 聚合语义 | | --- | --- | --- | | full-attention token 复制 | `copy_full_att_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | -| bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页 | +| bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页 | | bounded 状态页保存 | `save_bounded_att_runtime_to_pages(...)` | 按 state_token_lens 保存全部参与 bounded State 的同类 pool、同一组槽位 | | bounded 状态页恢复 | `restore_bounded_att_runtime_from_pages(...)` | 按 state_token_lens 将选定快照恢复到对应请求运行态及 ring 位置 | | bounded 验证后提交 | `commit_bounded_att_state(req_indices, accepted_steps, *, role)` | 将指定角色的接受步数分发给参与 bounded State；由 State 选择已接受运行态，不暴露游标 Tensor | -| attention tail 初始化与无快照恢复 | `reset_attention_tail_state(req_indices)` | 清零指定请求在全部参与 tail 中的状态；调用方保证空前缀或可重置的缓存边界，保留其他状态 | | CPU 写出 / 读回 | `write_req_to_cpu(...)` / `read_req_from_cpu(...)` | 组合 KV 和状态页，按调用参数提供的长度读写各 State 在外部页中的区域 | -| DP 复制 | `copy_from_dp(remote_att_manager, src_token_indices, dst_token_indices, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，清零目标 tail；不分配资源 | +| DP 复制 | `copy_from_dp(remote_att_manager, src_token_indices, dst_token_indices, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，复制目标窗口运行态；不分配资源 | | PD 写出 / 读回 | `write_to_pd_page(...)` / `read_from_pd_page(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | #### 必须实现的聚合接口 -基类源码使用醒目的分隔线划分八个大块：基础管理、Full attention、Bounded attention、Indexer tail、CPU Cache 数据交换、DP 迁移、PD 分离传输、私有辅助实现。Bounded attention 分块依次放置运行态重置、大小页保存与恢复、MTP 提交接口。各块内保留接口的具体说明，接口顺序保持 public 在前、private 在后。 +基类源码划分七个大块：基础管理、Full attention、Bounded attention、CPU Cache 数据交换、DP 迁移、PD 分离传输、私有辅助实现。Bounded attention 分块依次放置运行态重置、大小页保存与恢复、MTP 提交接口。各块内保留接口的具体说明，接口顺序保持 public 在前、private 在后。 -管理器的分类操作使用 `full_att`、`bounded_att`、`att_tail` 标明作用范围；CPU/DP/PD 接口继续表示跨资源传输。具体 State 的分类操作与管理器使用相同命名；State 只操作自身数据，管理器负责聚合参与对象。 +管理器的分类操作使用 `full_att`、`bounded_att` 标明作用范围；CPU/DP/PD 接口继续表示跨资源传输。具体 State 的分类操作与管理器使用相同命名；State 只操作自身数据，管理器负责聚合参与对象。 以下仅展示接口签名，实际方法在基类中提供聚合实现。 @@ -837,7 +830,7 @@ class AttentionStateManager: def __init__(self, request_capacity: int): ... def register_state( - self, model_role: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState, + self, model_role: str, state: FullAttentionState | BoundedAttentionState, ) -> None: ... # ======================================================================== @@ -877,9 +870,7 @@ class AttentionStateManager: ) -> None: ... # ======================================================================== - # Indexer tail：空前缀初始化与无快照恢复 # ======================================================================== - def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... # 基类使用 lru_cache 缓存布局，存储初始化后清除缓存。 def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... @@ -894,7 +885,7 @@ class AttentionStateManager: small_page_slot_ids: List[int] | None = None, ) -> None: ... - # 一次读取同一前缀的有序页面；全部页面读回后清零 req_idx 的 tail。 + # 一次读取同一前缀的有序页面；全部页面读回后从末页恢复请求运行态。 def read_req_from_cpu( self, cpu_page_buffers, cpu_page_indices, token_indices, *, layouts: Mapping[str, StateTransferLayout], @@ -930,7 +921,7 @@ class AttentionStateManager: ) -> None: ... ``` -以上方法均有基类默认实现，直接按已创建的 State 聚合执行。子类仅实现模型特有部分。未支持的 CPU/DP/PD 路径须明确报错，启用相应功能时检查参与 State 的完整支持能力。未启用某种页分配域时，保存和恢复该域明确报错；纯 full-attention 模型无需创建大小页 pool，其 `reset_bounded_att_state` 可以为空操作；没有 indexer tail 的组合，其 `reset_attention_tail_state` 也可以为空操作。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 `page_kind="big"` 对应 `bounded_att_big_page_allocator` 和参与 State 的 `big_pages_pool`；`"small"` 对应小页 allocator 与 `small_pages_pool`。两者是独立分配域，即使槽位数字相同，也不是同一份资源。 @@ -943,11 +934,10 @@ class AttentionStateManager: `reset_bounded_att_state(req_indices: List[int])` 用于从前缀长度 0 开始计算时初始化 bounded 运行态。`req_indices` 仅接受 Python 整数列表，不支持 `torch.Tensor`；实现直接使用列表，不做类型转换。调用方已通过 ReqManager 取得请求槽位，并确认这些槽位此前的计算和传输使用结束；管理器只按传入索引初始化 bounded 数据，不查询缓存命中情况或请求对象。 - 对全部参与的 KDA/GDN 等状态调用 `reset_bounded_att_state(req_indices)`，将 Conv/SSM 恢复到零初值，并重置 MTP 游标及影响后续读取的候选信息。未来算法按自身契约恢复空前缀初态。 -- 不修改 AttentionTailState。没有缓存命中时，调用方另行调用 `reset_attention_tail_state(req_indices)`，完成 tail 的空前缀初始化。 - full-attention 没有需要清零的空前缀递推状态；该操作不遍历或清空历史 KV，不更改请求映射，也不失效大小页快照或归还索引。纯 full-attention 管理器合法地不做数据操作。 - 目标请求槽位须唯一且在有效范围内，main/draft 所需运行态全部初始化后才可使用。重复执行在没有并发访问时具有相同结果；方法返回仍遵守 stream 完成约定。 -这是长度 0 的数据初始化接口，不接收非零前缀长度。命中非零前缀时使用 `restore_bounded_att_runtime_from_pages` 或 CPU/DP/PD 读回；没有 tail 快照时在保证可重置边界后通过 `reset_attention_tail_state` 清零 tail；不能在 bounded 运行态恢复后调用此方法，否则会清除已经恢复的 bounded 数据。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 #### 验证后运行态提交 @@ -957,25 +947,12 @@ main/draft 共享管理器，因此必须显式指定 `role="main"` 或具体 dr `mtp_state_index` 及其他提交控制信息只由具体 State 持有。BaseModel 不维护 `mtp_state_cursors` 或其他游标引用集合，验证流程也不直接读写游标。State 自行决定通过游标选择、数据搬移或其他方式提交；计算 backend 可通过 State 或层视图读取所需信息。 -#### 无 tail 快照时清零尾部状态 - -`reset_attention_tail_state(req_indices)` 清零指定请求的 tail，用于空前缀初始化或可重置的缓存前缀恢复。请求槽位使用 `list[int]`，目标唯一且已分配；接口不接收前缀长度，不读取请求映射或缓存节点。 - -管理器依次调用各 AttentionTailState 的 `reset_attention_tail_state(req_indices)`。K-pool 直接清零指定请求的完整 ring，不在 reset 中重复检查长度。模型初始化与缓存流程保证恢复边界对齐：缓存页大小必须是各参与 tail 池化步长的整数倍，main/draft 使用不同方案时分别满足各自约束。异常直接透传，不保证跨 State 或整个 CPU/DP 恢复流程的原子性。 - -该操作只重置 tail 及其必要控制信息，保留已恢复的 Conv/SSM、MTP 游标、KV 和大小页快照。没有 tail 的管理器合法地不做数据操作。调用方保证此前使用已结束，并在相关数据操作完成后继续计算；这不是运行中 MTP 回退的替代操作。 - -前缀页命中时，调用方恢复 KV 和 bounded 页后，调用本接口清零未保存的 tail。PD 已传入完整 tail 时直接使用该快照，不再重置。DP fetch 只迁移可恢复的缓存前缀，由管理器在目标端清零 tail；边界选择属于缓存流程。 - -这里的“最后一页”按本次命中前缀的逻辑顺序确定，与 CPU 物理页编号大小无关。全部页面先恢复 KV 和 linear 页数据，再由末页恢复 linear 运行态；中间长度无需逐个重置 tail，只有最终恢复长度需要满足 tail 的可重置边界。批量恢复多个请求时，调用方按请求分别调用，各自选择对应前缀的末页。 - #### full-attention 的聚合 token 复制 `copy_full_att_tokens(src_indices, dst_indices)` 的两个索引均为存储设备上的一维 CUDA 整数 Tensor；直接使用输入 dtype 和 stride，不转换设备、类型或连续布局。源和目标位置不重叠。该接口在本管理器持有的 full-attention 存储内部执行位置复制。参与集合在初始化时按共享 token 索引域确定，可以包含 main/draft 的不同布局，也可以包含同一模型的多个 FullAttentionState。管理器将同一组物理索引交给各参与 State 的 `copy_full_att_tokens`；每个 State 在自己的存储中从源位置复制到目标位置，不在不同模型之间互拷数据。 - 复制覆盖全部参与存储层的 KV、动态 scale 和 packed indexer key/scale 等随 token 变化的数据，保持各自存储表示；共享物理区域的别名按所有权去重。 - 源和目标为逐项对应的一维整数索引 Tensor，长度相等、目标唯一，且满足各布局的复制粒度与对齐要求；基础契约要求源目标不重叠。源数据须已就绪并保持有效，目标物理位置须已分配且允许写入。 -- 只搬运 full-attention token 数据，不分配或回收索引、不修改 ReqManager 映射，也不复制 linear 运行态、大小页快照或 AttentionTailState 的原始 tail。 - 调用前校验全部参与 State 的参数和能力；无 full-attention 参与者时明确报错。方法在当前 stream 提交并返回 None，全部参与复制完成前不能消费目标数据；中途失败时目标视为未就绪。 例如 main 和 draft 各有一个 FullAttentionState，复制物理位置 `[10, 11]` 到 `[20, 21]` 时，一次管理器调用会分别搬运 main 和 draft 各自的对应数据。共享的是位置变换，不要求两者具有相同存储布局。 @@ -997,7 +974,7 @@ bounded_att_big_page_allocator 分配一次 slot_id = 17 调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_runtime_from_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 -full-attention 历史 KV 和没有独立 pool 的 AttentionTailState 不属于这组状态页操作。需要同时恢复这些数据的调用方，使用相应资源操作，并保证历史 KV、bounded 快照和 tail 对应同一前缀且均已就绪。 +full-attention 历史 KV 和没有独立 pool 的 WindowCompressorState 不属于这组状态页操作。需要同时恢复这些数据的调用方，使用相应资源操作，并保证历史 KV、bounded 快照和窗口历史 对应同一前缀且均已就绪。 #### 保存与直接归还槽位的完整性 @@ -1016,7 +993,7 @@ GPU 数据操作遵守当前 stream 的依赖约定。调用方通过 stream/eve `layouts` 是按完整 `state_name`（包含 main/draft 角色）索引的 `StateTransferLayout` 映射，用于 CPU/PD 外部页。DP 不构造字段映射：manager 直接遍历本地 State，通过同名成员取得远端 State，将对应 token、页或请求索引交给具体实现。DP 不再维护映射缓存，也不通过 helper 转发。 -CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。CPU 读写由管理器按 full-attention、bounded 状态页、末页运行态恢复及 tail 重置分阶段组织，高频读写路径不记录流程日志；buffer 与字段搬运约束交给 State，管理器检查恢复所需的目标请求、非空前缀及完整页槽位数量。异步搬运的完成仍由调用方的 stream/event 保证。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 +CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。CPU 读写由管理器按 full-attention、bounded 状态页及末页运行态恢复分阶段组织，高频读写路径不记录流程日志；buffer 与字段搬运约束交给 State，管理器检查恢复所需的目标请求、非空前缀及完整页槽位数量。异步搬运的完成仍由调用方的 stream/event 保证。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 **CPU 联合页** @@ -1026,30 +1003,27 @@ CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能 - 格式包含 linear 页时，管理器写出接口分别接收 `big_page_slot_ids` 与 `small_page_slot_ids`；`cpu_page_indices` 前段对应大页、后段对应小页，每段与槽位列表一一对应。未参与的页池传 None 或空列表。full-attention KV 按完整页面列表写出一次，各 bounded State 再分别从大小页池写出快照；读回接口固定写入大页池，不接收 `page_kind`。管理器的 `big_page_slot_ids: List[int]` 只包含完整 CPU 页对应的普通大页槽位，不包含末尾碎页的临时槽位；没有 linear 页时为 None。槽位参数仅接受 Python 列表，不接受 Tensor。快照搬运不需要逐页长度数组。 - `cpu_page_ready_flags: List[bool] | None` 是本批 CPU 页的跳过写入标记，顺序与 `cpu_page_indices` 一致；True 保留整张联合页，None 表示全部写出。缓存层保留完整 token/page 序列并一次调用资源管理器，不逐页发起传输，也不通过过滤页 ID 改变 token 映射。资源层不查询或发布 ready 状态，全部 ready 时直接返回。 - FullAttentionState 将 ready 标记传给各自 CPU/PD 共用的传输入口；CPU 写出按 State 构造一次 GPU bool Tensor，全部字段复用。在 kernel 内按逻辑页屏蔽源 KV/量化字段读取和目标写入；若当前 block 没有待写元素，直接跳过后续索引读取、地址计算和数据搬运。跨页 block 混有 ready 与未就绪元素时保留逐元素 mask。PD 与读回不传标记、不分配此 Tensor，也不执行 ready 归约判断。bounded State 接收对应大小页段的 ready 列表，在取快照和构造复制视图前跳过 ready 页。 -- CPU 页不保存 tail 数据，写出接口不接收请求索引；布局只包含 full-attention 与 bounded-attention。 +- CPU 写出只读取大小页快照，不接收请求索引；无重叠的窗口 State 返回空区域，不占联合页。 - `read_req_from_cpu` 一次处理同一前缀的有序页面，将全部 KV 写到给定物理位置，将各页 linear 快照写入对应 pool 槽位。管理器根据 `token_indices.numel()` 和布局的页 token 数判断是否存在末尾碎页，将其完整 linear 快照写入预留 LOAD 大页；完整页写入 `big_page_slot_ids` 指定的大页槽位。读回后直接使用最后一页恢复 `req_idx` 的 linear 运行态，不修改调用方的槽位列表。仅含碎页时传 `big_page_slot_ids=[]`；只传状态的格式使用空 token 索引，按传入的普通槽位恢复。 -- 存在 bounded 或 tail 时，调用方提供单个目标 `req_idx`；管理器在全部页面读回后调用 `reset_attention_tail_state([req_idx])`。调用方保证最终前缀处于可重置边界。仅有 full-attention 时，`req_idx=None`。 -- 恢复终点必须对齐 tail 的可重置边界。最后一页可以仅使用大页的一部分，但仍须落在允许恢复的小页边界；不能恢复需要真实 tail 快照的非对齐位置。 +- 恢复终点须满足各 State 的快照边界；窗口压缩器要求完成一个 compress_stride，最后一页可以是大页的一部分。 - 任一资源读入失败时异常直接透传，已写入的数据不得作为就绪数据使用。写出失败时不得发布不完整联合页。 -例如恢复三个 CPU 页，终点长度依次为 `[P, 2P, 3P]`：全部 KV 和 linear 快照读回后，只对最终长度 `3P` 重置 tail，不逐页重置。管理器在重置 tail 前已将最后一页的 linear 快照恢复为运行态，调用方不再额外调用恢复接口。 +例如恢复三个 CPU 页，终点长度依次为 `[P, 2P, 3P]`：先读回各页 KV 和状态快照，再从最后一个大页槽位恢复请求在 `3P` 的运行态。无重叠窗口在该边界无需历史快照，后续计算使用新输入；调用方不再额外恢复一次。 **DP 直接复制** -当前 DP fetch 的来源是 radix cache 命中的前缀。缓存页大小须为每个 tail 池化步长的整数倍；GLM 启动时检查 `linear_att_hash_page_size % index_kpool == 0`。管理器复制 KV 与 bounded 运行态后清零目标 tail，不额外接收前缀长度。tail 不参与 DP 映射，也不读取源 tail；PD 的结束位置可能未对齐，仍传输完整快照。 +DP fetch 的来源是 radix cache 命中的前缀。缓存页大小须对齐压缩窗口的产出步长；GLM 启动时检查 `linear_att_hash_page_size % index_kpool == 0`。manager 接收远端管理器、CUDA token 索引、源/目标大页 ID 及源/目标请求索引，按同名 State 依次复制 token、大页快照与请求运行态。窗口 State 和其他 bounded State 使用同一流程，不再另设窗口重置阶段。 -`src_manager` 明确来源 DP/TP 存储实例，也可以是传输进程中已绑定共享存储的管理器。源和目标 token 索引等长、逐项对应；`src_req_indices`、`dst_req_indices` 仅接受 `List[int]`，成对提供且数量相等，用来迁移 prefill 基准 bounded 运行态；KDA/GDN 源游标必须为 0。有 tail 时直接清零目标请求的 tail，不传递额外长度。token 索引与请求槽位是不同的寻址维度，无需等长；前缀对应关系及可重置边界由调用方保证。 - -只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数时须包含所需 bounded 映射，或存在需要重置的 tail。一个目标分片需要多个来源时，调用方按来源分批调用，`AttentionStateManager` 将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 +KDA/GDN 的源游标须为 0；窗口 State 的两端配置和 MTP 容量须相同。复制操作不创建历史快照，调用方负责前缀对应关系、合法缓存边界、源数据就绪及引用有效期。不需要复制 token 或大页时传入空 Tensor 或空列表；指定请求的运行态仍会复制。 **PD 页读写** | `kind` | 必需参数 | 处理对象 | | --- | --- | --- | | `"kv"` | `token_indices`；`req_idx=None` | 按单页容量读写全部参与 full-attention 的 KV、scale 和 packed 字段 | -| `"att_state"` | `req_idx`；有 tail 时还需 `state_token_len`；`token_indices=None` | 读写 prefill 基准 bounded 状态与固定 tail 逻辑窗口，不经过大小页 pool | +| `"att_state"` | `req_idx`；`state_token_len`；`token_indices=None` | 读写 prefill 基准 bounded 状态（包括固定压缩窗口），不经过大小页 pool | -PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再转交 bounded 和 tail State，用于各自的窗口定位及恢复。full-attention 的 token 传输接口保持不变；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 +PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再转交各 bounded State，用于各自的窗口定位及恢复。full-attention 的 token 传输接口保持不变；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 @@ -1063,11 +1037,11 @@ CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，`Atten Model 和 InferStateInfo 通过 `att_state_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_gqa_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_state_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 -视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与 tail 的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 +视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与池化窗口的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 普通 KV 层可支持多种存储精度。推理层使用 `get_state(model_role=infer_state.att_state_role, state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState)` 取得资源。MLA 推理层和 GLM indexer 使用 `state_type_filter="mla"`，indexer 使用自己的 `layer_idx`。通用写入模板使用空字符串和 `FullAttentionState` 类型，不推断算法。匹配结果由管理器缓存，推理层和 indexer 不保存重复成员名，也不需要模型启动时逐层绑定。不同精度的 State 由逻辑层归属区分。 -Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 State 将共享消费者映射到源物理层。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按自身压缩比例访问对应角色的 tail 成员。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 +Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 State 将共享消费者映射到源物理层。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按窗口和步长访问对应角色的 WindowCompressorState。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 成员访问只读取 Python 属性，不复制 buffer，也不执行 GPU→CPU 同步。KV 写入直接调用该 State 的方法，packed KV/indexer 通过其层视图访问。DFlash 的特殊 KV 写入上下文同样设置管理器和 draft 标记。需要主机标量的 backend 在 State 初始化时准备标定镜像，避免 forward 中逐层执行 GPU→CPU 同步。 @@ -1083,7 +1057,7 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S #### 位置映射的范围 -传输需要分别描述逻辑层到物理层的映射、组合页中的资源区域、各字段的 TP 分片与副本关系，以及本次操作的 token/page 索引。KV、scale、Conv、SSM 和 tail 可以使用不同的 dtype、stride、对齐和轴划分，不能统一简化为一个 `offset` 或一个全局 `tp_rank`。 +传输需要分别描述逻辑层到物理层的映射、组合页中的资源区域、各字段的 TP 分片与副本关系，以及本次操作的 token/page 索引。KV、scale、Conv、SSM 和窗口 可以使用不同的 dtype、stride、对齐和轴划分，不能统一简化为一个 `offset` 或一个全局 `tp_rank`。 #### 信息由谁维护 @@ -1091,7 +1065,7 @@ Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 S | --- | --- | --- | | 本地存储布局与分片事实 | 具体 State 的配置、存储参数及内部 Tensor | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | | 传输字段及本地 TP 切片 | 具体 State 的 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` | 字段选择、全局形状、State 内字段的起止区间和各 rank 的读写切片 | -| 组合页布局及资源间分区 | 模型专属 `AttentionStateManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、tail 各区域的起止位置、对齐和整页容量 | +| 组合页布局及资源间分区 | 模型专属 `AttentionStateManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、窗口各区域的起止位置、对齐和整页容量 | | 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `StateTransferLayout`；按资源与 CPU/PD 路径分别保存 | | 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | @@ -1268,7 +1242,7 @@ DP 请求共享信息分为两个数组，均按请求在当前批次中的位 目标 DP 只复制从本地已有 KV 长度到来源 KV 长度之间新增的完整大页，不重复复制已有前缀页。先申请目标大页再申请 KV；目标大页分配必须成功，通过断言检查。复制完成后，将目标页按前缀长度登记到 `req.hybrid_len_to_big_page_id`，后续 radix 插入和释放沿用请求的现有归属管理。未跨越新大页边界时，大页列表为空。 -DP 副本之间使用相同的 State 成员名。`copy_from_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.state_name)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_runtime_state_from_dp`，再清零目标 tail。调用方保证源 MTP 游标为 0、前缀处于 tail 可重置边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 +DP 副本之间使用相同的 State 成员名。`copy_from_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.state_name)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_runtime_state_from_dp`。调用方保证源 MTP 游标为 0、前缀满足各 State 的缓存边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 full-attention 使用 `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_pages_from_dp` 与 `copy_runtime_state_from_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 @@ -1328,7 +1302,7 @@ byte 32768 : draft0_mla_fp8_state 的数据区域，按 draft 自身格式定 如果 draft 使用 TP 副本形式的 MLA，draft 区域由选定的 P 端副本写一次，D 端各副本读取该区域。它不沿用 main 的 head 偏移公式。main 与 draft 即使共享 token allocator，也分别取得自己的区域映射。 -CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索引数量确定、格式由外层描述；PD 使用 `kv` 与 `att_state` 两类任务，各自拥有布局，不强制把 KV、Conv、SSM 和 tail 放进同一页。采用按 TP rank 分块的格式时，必须明确其 TP 配置约束；跨 TP 配置恢复需要相应格式和转换能力。 +CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索引数量确定、格式由外层描述；PD 使用 `kv` 与 `att_state` 两类任务，各自拥有布局，不强制把 KV、Conv、SSM 和窗口 放进同一页。采用按 TP rank 分块的格式时，必须明确其 TP 配置约束；跨 TP 配置恢复需要相应格式和转换能力。 #### 初始化与提交时的校验 @@ -1410,7 +1384,7 @@ d_state.read_from_pd_page( # 所有目标资源/分片恢复完成后允许 decode。 ``` -调度层负责网络、同步和生命周期。方法返回不等于数据已经就绪，发送页在网络发送完成前不能复用，源槽位在读取完成前不能释放。KDA/GDN 状态的统一操作见 6.7 节，indexer tail 的操作见 6.8 节；它们共用适用的位置描述，不继承 `FullAttentionState`。 +调度层负责网络、同步和生命周期。方法返回不等于数据已经就绪，发送页在网络发送完成前不能复用，源槽位在读取完成前不能释放。KDA/GDN 状态的统一操作见 6.7 节，窗口压缩器的操作见 6.8 节；它们共用适用的位置描述，不继承 `FullAttentionState`。 ### 6.7 BoundedAttentionState：有界状态的统一操作接口 @@ -1562,11 +1536,11 @@ class BoundedAttentionState(ABC): ) -> None: ... ``` -bounded State 使用 `write_state_pages_to_cpu` / `read_state_pages_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复及 tail 重置。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `page_slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 +bounded State 使用 `write_state_pages_to_cpu` / `read_state_pages_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `page_slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 CPU 页读写的 `cpu_page_indices`、写出的 `page_slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`state_token_lens`、`page_slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。`state_token_lens` 表示各快照已经执行到的绝对 token 长度，不包含新采样但尚未执行的 token，也不是本批搬运的 token 数。调用方保证长度与选定快照一致；页池不新增长度记录。KDA/GDN 的 Conv/SSM 布局不依赖长度定位，保留此参数供 ring 等实现使用。 -PD 状态页按单请求处理，bounded 和 tail 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`page_slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 +PD 状态页按单请求处理，全部 bounded State 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`page_slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 #### 运行态与推理访问 @@ -1625,7 +1599,7 @@ CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字 #### 依赖、同步与使用示例 -资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 AttentionTailState；`AttentionStateManager`单独协调 tail 在恢复边界上的重置或迁移。 +资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 WindowCompressorState；管理器统一遍历 bounded State，各实现分别处理自身快照和运行态迁移。 GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 @@ -1647,7 +1621,7 @@ state.restore_bounded_att_runtime_from_pages( page_kind="big", page_slot_ids=page_slot_ids, ) -# target 请求在相关 KV、linear、tail 均就绪后继续执行。 +# target 请求在相关 KV、linear、窗口 均就绪后继续执行。 # 验证 forward 完成后，将接受步数按请求提交。 state.commit_bounded_att_state(req_indices, accepted_steps) @@ -1656,102 +1630,40 @@ state.commit_bounded_att_state(req_indices, accepted_steps) 各具体实现必须满足数据完整性、接受位置选择和跨布局复制约束。 -### 6.8 AttentionTailState：原始尾部窗口的操作接口 - -`AttentionTailState(ABC)` 约束注意力相关模块的尾部运行态操作，具体实现负责自身窗口、位置规则和搬运 kernel。例如 `KPoolAttentionTailState` 定义 GLM K-pool 的尾部布局。通用分类使用 `att_tail`，具体类及 `kpool_c4_bf16` 等资源标记保留算法信息。该分类不包含算子临时 workspace、完整递推状态或随历史增长的索引 key cache；对齐后可重置是 K-pool 的具体规则，不作为所有 tail 实现的共同假设。它不继承 `FullAttentionState` 或 `BoundedAttentionState`，也不持有 allocator、请求 KV 索引表或独立大小页 pool。 - -#### KPoolAttentionTailState 布局与状态边界 - -设本实例层数为 `L`、请求容量为 `N`（包含 HOLD）、池化步长为 `P`、MTP 额外容量为 `S`、index head 维度为 `D`，原生 ring 布局为 `[L, N, P+S, 2*D]`，末维分别存 raw key 与 gate score。GLM K-pool 示例取 `P=4`、`D=128`，具体实现须声明其 kernel 支持的配置范围。 - -计算先用此前 tail 和本轮 raw 输入完成跨 chunk 池化，再更新 tail。绝对位置 `t` 写入 `t % (P+S)`。完成池化的压缩 key/scale 写入 full-attention packed buffer，由 `FullAttentionState` 管理；tail 保留增量计算所需的原始数据。 - -PD 只传输 prefill 结束时最近 `P` 个逻辑位置，不传输完整 `P+S` 物理 ring。CPU cache 与 DP fetch 恢复对齐的缓存前缀，直接重置目标 tail。D 端仍按本地 `P+S` 分配 ring，保留后续 MTP 验证、拒绝回退和 draft 回访所需的候选历史。PD 逻辑窗口仅用于 prefill 状态交接，不用于迁移已经执行 MTP 的运行态；前缀恢复的可重置边界由模型配置和缓存流程保证。 - -#### 必须实现的接口 - -构造函数直接接收配置、请求容量 `request_capacity`、MTP 步数 `mtp_step` 和分配设备 `device`，不再使用 `TailStorageSpec`。请求容量包含 HOLD 槽位；容量与 MTP 步数各在 State 上保存一份，设备从 `tail_buffer.device` 获取。K-pool 在构造阶段分配完整的 `pool_size + mtp_step` ring，不依赖后续 KV token 容量规划。 - -基类接口仅包含类型标记、层访问、边界重置及 PD 区域/读写；构造参数由各具体 State 和算法 create 入口定义，不要求统一使用 KPoolStateConfig。CPU cache 不保存 tail，CPU/DP 前缀恢复统一调用重置接口；tail 不提供 CPU 读写或 DP 复制接口。 - -```python -from abc import ABC, abstractmethod -from typing import Any - - -class AttentionTailState(ABC): - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - @abstractmethod - def get_state_type(self) -> str: ... - - # ======================================================================== - # 二、Attention tail：计算视图与无快照重置 - # ======================================================================== - @abstractmethod - def contains_logical_layer(self, logical_layer_id: int) -> bool: ... - - def get_layer_view(self, logical_layer_id: int) -> Any: ... - - @abstractmethod - def reset_attention_tail_state(self, req_indices: list[int]) -> None: ... - - # ======================================================================== - # 三、PD 分离传输 - # ======================================================================== - @abstractmethod - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... - - @abstractmethod - def write_to_pd_page( - self, req_idx, pd_page_buffers, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, - ) -> None: ... - - @abstractmethod - def read_from_pd_page( - self, pd_page_buffers, req_idx, *, pd_page_index: int, state_token_len: int, layout: StateTransferLayout, - ) -> None: ... -``` - -重置接口只接收请求索引列表并清零对应 tail，不接收前缀长度。PD 按单请求和显式 `pd_page_index` 搬运固定逻辑窗口，`state_token_len` 表示该 State 已完成 prefill 的前缀长度。请求槽位由 ReqManager 分配，State 不保存请求长度表。 - -#### 容量、访问与恢复 - -| 接口 | 语义 | -| --- | --- | -| 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | -| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | -| `get_layer_view(logical_layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | -| `reset_attention_tail_state(req_indices)` | 清零指定请求的 tail，适用于新请求或可重置的缓存前缀；K-pool 清零完整 ring，其他实现清零对应数据与控制信息。模型配置和缓存流程负责保证边界，接口不接收长度 | - -`get_layer_view` 供计算 backend 直接使用；池化与写 tail 仍在模型 kernel 中完成，不增加通用的 `append_token` 或 attention forward 包装。执行顺序必须满足“先池化、再更新 tail”。 - -`reset_attention_tail_state` 是新请求/前缀恢复入口，不是 MTP 回滚入口。即使已接受长度恰好对齐 P，运行中的 draft 重访仍可能需要 ring 内历史，不能用该条件代替 MTP 状态处理。模型初始化与缓存流程分别保证 main/draft 的恢复条件；调用方还需确认对应 full-attention、linear 等资源全部恢复到约定位置。 - -#### PD：传输 prefill 逻辑窗口并适配本地 ring - -| 接口 | 数据路径及约束 | -| --- | --- | -| `write_to_pd_page(...)` | 按 `state_token_len` 提取指定请求的逻辑窗口 → `pd_page_index` 指定的 PD 状态页区域，使用 layout 的资源偏移与 TP 写入规则 | -| `read_from_pd_page(...)` | 已接收的逻辑窗口 → 按 `state_token_len` 和本地 ring 容量映射到目标槽位；窗口外的 MTP 槽位保持原值，调用方建立接收完成依赖 | - -这些方法传输 tail 自身的 prefill 逻辑窗口，不承担历史 KV、池化 key/scale、bounded 状态或 MTP proposal 的迁移。发送端取最近 `P` 个位置并按时间顺序排列；接收端同样按逻辑窗口遍历，单次 kernel 恢复窗口数据及短前缀补零，窗口之外的 MTP 槽位与其他请求保持不变。后续 MTP 候选历史由 D 端执行过程中产生。 - -`StateTransferLayout` 描述外部页位置与切片。K-pool region 的轴为 `("layer", "window", "channel")`，形状为 `[layers, P, 2*head_dim]`，与本地 `mtp_step` 无关。`window` 按时间顺序对应 `[state_token_len-P, state_token_len)`；短前缀中小于 0 的位置在发送和恢复时均补零，窗口外槽位保持原值。独立状态页的 `tokens_per_page=None`，该窗口不是 full-attention 的页内 token 轴。 - -K-pool 的 `window` 切片设置 `local_circular=True`。`kpool_transfer.py` 从 `state_token_len` 计算本次本地窗口起点；可通过 `axis_mapping_overrides` 覆盖本次轴映射，不修改缓存的 layout。窗口起点及短前缀的无效位置数量作为运行时参数传入定制 Triton kernel，不随请求位置编译新变体。写出和读回各一次 launch，处理量均为 `layers * pool_size * channels`，不遍历窗口之外的 MTP 槽位；连续 channel 在对齐允许时使用最大 8 字节整数视图搬运，保留原始位模式。 - -P/D 使用相同的模型及 main/draft 资源集合，但允许不同的 `mtp_step`。位置 `pos` 在 P 端从 `pos % (P+S_p)` 读取，在 D 端写入 `pos % (P+S_d)`；不能直接截取 ring 的前 P 个物理槽位。模型算法、池化步长、层顺序、字段精度和逻辑维度仍须兼容。现有 main/draft prefill 保持相同的序列位置，因此可共用任务的 `end_kv_index`；位置变化不影响缓存的布局描述。 - -例如 `P=4`、`S_p=0`、`S_d=2`、前缀长度 11，外部窗口保存位置 `[7,8,9,10]`；P 端读取槽位 `[3,0,1,2]`，D 端恢复到 `[1,2,3,4]`。D 端其余 ring 槽位保持原值。PD 块仍使用 KV 页确定的固定容量，减少 tail 有效区域和本地复制量不代表网络整页传输字节数同步减少。 - -长度、前缀身份和格式来自保存/传输控制元数据，不能从原始 tail 字节反推。调用层负责核对元数据并维持 main/draft 的一致恢复点；State 负责布局和数据操作。所有 GPU 操作遵循当前 stream，返回不表示完成；源请求与目标页在读写完成前不能复用。 - -#### 与 linear 页接口和 MTP 提交的区别 - -- 不提供 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages`：tail 没有独立大小页 pool。CPU/DP 缓存前缀在对齐边界通过 reset 恢复；非对齐的 PD 交接使用 prefill 逻辑窗口。 -- 不要求 `commit_bounded_att_state`：K-pool 没有类似 KDA 的 SSM 选择游标，按绝对位置访问 ring，MTP 拒绝后由执行流程选择后续位置。增加空的提交方法不能替代窗口容量与重访范围的正确性；未来出现确实需要提交控制状态的实现时，再扩展对应契约。 +### 6.8 WindowCompressorState：窗口候选与 gate + +K-pool 与 V4 的压缩器共用 `bounded_att/WindowCompressorState`，具体池化、归一化、 +RoPE 和量化仍由各模型自己的计算算子完成。资源层只负责窗口存储、快照及传输。 +模型通过 `create_bounded_att_states()` 返回窗口 State，与 KDA/GDN、SWA 使用同一管理流程。 + +| 场景 | window_size | compress_stride | candidate_dim | 缓存边界历史 | +| --- | ---: | ---: | --- | --- | +| GLM K-pool | 4 | 4 | indexer key 的 head_dim | 无重叠，无需快照 | +| V4 CSA KV / indexer | 8 | 4 | 输出向量维度的两倍 | 保存最近 4 个位置 | +| V4 HCA KV | 128 | 128 | 输出向量维度 | 无重叠，无需快照 | + +`candidate_dim` 不包含 gate；候选和 gate 具有相同形状和存储精度，由同一个 `dtype` +描述。GLM K-pool 随模型精度配置（通常 BF16），V4 使用 FP32。运行态保持一个 Tensor，形状为 +`[layer, request, window_size + mtp_step, 2 * candidate_dim]`，末轴前半为候选、后半为 gate。 +`get_layer_view(logical_layer_id)` 返回该层的拼接视图，计算算子的输入方式保持不变。 +`target="kv"` 表示 attention KV 候选，`target="indexer"` 表示索引 key 候选， +仅区分 State 成员名，不选择模型算法。例如 K-pool 使用 `compressor_indexer_w4_s4_bf16`。 + +大小页保存和恢复必须对齐 `compress_stride`。有重叠时，通过 `WindowPagesPool` 保存 +时间有序的重叠历史，CPU 页搬运同一份拼接字段；无重叠时不分配页池、不占 CPU 页， +整组缓存边界后的计算只依赖新输入。空前缀仍通过 `reset_bounded_att_state()` 初始化， +候选清零，gate 设为负无穷。DP 运行态复制使用相同配置和 MTP 容量的物理 ring。 + +管理器只为声明了 CPU 快照区域的 bounded State 调用 `initialize_page_pools()`。 +窗口 State 在返回布局前断言小页 token 数可被 `compress_stride` 整除；大页由整数个小页组成， +因此也满足此边界。配置了 CPU 页 token 数时，同样检查该页边界，即使窗口没有快照区域也不能跳过。 +共享槽位 allocator 继续供缓存引用和其他 State 使用；不需要快照的 State 不分配对应的页池 Tensor。 + +PD 在任意前缀长度导出完整 `window_size` 逻辑窗口，例如 K-pool 的长度 6 前缀仍需 +携带未完成池化组的原始 key/gate。P/D 的 MTP 容量不必相同：搬运按绝对前缀长度映射 +到各自的 ring，不传输或清理逻辑窗口外的候选槽位。环形搬运 kernel 位于 +`bounded_att/window_compressor_transfer.py`，一次 launch 搬运全部层的候选与 gate。 +MTP 回退依赖保留的历史和绝对位置，不调用缓存恢复入口。 ## 7. 资源初始化与调用约定 @@ -1763,34 +1675,32 @@ BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己 | --- | --- | | `get_req_manager_max_seq_len()` | 返回请求映射所需长度容量，取本模型 `batch_max_tokens`、`max_seq_length` 与 0 的最大值；未配置的值按 0 处理，不表示模型最大上下文能力 | | `bind_req_manager(req_manager)` | 绑定调用方创建的共享实例，不判断 main/draft，也不通过 `main_model` 获取资源 | -| `prepare_attention_states()` | main 直接创建空的 `AttentionStateManager`，draft 引用 main 的实例；在 KV 内存分区内调用三类 State 工厂并统一注册返回的列表 | +| `prepare_attention_states()` | main 直接创建空的 `AttentionStateManager`，draft 引用 main 的实例；在 KV 内存分区内调用两类 State 工厂并统一注册返回的列表 | | `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | 1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 -2. backend 按 main、各 draft 的顺序分别调用 `prepare_attention_states()`。公共准备流程为 main 创建唯一管理器；各 draft 的准备入口直接引用 main 的管理器，公共流程调用其 `create_full_att_states()`、`create_bounded_att_states()`、`create_att_tail_states()`，再通过 `register_state` 注册返回的全部 State。每个模型一次准备所需的 full-attention meta State、GDN/KDA 固定运行态和 indexer tail，各 State 的 config 保存容量规划与传输所需的完整布局信息。此阶段只准备布局，不读取静态标定文件；模型工厂不再接收或计算标定偏移。draft 可以拥有自己的 linear attention，不从 main 的 State 结构推导。main 和 draft 分别确定 KV head 数、head dimension、编码和层号，不提前在 main buffer 中预留 draft 层，也不在 Model 上保存 State 引用。 3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_registration()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 必须有 K/V 标定，Q 标定可选：有 Q 标定时使用静态量化，未提供时按本次请求的 Q 分组动态量化，prefill/decode 使用同一规则。随后 backend 调用 `initialize_storage(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 -BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `state_name`。 +BaseModel 提供两个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `state_name`。 | 创建接口 | 返回值与职责 | | --- | --- | | `create_full_att_states(...)` | 返回 `list[FullAttentionState]`；只准备布局，不加载静态标定或分配 KV 容量；可返回多种量化或存储方案 | | `create_bounded_att_states(...)` | 返回 `list[BoundedAttentionState]`；按请求容量创建运行态，大小页延后分配；CPU 查询使用 `device="meta"` | -| `create_att_tail_states(...)` | 返回 `list[AttentionTailState]`；按请求容量创建本模型 tail，CPU 联合页不调用此接口 | `model_logical_layer_ids` 表示当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；main/draft 保留现有全局编号机制。具体模型自行选择各类 State 的参与层和共享分组，该参数不代表物理存储层数，也不要求每个 State 使用全部层。 `Qwen3NextResourceConfig` 及其 GLM 配置只描述当前模型，保留 `from_model_config()`、`load_from_args()` 和 `get_full_att_layer_indices()` 三个入口。层查询返回模型内部序号，由模型工厂转换为实际逻辑层 ID；不追加 draft 层，不计算物理层映射、Conv/SSM 存储形状或字节数。配置只保存全局 head 数，本地 TP 分片由 State 或后端校验调用方计算。full-attention 维度保存原始特征宽度，GLM 的 indexer 字段和 packed padding 由具体 State 的创建入口处理。 -三个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded/tail 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_attention_states()` 为 main 直接创建空的 `AttentionStateManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttentionStateManager`，模型差异由三类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 +两个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_attention_states()` 为 main 直接创建空的 `AttentionStateManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttentionStateManager`，模型差异由两类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_attention_states()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig` 或 `MLAAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded/tail 工厂同样直接调用算法 create；GLM main/draft 分别选择 tail 的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 +`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig` 或 `MLAAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded 工厂同样直接调用算法 create；GLM main/draft 分别选择池化窗口的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 `AttentionStateManager.initialize_storage()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `allocate_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 @@ -1800,13 +1710,13 @@ BaseModel 提供三个默认返回空列表的 `classmethod` 扩展点，具体 PD 容量规划使用 `manager.get_pd_page_bytes() * page_count` 计算预留空间，依赖方向是先布局、后页大小。FullAttentionState 的 `get_pd_transfer_regions()` 根据全局 `pd_kv_page_size` 返回实际字段区域，包含各 field 的 dtype 对齐及实现要求的内部 padding。manager 同时平移字段区间的起点与终点，以 `max(region.byte_end)` 确定当前 State 的数据结束位置，向上对齐到 8 字节后记录为 layout 的 `byte_end`，并作为下一个 State 的 `byte_start`。所有 State 排列完成后，`max(layout.byte_end)` 就是 KV 页大小。GPU packed 行中未被任何传输字段覆盖的 padding 不进入 PD 页。 -bounded/indexer 的运行态布局使用同一规则：按字段的实际最大结束位置确定 State 占用，不另外维护纯数据字节数估算，也不累加多个映射共用的同一区域。它始终返回各运行态 State 自身的起止区间，不依赖 full-attention 布局，运行态总占用取 `max(layout.byte_end)`。`get_pd_page_bytes()` 返回 KV 页与完整运行态页占用的最大值；没有参与资源时返回 0。有 KV 且 KV 页小于运行态页时记录英文 warning，自动扩大实际传输页，保持 `pd_kv_page_size` 的 token 数和两类布局偏移不变，不再因两类页大小关系而中止初始化。PD 调用方按这个返回值分配和选取实际传输页，再以该页起点执行对应任务的布局；例如运行态布局为 80 字节、KV 页为 128 字节时，实际传输块为 128 字节，运行态仅操作自己的 80 字节范围。 +bounded（包括压缩窗口） 的运行态布局使用同一规则：按字段的实际最大结束位置确定 State 占用，不另外维护纯数据字节数估算，也不累加多个映射共用的同一区域。它始终返回各运行态 State 自身的起止区间，不依赖 full-attention 布局，运行态总占用取 `max(layout.byte_end)`。`get_pd_page_bytes()` 返回 KV 页与完整运行态页占用的最大值；没有参与资源时返回 0。有 KV 且 KV 页小于运行态页时记录英文 warning，自动扩大实际传输页，保持 `pd_kv_page_size` 的 token 数和两类布局偏移不变，不再因两类页大小关系而中止初始化。PD 调用方按这个返回值分配和选取实际传输页，再以该页起点执行对应任务的布局；例如运行态布局为 80 字节、KV 页为 128 字节时，实际传输块为 128 字节，运行态仅操作自己的 80 字节范围。 PD master 负责校验连接参数，长度通过传输任务的 `end_kv_index` 传递；资源层不附加 PD 页布局指纹及长度、版本字段。两个布局 getter 与页大小查询均使用 `lru_cache`，存储初始化后一起清除缓存。 P/D 传输进程为每次 KV 任务创建普通 CPU token 索引 Tensor，在 `copy_cuda_stream` 中调用 `cuda(non_blocking=True)` 上传，后续 State 拷贝沿用同一流；预热使用相同路径。不预分配或保留 pinned 索引页缓冲，普通 CPU 内存也不保证 H2D 全程异步。`att_state` 任务按请求槽位搬运，不创建或上传 token 索引。 -CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_page_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。Attention tail 不参与 CPU 联合页，读回时按最终前缀边界重置；快照接口仅用于 PD。 +CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_page_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。无重叠窗口不参与 CPU 联合页；有重叠窗口保存必要历史，PD 则传输完整窗口。 各具体 State 的 `get_cpu_transfer_regions()` 使用 `@lru_cache(maxsize=None)` 缓存不可变区域元组， 并在其外层添加 `@deepcopy_result`，每次先取得缓存结果、再返回深拷贝，调用方不持有缓存原对象。 @@ -1855,13 +1765,13 @@ CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 fu ### 7.2 缓存保存与恢复 -没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)` 和 `reset_attention_tail_state(req_indices)`，分别初始化 bounded 与 tail 运行态。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 +没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)`，初始化全部 bounded 运行态（包括窗口压缩器）。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_runtime_to_pages`。`AttentionStateManager` 将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 -radix 沿用原有判断：大页节点具有状态，小页节点通过 `small_page_buffer_idx is not None` 判断是否可恢复；图片边界仍通过 `clamp_mm_cache_len` 独立检查。匹配得到节点引用后，调用方直接调用 `restore_bounded_att_runtime_from_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_attention_tail_state(req_indices)` 清零 tail，保留已恢复的 linear 状态。调用方还需确认历史 KV、tail 和 bounded 快照均对应恢复前缀，并保证数据就绪。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 -使用 K-pool 的 tail 在完整池化边界上通过 `reset_attention_tail_state` 重置；非对齐恢复需要原始 tail 数据。CPU 读回后，管理器在最终前缀数据恢复后调用 `reset_attention_tail_state`，不读取 tail 快照。DP fetch 在对齐的缓存前缀上重置 tail；PD 按第 6.8 节传输固定逻辑窗口并重排到本地 ring。前缀恢复的清零规则不用于 MTP 回退。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 MTP 接受位置由模型执行流程选择，供后续计算使用；当前大小页保存和 DP/PD 导出发生在 prefill 基准状态阶段。KDA/GDN 只保存实际捕获的状态；跨越但未保存的中间位置不可作为有效页，也不能发布为 CPU 缓存页；执行流程需在待缓存页边界安排真实状态捕获。管理器的数据操作不承担一次 forward 的提交或请求长度更新。 @@ -1890,7 +1800,6 @@ CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合 - full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 - `KDAState`、`GDNState` 实现 `BoundedAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移 prefill 基准请求状态；管理器组织参与对象及位置映射。 -- `AttentionTailState` 提供层访问、按请求清零及 PD 直接读写方法，详见 6.8 节；CPU cache 和 DP fetch 通过重置恢复对齐前缀的 tail。GLM K-pool indexer 在 TP ranks 间为副本，其他实现自行描述其布局。 - 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttentionStateManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 跨进程导入的管理器持有借用的运行态存储和索引容量描述，不持有 allocator 或 CPU 状态池。请求映射由 server 层作为 ReqManager 所有的独立存储导出，不能附加到资源管理器。CUDA IPC 句柄按实际读取进程数发布，每个读取者消费独立引用；原进程在全部读取者结束使用前保持存储存活。 @@ -1902,15 +1811,15 @@ PD 任务分为 `kv` 和 `att_state` 两类；涉及 main/draft 时按模型角 | 任务 | 组成 | 处理方式 | | --- | --- | --- | | `kv` | full-attention 历史 KV，包含布局所需的 packed indexer 数据 | 按 token 页复制 | -| `att_state` | linear Conv/SSM + 可选 indexer tail | 模型专属 `AttentionStateManager` 按请求恢复位置调用对应 State 方法 | +| `att_state` | Conv/SSM、SWA 和窗口压缩器运行态 | 模型专属 `AttentionStateManager` 按请求恢复位置调用对应 State 方法 | -P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，indexer tail 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将 tail 写入目标 TP 各 rank 的本地请求槽位。 +P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，索引池化窗口 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将窗口写入目标 TP 各 rank 的本地请求槽位。 -模型专属 `AttentionStateManager` 实现状态页的布局组织和容量校验，调用 linear State 与 indexer tail State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 +模型专属 `AttentionStateManager` 实现状态页的布局组织和容量校验，调用 linear State 与 窗口压缩器 State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 -网络层只接收固定容量的字节页，不按 KV 的 head 数或 dtype 推断组合布局。两端在创建传输句柄前核对页容量，PD master 在连接时校验任务类别、资源身份、算法与编码、层映射和全局字段形状；页内不附加格式指纹、长度或版本。布局使用全局字段描述以支持不同 TP 大小，tail 布局只描述固定逻辑窗口，不包含本地 MTP 候选容量。运行态的前缀长度由传输任务提供。 +网络层只接收固定容量的字节页，不按 KV 的 head 数或 dtype 推断组合布局。两端在创建传输句柄前核对页容量，PD master 在连接时校验任务类别、资源身份、算法与编码、层映射和全局字段形状；页内不附加格式指纹、长度或版本。布局使用全局字段描述以支持不同 TP 大小，窗口布局只描述固定逻辑窗口，不包含本地 MTP 候选容量。运行态的前缀长度由传输任务提供。 -P/D 使用相同模型和 main/draft 资源集合，允许不同 TP 与 `mtp_step`。full-attention 传输有效 KV 和量化字段；bounded 状态按 6.7 节传输 prefill 基准 Conv/SSM，恢复到本地基准位置并将游标归零；tail 按 6.8 节传输固定逻辑窗口并按本地 ring 容量重排。MTP 候选容量不进入 PD 字段形状，单个 task 仍使用既有的前缀位置描述。 +P/D 使用相同模型和 main/draft 资源集合，允许不同 TP 与 `mtp_step`。full-attention 传输有效 KV 和量化字段；bounded 状态按 6.7 节传输 prefill 基准 Conv/SSM，恢复到本地基准位置并将游标归零；窗口压缩器按 6.8 节传输固定逻辑窗口并按本地 ring 容量重排。MTP 候选容量不进入 PD 字段形状，单个 task 仍使用既有的前缀位置描述。 ### 8.3 CPU/Disk 与 DP fetch @@ -1918,11 +1827,11 @@ CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布 GPU→CPU 保存状态页完成后，才能在 CPU 上重排或复制这些快照；主机侧读取必须等待保存完成，不能只向另一个 CUDA stream 排入依赖。CPU 页的 ready 发布仍由调用流程在全部参与资源与 TP 分片传输完成后执行。 -模型专属 `AttentionStateManager` 实现 CPU 联合页的布局与读写组织，调用所需 main/draft State 的 CPU 方法，组合 full-attention KV 和 linear 页；按调用方提供的槽位写出或读回快照，pool 和 CPU buffer 均不保存长度副本。CPU 缓存调度提供外部页和任务生命周期；每个已发布 CPU 页都具备该模型组合所需的有效 linear 状态。一次读回处理同一前缀的页面，并通过 `reset_attention_tail_state` 清零目标 tail；恢复终点必须满足具体 tail 的可重置边界。 +缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大页 pool 分别提供存储。各参与方的暂存数据消费完成后，后续请求才能复用该槽位。磁盘存储使用 CPU 页格式；main/draft 异构组合需定义各自资源段的位置、字段形状及 TP 切片。 -DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并在对齐边界重置目标 tail。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 +DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并按各 State 的规则恢复窗口历史。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 ### 8.1 DeepSeek V4：压缩 KV 与有界窗口 @@ -1936,7 +1845,7 @@ State 创建、backend 和提前预留 workspace 共用这些配置入口。编 `get_attention_view()` 均由对应原子配置实现，backend 不再维护一套物理格式计算。 SWA 使用独立的 `FP8E8M0SlidingWindowConfig`，只记录逻辑层、近期窗口、KV 编码与 TP 信息。 压缩器使用独立的 `WindowCompressorStateConfig`，明确记录 `window_size`、`compress_stride`、 -`candidate_dim`、`target`、精度和 TP 信息；没有 C4/C128 层列表或 SWA 字段。 +`candidate_dim`、`target`、`dtype` 和 TP 信息；没有 C4/C128 层列表或 SWA 字段。 `window_size` 是每次聚合覆盖的原始 token 数，`compress_stride` 是每隔多少原始 token 产出一条， 大小页需保留的跨步历史由 `overlap_tokens = window_size - compress_stride` 推导，不额外存储重复配置。 @@ -1992,7 +1901,7 @@ C4 保存上一组的四个位置，供下一组重叠压缩使用。运行态 PD 运行态采用与本地 MTP 容量无关的逻辑窗口：SWA 为配置窗口，C4/indexer 为八个位置， C128 为 128 个位置。收发两端根据传入的绝对前缀长度各自计算 ring 位置，因此允许不同的 MTP 配置。 SWA、压缩 KV 和 compressor 在 TP 内完整复制，仅 rank 0 写出共享外部区域，各 rank 可读回。 -V4 的这些窗口需要真实快照，不归入仅能在页边界清零的 `AttentionTailState`。 +窗口是否需要真实快照由重叠历史决定，GLM K-pool 和 V4 压缩器使用相同的 bounded State 接口。 checkpoint 中的权重层号与 draft 自有逻辑层数分别记录；多个复用相同 MTP 权重的 draft 仍注册独立 State。 模型和 backend 根据逻辑层查询 State 视图，不在 ReqManager 中附加 ring，不在 backend 中保存运行态副本。 @@ -2002,15 +1911,14 @@ checkpoint 中的权重层号与 draft 自有逻辑层数分别记录；多个 | 问题 | 设计约定 | 待确定内容 | | --- | --- | --- | | main/draft 管理器装配 | 模型组合实现自己的管理器；各角色和具体方案拥有唯一 State，同一分配域共享唯一 allocator 实例 | 新模型的 main/draft 装配时序及独立存储分配；调用方覆盖全部参与资源的引用解除及单次归还责任 | -| AttentionStateManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、清零指定请求的 tail、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | +| AttentionStateManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | | 内部 attention 资源的暴露 | 具体模型在初始化阶段绑定 State 或层视图，不要求统一模型视图获取接口 | 各模型/backend 的具体绑定方式与按层访问结构 | | FullAttentionState 接口与继承 | 必需语义：容量计算、初始化、层视图、KV 写入、完整 token 复制及固定字段传输布局；CPU/DP/PD 由子类按支持路径直接实现读写方法 | 新布局的 TensorField 覆盖范围、StateTransferLayout 映射约束及具体搬运实现 | | BoundedAttentionState 接口与继承 | KDAState、GDNState 实现统一抽象操作；已接受运行态、checkpoint 和传输分别定义，allocator 仍归管理器 | 按 6.7 节细化各子类的描述字段和 kernel，核对 MTP 状态选择、跨 TP/CPU 格式兼容性 | -| AttentionTailState 接口与继承 | 按 6.8 节实现层访问、按请求清零及 PD 传输；prefill 逻辑窗口迁移，无独立页 pool | 各子类的恢复规则及跨 ring 布局转换能力 | | allocator 能力描述 | 公共入口按需初始化，不需要的 allocator 为 None | 恢复边界和传输能力使用哪些最小描述，避免重复配置导致不一致 | | checkpoint 引用形式 | 调用方维护大小页类别、共享槽位 ID、前缀长度及参与模型资源；各侧 pool 只保存快照数据 | 引用对象或参数组合的具体类型，以及参与资源集合的表达 | | 配置类型拆分 | main/draft 分别产生实现和布局参数，共享 allocator 接收统一容量与索引配置 | 各类 State 的配置字段、跨资源共享字段与校验责任 | -| PD tail 的协议布局 | 固定 P 个逻辑位置，按任务前缀长度映射两端 ring；兼容不同 MTP 步数 | 保持相同模型资源集合与 prefill 状态交接约束 | -| CPU 缓存页的组合方式 | 每页包含全部有效 linear 状态；管理器组织读写，并在全部页面恢复后清零目标 tail | 不同 main/draft 实现的资源段布局及恢复边界兼容 | +| PD 窗口压缩器的布局 | 固定 P 个逻辑位置，按任务前缀长度映射两端 ring；兼容不同 MTP 步数 | 保持相同模型资源集合与 prefill 状态交接约束 | +| CPU 缓存页的组合方式 | 每页包含全部有效 linear 状态；管理器组织读写，并在全部页面恢复后从末页恢复请求运行态 | 不同 main/draft 实现的资源段布局及恢复边界兼容 | 以上待细化事项须遵守本文的资源归属、共享分配域和恢复语义约定。 diff --git a/lightllm/common/attention_state/att_tail/__init__.py b/lightllm/common/attention_state/att_tail/__init__.py deleted file mode 100644 index 05d1c630e1..0000000000 --- a/lightllm/common/attention_state/att_tail/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .create_kpool_state import create_kpool_state -from .base_state import AttentionTailState -from .kpool_state import KPoolAttentionTailState, KPoolStateConfig - -__all__ = ["create_kpool_state", "AttentionTailState", "KPoolAttentionTailState", "KPoolStateConfig"] diff --git a/lightllm/common/attention_state/att_tail/base_state.py b/lightllm/common/attention_state/att_tail/base_state.py deleted file mode 100644 index e3e300cd87..0000000000 --- a/lightllm/common/attention_state/att_tail/base_state.py +++ /dev/null @@ -1,92 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any -from abc import ABC, abstractmethod - -from lightllm.common.attention_state.types import ( - StateTransferLayout, - TransferRegion, -) - -if TYPE_CHECKING: - import torch - - -class AttentionTailState(ABC): - """注意力相关模块的尾部状态接口，保存增量计算所需的尾部数据及控制信息。""" - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - @abstractmethod - def get_state_type(self) -> str: - """返回算法与存储方案标记，用于生成资源成员名。 - - 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， - 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 - 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 - 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 - 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 - - K-pool 格式为 kpool_c{压缩比例}_{存储精度}，例如 kpool_c4_bf16， - 对应 main_kpool_c4_bf16_state；同压缩比例的不同存储精度使用不同名称。 - 新实现使用具体算法及必要方案参数，不仅返回 indexer 或 tail 这种通用名称。 - """ - ... - - # ======================================================================== - # 二、Attention tail：计算视图与无快照重置 - # ======================================================================== - - @abstractmethod - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" - ... - - @abstractmethod - def get_layer_view(self, logical_layer_id: int) -> Any: - """按逻辑层号返回尾部状态计算所需的存储视图。""" - ... - - @abstractmethod - def reset_attention_tail_state( - self, - req_indices: list[int], - ) -> None: - """清零指定请求的 tail；调用方保证用于空前缀或可重置的缓存边界，不用于 MTP 回滚。""" - ... - - # ======================================================================== - # 三、PD 分离传输：prefill 逻辑窗口 - # ======================================================================== - - @abstractmethod - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回独立于本地 MTP 容量的 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" - ... - - @abstractmethod - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """按 state_token_len 对应的 prefill 终点导出逻辑窗口，不导出推测候选历史。""" - ... - - @abstractmethod - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """按 state_token_len 将逻辑窗口恢复到本地运行态，适配本地 MTP 容量。""" - ... diff --git a/lightllm/common/attention_state/att_tail/create_kpool_state.py b/lightllm/common/attention_state/att_tail/create_kpool_state.py deleted file mode 100644 index b4158e3785..0000000000 --- a/lightllm/common/attention_state/att_tail/create_kpool_state.py +++ /dev/null @@ -1,22 +0,0 @@ -"""K-pool attention tail State 的统一创建入口。""" - -from __future__ import annotations - -import torch - -from .kpool_state import KPoolAttentionTailState, KPoolStateConfig - - -def create_kpool_state( - config: KPoolStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", -) -> KPoolAttentionTailState: - """按池化配置和请求容量创建 K-pool 尾部运行态，包含 MTP 额外槽位。""" - return KPoolAttentionTailState( - config=config, - request_capacity=request_capacity, - mtp_step=mtp_step, - device=device, - ) diff --git a/lightllm/common/attention_state/att_tail/kpool_state.py b/lightllm/common/attention_state/att_tail/kpool_state.py deleted file mode 100644 index 6c50c7cc0b..0000000000 --- a/lightllm/common/attention_state/att_tail/kpool_state.py +++ /dev/null @@ -1,174 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING -from dataclasses import dataclass - -import torch - -from ..utils import dtype_name -from ..types import AxisMapping, TransferRegion -from .base_state import AttentionTailState -from .kpool_transfer import transfer_kpool_pages - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import StateTransferLayout - - -@dataclass(frozen=True) -class KPoolStateConfig: - # 本 State 管理的逻辑层号；当前每层独占一份 tail，列表顺序对应物理层轴，不要求编号连续。 - logical_layer_ids: tuple[int, ...] - # K-pool 的池化步长/压缩比例，单位为原始 token；例如 4 表示每 4 个 token 形成一个池化组。 - # 无快照前缀恢复要求长度对齐该步长；ring 的实际槽位数为 pool_size + State.mtp_step。 - pool_size: int - # 单个原始 indexer key 的特征维度，不是 head 数，也不包含 gate 部分。 - # ring 的 channel 轴拼接 key 和同维度的 gate，共 2 * head_dim 个元素。 - head_dim: int - # 原始 key/gate ring 的存储精度，不代表池化后 packed indexer key 的量化编码类型。 - dtype: torch.dtype - # 当前模型 TP 组内的 rank。K-pool tail 在组内完整复制，传输时仅 rank 0 写出。 - tp_rank: int = 0 - # 当前模型的 TP 组大小，用于约束 rank 范围；不按此数切分 tail 的 head_dim。 - # 外部传输布局只保留一份 pool_size 长的逻辑窗口，所有 rank 可恢复到自己的 ring。 - tp_world_size: int = 1 - - def __post_init__(self) -> None: - """构造配置时检查维度、层号和 TP 范围。""" - assert min(self.pool_size, self.head_dim) > 0, "invalid K-pool dimensions" - assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids), ( - "tail layers must be nonempty and unique" - ) - assert 0 <= self.tp_rank < self.tp_world_size, "invalid TP rank" - - -class KPoolAttentionTailState(AttentionTailState): - """保存 K-pool 原始 key/gate ring，完整保留 pool_size + mtp_step 个槽位。""" - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - def __init__( - self, - config: KPoolStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", - ) -> None: - """为各层、各请求创建原始 key/gate ring，包含 MTP 所需额外历史。""" - assert request_capacity > 0 and mtp_step >= 0, "invalid request capacity or MTP step" - self.config = config - self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 - self.mtp_step = mtp_step # ring 额外保留的 MTP 历史槽位数。 - self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 - self._logical_to_storage_layer = { - logical_layer_id: storage_layer_index - for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) - } - self.tail_buffer = torch.zeros( - ( - len(config.logical_layer_ids), - request_capacity, - config.pool_size + mtp_step, - 2 * config.head_dim, - ), - dtype=config.dtype, - device=device, - ) - - def get_state_type(self) -> str: - """返回 K-pool 压缩比例及存储精度。""" - return f"kpool_c{self.config.pool_size}_{dtype_name(dtype=self.config.dtype)}" - - # ======================================================================== - # 二、Attention tail：计算视图与无快照重置 - # ======================================================================== - - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" - return logical_layer_id in self._logical_to_storage_layer - - def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: - """返回该层所有请求的原生 ring 视图，末轴包含 key 和 gate。""" - return self.tail_buffer[self._logical_to_storage_layer[logical_layer_id]] - - def reset_attention_tail_state( - self, - req_indices: list[int], - ) -> None: - """清零指定请求的完整 ring；调用方保证用于空前缀或对齐池化边界，不用于 MTP 回滚。""" - assert isinstance(req_indices, list), "request indices must be a list" - for req_idx in req_indices: - self.tail_buffer[:, req_idx].fill_(0) - # ======================================================================== - # 三、PD 分离传输 - # ======================================================================== - - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """描述最近 pool_size 个逻辑位置的 key/gate，由 TP rank 0 写出；不包含 MTP 额外槽位。""" - config = self.config - window_size = config.pool_size - channels = 2 * config.head_dim - tail_region = TransferRegion( - field="tail", - byte_start=0, - byte_end=len(config.logical_layer_ids) * window_size * channels * config.dtype.itemsize, - dtype=config.dtype, - axes=("layer", "window", "channel"), - shape=(len(config.logical_layer_ids), window_size, channels), - axis_mappings=( - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=window_size, - local_circular=True, - ), - AxisMapping( - axis="channel", - local_start=0, - external_start=0, - length=channels, - ), - ), - write_enabled=config.tp_rank == 0, - ) - return (tail_region,) - - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """将 prefill 末尾的逻辑窗口写入 PD 页；state_token_len 为已执行的前缀长度。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - transfer_kpool_pages( - tail=self.tail_buffer[:, req_idx], - page_buffers=pd_page_buffers, - page_index=pd_page_index, - state_token_len=state_token_len, - layout=layout, - write=True, - ) - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """按前缀长度将逻辑窗口放回本地 ring，保留窗口外槽位，允许两端 MTP 步数不同。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - transfer_kpool_pages( - tail=self.tail_buffer[:, req_idx], - page_buffers=pd_page_buffers, - page_index=pd_page_index, - state_token_len=state_token_len, - layout=layout, - write=False, - ) diff --git a/lightllm/common/attention_state/att_tail/kpool_transfer.py b/lightllm/common/attention_state/att_tail/kpool_transfer.py deleted file mode 100644 index 235246f0d0..0000000000 --- a/lightllm/common/attention_state/att_tail/kpool_transfer.py +++ /dev/null @@ -1,186 +0,0 @@ -"""K-pool 的 PD 逻辑窗口传输：一次 kernel 完成环形映射及窗口数据搬运。""" - -import torch -import triton -import triton.language as tl - -from ..types import AxisMapping, StateTransferLayout -from ..utils import view_max_copy_dtype - - -# ======================================================================== -# 一、K-pool 传输：固定布局 + 本次窗口位置，直接操作原始存储 -# ======================================================================== - - -def transfer_kpool_pages( - tail: torch.Tensor, - page_buffers: torch.Tensor, - page_index: int, - state_token_len: int, - layout: StateTransferLayout, - write: bool, - axis_mapping_overrides: tuple[AxisMapping, ...] = (), - grid_num: int = 24, -) -> None: - """在当前 CUDA stream 交换一个请求的 tail 和选定 PD 页，不创建中间 Tensor。 - - 本地 tail 为 [layer, ring_slot, channel]，ring 容量为 pool_size + 本地 mtp_step； - 外部页保存最近 pool_size 个逻辑位置，按时间顺序排列。P/D 的 ring 容量可以不同。 - 窗口起点由 state_token_len 推导；axis_mapping_overrides 可覆盖本次 window 映射， - 其余轴沿用缓存布局。K-pool 只支持 window 环形寻址，channel 保持完整、同序对应。 - 两个方向都只遍历逻辑窗口，短前缀对应的窗口内位置补零； - 本地窗口之外的 MTP 槽位保持原值，不在传输时清理。 - 调用方负责传输完成依赖及存储生命周期；返回不表示数据已经就绪。 - """ - # 阶段一：解析选定页的 tail 区域，统一为 [layer, window, channel] 视图。 - # K-pool 布局由 State 创建；此处直接使用约定的单个 tail region。 - region = layout.regions[0] - if write and not region.write_enabled: - return - - page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) - field_bytes = page_bytes[region.byte_start:region.byte_end] - external_tail = field_bytes.view(region.dtype).view(region.shape) - axis_order = tuple(region.axes.index(axis) for axis in ("layer", "window", "channel")) - external_tail = external_tail.permute(axis_order) - local_ring_size = tail.shape[1] - external_ring_size = external_tail.shape[1] - channel_count = tail.shape[2] - - # 阶段二：确定本次窗口。显式覆盖的起点直接使用；否则从 prefill 长度推导。 - # 只计算整数位置，不修改 layout，也不创建新的 AxisMapping 对象。 - layout_mappings = {item.axis: item for item in region.axis_mappings} - mapping_overrides = {item.axis: item for item in axis_mapping_overrides} - window_mapping = mapping_overrides.get("window", layout_mappings["window"]) - channel_mapping = mapping_overrides.get("channel", layout_mappings["channel"]) - window_size = window_mapping.length - local_window_start = window_mapping.local_start - external_window_start = window_mapping.external_start - - if window_mapping.local_circular: - if "window" not in mapping_overrides: - # 外部 window[i] 对应绝对位置 state_token_len - window_size + i。 - local_window_start += state_token_len - window_size - local_window_start %= local_ring_size - if window_mapping.external_circular: - external_window_start %= external_ring_size - invalid_prefix = max(window_size - state_token_len, 0) - - # 定制 kernel 完整处理 key/gate；窗口不能重复覆盖或越过线性存储边界。 - assert ( - channel_mapping.local_start == channel_mapping.external_start == 0 - and channel_mapping.length == channel_count - and not channel_mapping.local_circular and not channel_mapping.external_circular - ), "K-pool transfers require the complete channel axis" - assert ( - 0 < window_size <= min(local_ring_size, external_ring_size) - and (window_mapping.local_circular or 0 <= local_window_start <= local_ring_size - window_size) - and (window_mapping.external_circular or 0 <= external_window_start <= external_ring_size - window_size) - ), "window exceeds storage capacity" - - # 阶段三：只扩大连续 channel 的搬运单位，layer/window 维度和环形语义保持不变。 - # 例如 BF16 的 [layers, ring_size, 256] 按 INT64 搬运时变为 [layers, ring_size, 64]。 - local_tail, external_tail = view_max_copy_dtype(copies=[(tail, external_tail)])[0] - - # 阶段四：正反方向都只处理逻辑窗口，窗口外的 MTP 槽位不参与读写。 - # grid_num 是发射上限，大窗口由 kernel 的跳步循环处理。 - element_count = local_tail.shape[0] * window_size * local_tail.shape[2] - block_elements = 1024 - program_count = min(grid_num, triton.cdiv(element_count, block_elements)) - _transfer_kpool_pages[(program_count,)]( - LOCAL=local_tail, - EXTERNAL=external_tail, - local_window_start=local_window_start, - external_window_start=external_window_start, - invalid_prefix=invalid_prefix, - LOCAL_STRIDES=local_tail.stride(), - EXTERNAL_STRIDES=external_tail.stride(), - LAYERS=local_tail.shape[0], - LOCAL_RING_SIZE=local_ring_size, - EXTERNAL_RING_SIZE=external_ring_size, - WINDOW_SIZE=window_size, - CHANNELS=local_tail.shape[2], - ELEMENT_BYTES=local_tail.element_size(), - LOCAL_CIRCULAR=window_mapping.local_circular, - EXTERNAL_CIRCULAR=window_mapping.external_circular, - WRITE=write, - BLOCK_ELEMENTS=block_elements, - num_warps=4, - ) - - -# ======================================================================== -# 二、定制 kernel：计算窗口位置 → 读取有效历史 → 写出数据或零 -# ======================================================================== - - -@triton.jit(do_not_specialize=["local_window_start", "external_window_start", "invalid_prefix"]) -def _transfer_kpool_pages( - LOCAL, # 当前请求的本地 ring：[layer, ring_slot, channel]。 - EXTERNAL, # 已选定页中的 tail 字段：[layer, window, channel]。 - local_window_start, # 本次逻辑窗口在本地 ring 的起点，保持为运行时参数。 - external_window_start, # 本次窗口在外部字段中的起点，通常为 0。 - invalid_prefix, # 短前缀中不存在的窗口前部位置数量。 - LOCAL_STRIDES: tl.constexpr, - EXTERNAL_STRIDES: tl.constexpr, - LAYERS: tl.constexpr, - LOCAL_RING_SIZE: tl.constexpr, - EXTERNAL_RING_SIZE: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - CHANNELS: tl.constexpr, - ELEMENT_BYTES: tl.constexpr, - LOCAL_CIRCULAR: tl.constexpr, - EXTERNAL_CIRCULAR: tl.constexpr, - WRITE: tl.constexpr, - BLOCK_ELEMENTS: tl.constexpr, -): - if ELEMENT_BYTES == 8: - copy_dtype = tl.uint64 - elif ELEMENT_BYTES == 4: - copy_dtype = tl.uint32 - elif ELEMENT_BYTES == 2: - copy_dtype = tl.uint16 - else: - copy_dtype = tl.uint8 - LOCAL = LOCAL.to(tl.pointer_type(copy_dtype)) - EXTERNAL = EXTERNAL.to(tl.pointer_type(copy_dtype)) - local_window_start = tl.cast(local_window_start, tl.int64) - external_window_start = tl.cast(external_window_start, tl.int64) - element_count = LAYERS * WINDOW_SIZE * CHANNELS - program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS - grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS - lanes = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) - - for block_start in tl.range(program_start, element_count, grid_stride, num_stages=1): - # 两个方向共用逻辑窗口坐标，只在 load/store 时交换数据流向。 - element = block_start + lanes - valid = element < element_count - channel = element % CHANNELS - window = element // CHANNELS % WINDOW_SIZE - layer = element // (CHANNELS * WINDOW_SIZE) - local_slot = local_window_start + window - if LOCAL_CIRCULAR: - local_slot %= LOCAL_RING_SIZE - external_slot = external_window_start + window - if EXTERNAL_CIRCULAR: - external_slot %= EXTERNAL_RING_SIZE - has_history = (window >= invalid_prefix) & (window < WINDOW_SIZE) - - # 所有 stride 在参与地址乘加前转为 int64，支持请求切片产生的大 layer stride。 - local_offset = ( - layer * tl.cast(LOCAL_STRIDES[0], tl.int64) - + local_slot * tl.cast(LOCAL_STRIDES[1], tl.int64) - + channel * tl.cast(LOCAL_STRIDES[2], tl.int64) - ) - external_offset = ( - layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) - + external_slot * tl.cast(EXTERNAL_STRIDES[1], tl.int64) - + channel * tl.cast(EXTERNAL_STRIDES[2], tl.int64) - ) - if WRITE: - value = tl.load(LOCAL + local_offset, mask=valid & has_history, other=0) - tl.store(EXTERNAL + external_offset, value, mask=valid) - else: - value = tl.load(EXTERNAL + external_offset, mask=valid & has_history, other=0) - tl.store(LOCAL + local_offset, value, mask=valid) diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 689b61b60b..9796d49c0a 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -30,7 +30,7 @@ class BoundedAttentionState(ABC): def get_state_type(self) -> str: """返回算法与存储方案标记，用于生成资源成员名。 - 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 @@ -187,8 +187,14 @@ def copy_runtime_state_from_dp( # ======================================================================== @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" + """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 管理器分配存储后统一清理，子类覆盖时须保留缓存和副本语义。 + """ ... @abstractmethod diff --git a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py index 00a7862f1b..4400aa1f5b 100644 --- a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py +++ b/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py @@ -281,6 +281,8 @@ def copy_runtime_state_from_dp( # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回完整逻辑窗口，不包含本地 MTP 候选区。""" return self._get_transfer_regions() diff --git a/lightllm/common/attention_state/bounded_att/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn_state.py index c7d8d75649..c850067c87 100644 --- a/lightllm/common/attention_state/bounded_att/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn_state.py @@ -317,6 +317,8 @@ def copy_runtime_state_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" return self._get_transfer_regions() diff --git a/lightllm/common/attention_state/bounded_att/kda_state.py b/lightllm/common/attention_state/bounded_att/kda_state.py index 312f0dabee..b34ffdd3f4 100644 --- a/lightllm/common/attention_state/bounded_att/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda_state.py @@ -317,6 +317,8 @@ def copy_runtime_state_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" return self._get_transfer_regions() diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor_state.py index 4c3c7f983d..91b540760e 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor_state.py @@ -10,6 +10,8 @@ import torch +from lightllm.utils.envs_utils import get_env_start_args + from ..utils import deepcopy_result from ..types import PageKind, StateTransferLayout, TensorField, TransferRegion from .base_state import BoundedAttentionState @@ -21,21 +23,41 @@ @dataclass(frozen=True) class WindowCompressorStateConfig: - # 本实例服务的逻辑层；KV 与 indexer 的压缩器分别创建实例。 + """按固定窗口和产出步长保存候选向量及 gate，供模型的压缩或池化算子使用。 + + 这里只描述运行态的存储和历史保留方式，具体聚合、归一化及量化由模型算子完成。 + window_size / compress_stride 分别表示窗口覆盖的原始 token 数和相邻产出的间隔： + - GLM K-pool：W4/S4，每四个 key/gate 形成一组，无跨组重叠历史。 + - V4 C4：W8/S4，每四步产出一次，下一组仍需保留四个 token 的历史。 + - V4 C128：W128/S128，每 128 步产出一次，无跨组重叠历史。 + MTP 的额外候选槽位由 State.mtp_step 扩展，不计入窗口或产出步长。 + """ + + # 本实例服务的逻辑层，顺序对应运行态 Tensor 的 layer 轴。 + # attention KV 候选与索引 key 候选分别创建实例，可使用不同窗口、维度和精度。 logical_layer_ids: tuple[int, ...] - # 每次压缩聚合的原始 token 窗口，例如 C4 为 W8，C128 为 W128。 + # 单次聚合覆盖的原始 token 窗口，不包含 MTP 额外槽位。 + # 例如 K-pool 为 4，V4 C4 为 8；压缩比例同为 4 不代表窗口长度相同。 window_size: int - # 每执行多少个原始 token 产出一个条目，例如 C4 为 S4，C128 为 S128。 + # 相邻产出之间经过的原始 token 数；K-pool / V4 C4 为 4，V4 C128 为 128。 + # 等于 window_size 时无重叠，小于时保留 window_size - compress_stride 个历史位置。 + # 大小页快照在该步长的边界保存；PD 可在任意前缀终点传输完整逻辑窗口。 compress_stride: int - # 每个 token 保存的候选向量宽度，gate 另占同样宽度。 - # W8/S4 可为前后半窗口分别保存候选，所以当前 V4 使用 2 * 输出向量维度。 + # 每个原始 token 保存的候选向量宽度；gate 具有相同形状和存储精度。 + # K-pool 中等于 indexer key 的 head_dim；V4 W8/S4 为前后半窗口分别保存候选， + # 使用 2 * 输出向量维度。因此它不一定等于最终压缩条目的维度，也不包含 gate 的宽度。 candidate_dim: int - # 产出的用途用于区分规范 State 成员名；不改变窗口保存和恢复算法。 + # 候选数据的用途，仅用于区分规范 State 成员名，不选择聚合算法或改变存储和恢复行为。 + # kv 表示 attention KV 候选；indexer 表示索引 key 候选及其 gate，不代表整个 indexer。 target: Literal["kv", "indexer"] = "kv" - # 候选和 gate 的精度，当前压缩算子使用 FP32。 + # 候选与 gate 使用相同精度，沿最后一轴拼接为 candidate|gate，共用一个 Tensor。 + # 此精度描述压缩器输入历史，不是产出 KV / indexer key 的量化编码； + # GLM K-pool 随模型精度配置（通常 BF16），V4 的 KV/indexer 压缩器当前使用 FP32。 + # 当前接口不支持候选与 gate 使用不同 dtype。 dtype: torch.dtype = torch.float32 - # 运行态在 TP 内完整复制，只有 rank 0 向共享外部页写出。 + # 所属模型 TP 组内的 rank；运行态完整复制，仅 rank 0 向共享外部页写出。 tp_rank: int = 0 + # 保存相同运行态副本的 TP rank 数，不用于切分 candidate_dim 或窗口长度。 tp_world_size: int = 1 @property @@ -200,7 +222,17 @@ def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: to @deepcopy_result @lru_cache(maxsize=None) def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回快照布局；缓存边界无需保留历史的窗口返回空区域。""" + """校验缓存页对齐产出步长；无重叠历史的窗口不占 CPU 页，也不创建大小页池。""" + args = get_env_start_args() + # 大页是小页的整数倍；小页对齐产出步长后，大小页终点均不会落在未完成的组内。 + assert args.linear_att_hash_page_size % self.config.compress_stride == 0, ( + "state cache pages must align to the compressor stride" + ) + # CPU 服务只查询 meta 布局，也要在忽略无重叠窗口之前拒绝不安全的恢复边界。 + if args.cpu_cache_token_page_size is not None: + assert args.cpu_cache_token_page_size % self.config.compress_stride == 0, ( + "CPU cache pages must align to the compressor stride" + ) return self._get_transfer_regions(runtime=False) def write_state_pages_to_cpu( @@ -291,6 +323,8 @@ def copy_runtime_state_from_dp( # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回完整逻辑窗口，不包含本地 MTP 候选区。""" return self._get_transfer_regions(runtime=True) diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py b/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py index 43fd11b0bf..49e88896b0 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py @@ -2,9 +2,9 @@ import torch import triton +import triton.language as tl from ..utils import view_max_copy_dtype -from ..att_tail.kpool_transfer import _transfer_kpool_pages def copy_compressor_window( @@ -22,12 +22,11 @@ def copy_compressor_window( local_view, external_view = view_max_copy_dtype(copies=[(local, external)])[0] window = external.shape[1] elements = local_view.shape[0] * window * local_view.shape[2] - _transfer_kpool_pages[(min(24, triton.cdiv(elements, 1024)),)]( + _copy_compressor_window[(min(24, triton.cdiv(elements, 1024)),)]( LOCAL=local_view, EXTERNAL=external_view, - local_window_start=(state_token_len - window) % local.shape[1], + window_start=state_token_len - window, external_window_start=0, - invalid_prefix=max(window - state_token_len, 0), LOCAL_STRIDES=local_view.stride(), EXTERNAL_STRIDES=external_view.stride(), LAYERS=local_view.shape[0], @@ -42,3 +41,76 @@ def copy_compressor_window( BLOCK_ELEMENTS=1024, num_warps=4, ) + + +@triton.jit(do_not_specialize=["window_start", "external_window_start"]) +def _copy_compressor_window( + LOCAL, # 当前请求的本地 ring：[layer, ring_slot, channel]。 + EXTERNAL, # 已选定页中的候选/gate 拼接字段：[layer, window, channel]。 + window_start, # 未取模的绝对窗口起点；短前缀时可为负数，保持为运行时参数。 + external_window_start, # 本次窗口在外部字段中的起点，通常为 0。 + LOCAL_STRIDES: tl.constexpr, + EXTERNAL_STRIDES: tl.constexpr, + LAYERS: tl.constexpr, + LOCAL_RING_SIZE: tl.constexpr, + EXTERNAL_RING_SIZE: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + CHANNELS: tl.constexpr, + ELEMENT_BYTES: tl.constexpr, + LOCAL_CIRCULAR: tl.constexpr, + EXTERNAL_CIRCULAR: tl.constexpr, + WRITE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, +): + if ELEMENT_BYTES == 8: + copy_dtype = tl.uint64 + elif ELEMENT_BYTES == 4: + copy_dtype = tl.uint32 + elif ELEMENT_BYTES == 2: + copy_dtype = tl.uint16 + else: + copy_dtype = tl.uint8 + LOCAL = LOCAL.to(tl.pointer_type(copy_dtype)) + EXTERNAL = EXTERNAL.to(tl.pointer_type(copy_dtype)) + window_start = tl.cast(window_start, tl.int64) + external_window_start = tl.cast(external_window_start, tl.int64) + element_count = LAYERS * WINDOW_SIZE * CHANNELS + program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS + lanes = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) + + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=1): + # 两个方向共用逻辑窗口坐标，只在 load/store 时交换数据流向。 + element = block_start + lanes + valid = element < element_count + channel = element % CHANNELS + window_index = element // CHANNELS % WINDOW_SIZE + layer = element // (CHANNELS * WINDOW_SIZE) + # 先用绝对位置判断历史是否存在，再映射到物理 ring，避免取模丢失负数信息。 + position = window_start + window_index + has_history = position >= 0 + local_slot = position + if LOCAL_CIRCULAR: + # 窗口不超过 ring 容量；先加容量，使短前缀读回补零也落在非负槽位。 + local_slot = (position + LOCAL_RING_SIZE) % LOCAL_RING_SIZE + external_slot = external_window_start + window_index + if EXTERNAL_CIRCULAR: + external_slot %= EXTERNAL_RING_SIZE + + # 所有 stride 在参与地址乘加前转为 int64，支持请求切片产生的大 layer stride。 + local_offset = ( + layer * tl.cast(LOCAL_STRIDES[0], tl.int64) + + local_slot * tl.cast(LOCAL_STRIDES[1], tl.int64) + + channel * tl.cast(LOCAL_STRIDES[2], tl.int64) + ) + external_offset = ( + layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) + + external_slot * tl.cast(EXTERNAL_STRIDES[1], tl.int64) + + channel * tl.cast(EXTERNAL_STRIDES[2], tl.int64) + ) + if WRITE: + value = tl.load(LOCAL + local_offset, mask=valid & has_history, other=0) + tl.store(EXTERNAL + external_offset, value, mask=valid) + else: + value = tl.load(EXTERNAL + external_offset, mask=valid & has_history, other=0) + tl.store(LOCAL + local_offset, value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index 1a44a487ef..e591927ac9 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -31,7 +31,7 @@ class FullAttentionState(ABC): def get_state_type(self) -> str: """返回算法与存储方案标记，用于生成资源成员名。 - 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att/att_tail 类别前缀， + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 @@ -158,8 +158,15 @@ def copy_tokens_from_dp( # ======================================================================== @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。""" + """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 PD 页规格保持不变；管理器分配存储后统一清除缓存。 + 子类覆盖时须保留缓存和副本语义。 + """ ... @abstractmethod diff --git a/lightllm/common/attention_state/full_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att_state.py index 74722526b4..f1203b41c3 100644 --- a/lightllm/common/attention_state/full_att/compressed_att_state.py +++ b/lightllm/common/attention_state/full_att/compressed_att_state.py @@ -322,6 +322,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py index 69e6c4c289..71a71f2433 100644 --- a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py @@ -225,6 +225,8 @@ def copy_tokens_from_dp( # 五、PD 分离：传输固定的完整 token 编码 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """返回固定 FP8/E8M0 格式的 PD 页描述，由管理器安排联合页偏移。""" return self._get_transfer_regions(page_tokens=get_env_start_args().pd_kv_page_size) diff --git a/lightllm/common/attention_state/full_att/gqa_state.py b/lightllm/common/attention_state/full_att/gqa_state.py index ffcbed33db..5dea8622d2 100644 --- a/lightllm/common/attention_state/full_att/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa_state.py @@ -266,6 +266,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/full_att/group_quantized_state.py b/lightllm/common/attention_state/full_att/group_quantized_state.py index 9029ca1a0c..5d8e6532e3 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/group_quantized_state.py @@ -319,6 +319,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla_state.py index 71b34a05a6..da82e0cddd 100644 --- a/lightllm/common/attention_state/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla_state.py @@ -225,6 +225,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla_state.py index 52edb0e1dd..a596e5a97c 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla_state.py @@ -308,6 +308,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/full_att/static_fp8_state.py b/lightllm/common/attention_state/full_att/static_fp8_state.py index b874cc421d..b7df6f2cfe 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/static_fp8_state.py @@ -343,6 +343,8 @@ def copy_tokens_from_dp( # 五、PD 分离传输 # ======================================================================== + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" page_tokens = get_env_start_args().pd_kv_page_size diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index b895383da9..7f425aae07 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -21,7 +21,6 @@ from .allocators import StatePageAllocator, FullAttKVAllocator from .full_att import FullAttentionState from .bounded_att import BoundedAttentionState -from .att_tail import AttentionTailState from .utils import MODEL_ROLES if TYPE_CHECKING: @@ -35,7 +34,7 @@ class AttentionStateManager: """管理模型组合的注意力资源，提供缓存复用及 CPU/DP/PD 传输所需的聚合操作。 - 通过聚合查询操作各类 State；模型通过三类创建接口决定资源组合，main/draft 共享本管理器。 + 通过聚合查询操作各类 State；模型通过两类创建接口决定资源组合，main/draft 共享本管理器。 未支持的传输路径必须明确报错。 """ @@ -55,10 +54,9 @@ def __init__(self, request_capacity: int): self.bounded_att_big_page_allocator = None self.bounded_att_small_page_allocator = None self._request_capacity = request_capacity - self.cpu_tail_restore = False def register_state( - self, model_role: str, state: FullAttentionState | BoundedAttentionState | AttentionTailState + self, model_role: str, state: FullAttentionState | BoundedAttentionState ) -> None: """按模型角色和具体方案注册 State，保留原对象及其存储。 @@ -84,8 +82,8 @@ def get_state( model_role: str, state_type_filter: str, logical_layer_id: int, - state_class: type[FullAttentionState | BoundedAttentionState | AttentionTailState] | None = None, - ) -> FullAttentionState | BoundedAttentionState | AttentionTailState: + state_class: type[FullAttentionState | BoundedAttentionState] | None = None, + ) -> FullAttentionState | BoundedAttentionState: """按模型角色、算法/方案字符串和逻辑层选择唯一 State，缓存查询结果。 model_role 为 main 或 draft0..draft35，按完整角色前缀匹配，避免 draft1 匹配 draft10。 @@ -100,7 +98,7 @@ def get_state( state for name, state in self.__dict__.items() if name.startswith(model_role + "_") - and isinstance(state, (FullAttentionState, BoundedAttentionState, AttentionTailState)) + and isinstance(state, (FullAttentionState, BoundedAttentionState)) and state_type_filter in state.get_state_type() and (state_class is None or isinstance(state, state_class)) and state.contains_logical_layer(logical_layer_id) @@ -233,7 +231,7 @@ def allocate_storage( ) -> None: """在 State 对象创建完成后分配 full-attention 存储、共享 allocator 和 bounded 大小页。 - 请求容量对应的 bounded/indexer 运行态已由各模型创建，这里仅补齐缓存存储。 + 请求容量对应的 bounded 运行态（包括压缩窗口）已由各模型创建，这里仅补齐缓存存储。 """ assert self._state_registration_finished, "State registration must finish before storage initialization" assert self.full_att_allocator is None, "resources already initialized" @@ -256,6 +254,10 @@ def allocate_storage( StatePageAllocator(small_page_capacity) if small_page_capacity else None ) for state in bounded_states: + # 无快照区域的 State 在合法缓存边界无需历史，不为它分配大小页 Tensor。 + # 共享槽位 allocator 仍供其他 State 和缓存流程使用，不随单个 State 重复创建。 + if not state.get_cpu_transfer_regions(): + continue state.initialize_page_pools( big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=torch.device(device).type == "cuda", @@ -263,10 +265,10 @@ def allocate_storage( self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 - self.cpu_tail_restore = bool(self.get_attention_tail_states()) - # 存储分配可能改变 State 的物理页规格，清除准备阶段的布局及其派生 CPU/PD 页大小缓存。 + # 存储分配可能改变物理页规格，清除 State 的 CPU/PD 区域和管理器的派生布局缓存。 for state in (*full_states, *bounded_states): state.get_cpu_transfer_regions.cache_clear() + state.get_pd_transfer_regions.cache_clear() self.get_cpu_page_layouts.cache_clear() self.get_cpu_page_bytes.cache_clear() self.get_pd_token_page_layouts.cache_clear() @@ -323,10 +325,6 @@ def get_bounded_attention_states(self, role: str | None = None) -> tuple[Bounded """返回全部或指定角色的 bounded-attention State。""" return self._get_states(BoundedAttentionState, role) - def get_attention_tail_states(self, role: str | None = None) -> tuple[AttentionTailState, ...]: - """返回全部或指定角色的 attention-tail State。""" - return self._get_states(AttentionTailState, role) - def get_full_att_token_bytes(self) -> int: """累计当前 TP rank 全部 full-attention State 的单 token 存储字节数。""" total = 0 @@ -382,7 +380,7 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV、tail 和状态页。""" + """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" for state in self.get_bounded_attention_states(): state.reset_bounded_att_state(req_indices) @@ -447,23 +445,7 @@ def commit_bounded_att_state( state.commit_bounded_att_state(req_indices, accepted_steps) # ======================================================================== - # 四、Attention tail：空前缀初始化与无快照恢复 - # ======================================================================== - - # 没有 tail 快照时使用；长度 0 用于空前缀初始化，没有 tail 的模型可为空操作。 - def reset_attention_tail_state( - self, - req_indices: list[int], - ) -> None: - """依次清零指定请求的 tail，恢复边界由模型配置和缓存调用流程保证。 - - 保留已恢复的 KV 和 bounded 运行态，不提供跨 State 的原子性保证。 - """ - for state in self.get_attention_tail_states(): - state.reset_attention_tail_state(req_indices) - - # ======================================================================== - # 五、CPU Cache 数据交换 + # 四、CPU Cache 数据交换 # ======================================================================== @lru_cache(maxsize=None) @@ -478,7 +460,7 @@ def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: assert states, "a CPU page requires participating storage" # 按 full-attention、bounded 顺序组织全部角色；各类 State 内部按资源名排序。 - # Indexer tail 沿用前缀边界重置，不加入当前 CPU 联合页。 + # 无重叠窗口在缓存组边界无需快照，由 State 返回空 CPU 区域。 layouts, offset = {}, 0 for state in states: regions = tuple( @@ -585,12 +567,11 @@ def read_req_from_cpu( big_page_slot_ids 只包含完整 CPU 页对应的普通槽位；末尾碎页使用内部预留的 LOAD 大页。 存在 bounded 时须提供 state_token_len，即恢复后的绝对前缀长度，不是本批 token 数。 - 存在 bounded 或 tail 时必须提供 req_idx；调用方保证最终前缀位于 tail 可重置边界。 + 存在 bounded 时必须提供 req_idx；最终前缀须满足各 State 的快照或重置边界。 """ full_states = self.get_full_attention_states() bounded_states = self.get_bounded_attention_states() - tail_states = self.get_attention_tail_states() - if bounded_states or tail_states: + if bounded_states: assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" if bounded_states: assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" @@ -637,12 +618,8 @@ def read_req_from_cpu( page_slot_ids=[load_big_page_slot_ids[-1]], ) - # 阶段四：清零目标请求的 tail；调用方保证前缀终点满足重置边界。 - for state in tail_states: - state.reset_attention_tail_state([req_idx]) - # ======================================================================== - # 六、DP 迁移 + # 五、DP 迁移 # ======================================================================== def copy_from_dp( @@ -659,7 +636,7 @@ def copy_from_dp( token 索引由调用方准备在源、目标各自的 CUDA 设备上，目标索引为 int64。 不需要复制 token 或大页时传空 Tensor 或空列表；请求运行态仍按请求索引复制。 - 调用方保证源 MTP 游标为 0、目标前缀位于 tail 可重置边界，以及源页就绪和引用有效。 + 调用方保证源 MTP 游标为 0、目标前缀满足各 State 的缓存边界，以及源页就绪和引用有效。 """ assert type(remote_att_manager) is type(self), "DP replicas must use the same model resource manager" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" @@ -678,17 +655,16 @@ def copy_from_dp( dst_page_ids=dst_big_page_ids, remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), ) - # 运行态不由大页复制隐式恢复；按指定请求复制基准位置，随后重置目标 tail。 + # 运行态由具体 State 复制或按合法缓存边界重置，不由大页复制隐式恢复。 for state in self.get_bounded_attention_states(): state.copy_runtime_state_from_dp( src_req_indices=[src_req_idx], dst_req_indices=[dst_req_idx], remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), ) - self.reset_attention_tail_state([dst_req_idx]) # ======================================================================== - # 七、PD 分离传输 + # 六、PD 分离传输 # ======================================================================== @lru_cache(maxsize=None) @@ -724,7 +700,7 @@ def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: 每个 State 保存一份运行态，tokens_per_page 为 None；无 State 时返回空字典。 """ layouts, offset = {}, 0 - for state in (*self.get_bounded_attention_states(), *self.get_attention_tail_states()): + for state in self.get_bounded_attention_states(): regions = tuple( replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) for region in state.get_pd_transfer_regions() @@ -740,7 +716,7 @@ def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: offset = state_end return layouts - # PD 分离：kind="kv" 传输 token 数据，kind="att_state" 传输 bounded 运行态及 tail。 + # PD 分离：kind="kv" 传输 token 数据，kind="att_state" 传输 bounded 运行态。 # pd_page_buffers 第一维为页，pd_page_index 选择物理页；网络收发与同步由传输流程组织。 def write_to_pd_page( self, @@ -766,9 +742,8 @@ def write_to_pd_page( ) case "att_state": - # 前缀长度由 PD 任务提供，bounded/tail 据此定位各自的运行态窗口。 + # 前缀长度由 PD 任务提供，各 State 据此定位自己的运行态窗口。 assert req_idx is not None, "state PD pages require a request index" - tail_states = self.get_attention_tail_states() assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" # 阶段一：写出 bounded 的 prefill 基准运行态。 @@ -781,16 +756,6 @@ def write_to_pd_page( layout=layouts[state.state_name], ) - # 阶段二：按前缀长度写出 tail 逻辑窗口，适配两端各自的 MTP ring。 - for state in tail_states: - state.write_to_pd_page( - req_idx=req_idx, - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - state_token_len=state_token_len, - layout=layouts[state.state_name], - ) - case _: raise AssertionError(f"unknown PD page kind {kind}") @@ -818,9 +783,8 @@ def read_from_pd_page( ) case "att_state": - # 前缀长度由 PD 任务提供，bounded/tail 据此定位各自的运行态窗口。 + # 前缀长度由 PD 任务提供，各 State 据此定位自己的运行态窗口。 assert req_idx is not None, "state PD pages require a request index" - tail_states = self.get_attention_tail_states() assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" # 阶段一：恢复 bounded 的 prefill 基准运行态。 @@ -833,28 +797,18 @@ def read_from_pd_page( layout=layouts[state.state_name], ) - # 阶段二：按前缀长度恢复 tail 逻辑窗口，适配两端各自的 MTP ring。 - for state in tail_states: - state.read_from_pd_page( - pd_page_buffers=pd_page_buffers, - req_idx=req_idx, - pd_page_index=pd_page_index, - state_token_len=state_token_len, - layout=layouts[state.state_name], - ) - case _: raise AssertionError(f"unknown PD page kind {kind}") # ======================================================================== - # 八、私有辅助实现 + # 七、私有辅助实现 # ======================================================================== # 私有实现：缓存创建完成后的成员集合，模型逐层计算直接访问具体 State 成员。 @lru_cache(maxsize=None) def _get_states( self, - state_class: type[FullAttentionState] | type[BoundedAttentionState] | type[AttentionTailState], + state_class: type[FullAttentionState] | type[BoundedAttentionState], role: str | None, ): """按 State 类型和模型角色遍历规范成员，只缓存对象引用。""" diff --git a/lightllm/common/attention_state/types.py b/lightllm/common/attention_state/types.py index e2fe4c8bd6..9b578b2658 100644 --- a/lightllm/common/attention_state/types.py +++ b/lightllm/common/attention_state/types.py @@ -170,12 +170,12 @@ class TransferRegion: # 描述本地字段到外部区域的轴范围，每个非 layer/token 轴必须恰好出现一次。 # 无分片的轴也显式填写完整范围，例如 dim: 0 -> 0、length=head_dim。 # 两个方向共用同一组本地/外部起点，仅反转数据流向；映射应用于全部层和本次有效 token。 - # 环形寻址由支持它的具体 State 搬运实现处理；请求窗口起点通过本次调用的 - # axis_mapping_overrides 覆盖，不回写缓存布局，起点也不作为 kernel 的编译期常量。 + # 环形寻址由支持它的具体 State 搬运实现处理；请求窗口起点根据本次绝对前缀长度计算， + # 不回写缓存布局，也不作为 kernel 的编译期常量。 axis_mappings: tuple[AxisMapping, ...] # 当前 TP rank 是否负责向外部区域写入；False 不影响从该区域读回。 - # KV head 有 TP 副本时，每组只选一个 rank 写入；复制的 indexer tail 只由 rank 0 写。 + # KV head 有 TP 副本时，每组只选一个 rank 写入；复制的压缩窗口只由 rank 0 写。 # 标记属于 region，因为同一 State 内各字段的 TP 分片/复制规则可能不同。 write_enabled: bool = True diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index b7f1013a2f..e0c439cd8a 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -51,7 +51,6 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.attention_state.bounded_att import BoundedAttentionState - from lightllm.common.attention_state.att_tail import AttentionTailState logger = init_logger(__name__) @@ -167,12 +166,7 @@ def prepare_attention_states(self) -> None: request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, device="cuda", ) - tail_states = self.create_att_tail_states( - model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, - request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, - device="cuda", - ) - for state in (*full_states, *bounded_states, *tail_states): + for state in (*full_states, *bounded_states): self.att_state_manager.register_state(self.att_state_role, state) def init_runtime(self) -> None: @@ -282,18 +276,6 @@ def create_bounded_att_states( """ return [] - @classmethod - def create_att_tail_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", - ) -> list["AttentionTailState"]: - """返回本模型的 tail 运行态列表；默认没有此类资源。 - - model_logical_layer_ids 描述当前模型拥有的全部全局逻辑层，具体模型自行选择参与层。 - """ - return [] - def _check_mem_size(self): self.max_total_token_num = self.att_state_manager.full_att_allocator.size diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 8a49767838..5df510e5de 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -41,7 +41,7 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig role = infer_state.att_state_role tail = infer_state.att_state_manager.get_state( model_role=role, - state_type_filter=f"kpool_c{self.pool_size}", + state_type_filter=f"compressor_indexer_w{self.pool_size}_s{self.pool_size}", logical_layer_id=self.layer_idx, ).get_layer_view(logical_layer_id=self.layer_idx) state = infer_state.att_state_manager.get_state( diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 61aa9602c4..5d33a188eb 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -22,7 +22,6 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.attention_state.bounded_att import BoundedAttentionState - from lightllm.common.attention_state.att_tail import AttentionTailState class Glm5NextTpPartModel(TpPartBaseModel): @@ -44,33 +43,57 @@ def create_bounded_att_states( tp_world_size: int = 1, device: str = "meta", ) -> list["BoundedAttentionState"]: - """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" - from lightllm.common.attention_state.bounded_att import create_kda_state, ConvSSMStateConfig + """分别创建 KDA 运行态和索引 key 的无重叠池化窗口，不分配大小页。""" + from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + WindowCompressorStateConfig, + create_kda_state, + create_window_compressor_state, + ) + from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" cfg = Glm5NextResourceConfig.from_model_config(config, args) - logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) - if not logical_layer_ids: - return [] - state = create_kda_state( - config=ConvSSMStateConfig( - logical_layer_ids=logical_layer_ids, - global_k_heads=cfg.global_linear_k_heads, - global_v_heads=cfg.global_linear_v_heads, - key_dim=cfg.head_linear_k_dim, - value_dim=cfg.head_linear_v_dim, - conv_width=cfg.conv_kernel_size - 1, - conv_dtype=cfg.conv_state_dtype, - ssm_dtype=cfg.ssm_state_dtype, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - ) - return [state] + states = [] + linear_layers = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) + if linear_layers: + states.append(create_kda_state( + config=ConvSSMStateConfig( + logical_layer_ids=linear_layers, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + )) + full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) + if full_layers: + assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" + states.append(create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=full_layers, + window_size=config["index_kpool"], + compress_stride=config["index_kpool"], + candidate_dim=config["index_head_dim"], + target="indexer", + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + )) + return states @classmethod def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: @@ -140,42 +163,6 @@ def create_full_att_states( ) return [state] - @classmethod - def create_att_tail_states( - cls, - model_config: dict[str, object], - args: "StartArgs | EasyDict", - model_logical_layer_ids: tuple[int, ...], - request_capacity: int = 1, - tp_rank: int = 0, - tp_world_size: int = 1, - device: str = "meta", - ) -> list["AttentionTailState"]: - """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" - from lightllm.common.attention_state.att_tail import create_kpool_state, KPoolStateConfig - from lightllm.utils.torch_dtype_utils import get_torch_dtype - - config = cls.normalize_model_config(model_config) - cfg = Glm5NextResourceConfig.from_model_config(config, args) - logical_layer_ids = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) - if not logical_layer_ids: - return [] - assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" - state = create_kpool_state( - config=KPoolStateConfig( - logical_layer_ids=logical_layer_ids, - pool_size=config["index_kpool"], - head_dim=config["index_head_dim"], - dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, - tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - ) - return [state] - def _init_att_backend(self): from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 1d2ac2d764..fe18d2a888 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -15,7 +15,7 @@ from easydict import EasyDict from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.common.attention_state.full_att import FullAttentionState - from lightllm.common.attention_state.att_tail import AttentionTailState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState class Glm5NextMTPModel(TpPartBaseModel): @@ -105,7 +105,7 @@ def create_full_att_states( return [state] @classmethod - def create_att_tail_states( + def create_bounded_att_states( cls, model_config: dict[str, object], args: "StartArgs | EasyDict", @@ -114,9 +114,9 @@ def create_att_tail_states( tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", - ) -> list["AttentionTailState"]: - """直接选择本模型需要 K-pool tail 的层，不依赖 full State 创建过程。""" - from lightllm.common.attention_state.att_tail import create_kpool_state, KPoolStateConfig + ) -> list["BoundedAttentionState"]: + """直接选择本模型的 K-pool 窗口层，不依赖 full State 创建过程。""" + from lightllm.common.attention_state.bounded_att import create_window_compressor_state, WindowCompressorStateConfig from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -124,11 +124,13 @@ def create_att_tail_states( if not logical_layer_ids: return [] assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" - state = create_kpool_state( - config=KPoolStateConfig( + state = create_window_compressor_state( + config=WindowCompressorStateConfig( logical_layer_ids=logical_layer_ids, - pool_size=config["index_kpool"], - head_dim=config["index_head_dim"], + window_size=config["index_kpool"], + compress_stride=config["index_kpool"], + candidate_dim=config["index_head_dim"], + target="indexer", dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index ebb06ce913..ecd8d555e3 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -683,7 +683,6 @@ def _restore_attention_page(self, node, page_kind): page_kind=page_kind, page_slot_ids=[slot], ) - resources.reset_attention_tail_state([self.req_idx]) def _hybrid_match_radix_cache(self): assert ( @@ -800,7 +799,6 @@ def _hybrid_match_radix_cache(self): if self.cur_kv_len == 0: # 说明没有任何命中 g_infer_context.att_state_manager.reset_bounded_att_state([self.req_idx]) - g_infer_context.att_state_manager.reset_attention_tail_state([self.req_idx]) return def is_master_req(self): diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index e2aaa921a6..942f8002b9 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -47,7 +47,7 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: request_capacity=1, device="meta", ) - # tail 不占 CPU 页；无需创建 tail，也不需要实际模型或临时 CUDA 运行态。 + # 无重叠的池化窗口自行返回空 CPU 区域；这里只创建 meta State，不初始化 CUDA。 for state in (*full_states, *bounded_states): manager.register_state(role, state) logical_layer_start += count diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 961e65042e..69379e4150 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -78,7 +78,6 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP allocated_slots = list(source_slots) if hybrid else [] source_req = executor._alloc_req_indexes(1) resources.reset_bounded_att_state(source_req.tolist()) - resources.reset_attention_tail_state(source_req.tolist()) tokens = np.random.default_rng(42).integers( low=100, high=10000, @@ -107,8 +106,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP cpu_page_buffers=cpu, cpu_page_indices=pages, layouts=resources.get_cpu_page_layouts(), - page_kind="big" if hybrid else None, - page_slot_ids=source_slots, + big_page_slot_ids=source_slots if hybrid else None, ) # Host-side checkpoint repacking consumes the CPU page on the next call. torch.cuda.synchronize() @@ -138,15 +136,16 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP for kind in ("gpu_pages", "cpu_pages", "dp_copy", "pd_buffer"): request = executor._alloc_req_indexes(1) resources.reset_bounded_att_state(request.tolist()) - resources.reset_attention_tail_state(request.tolist()) executor._ensure_req_kv_capacity(request, torch.tensor([prefix_len], dtype=torch.int32)) indices = model.req_manager.req_to_token_indexs[int(request[0]), :prefix_len].clone() assert not set(indices.tolist()) & set(source_indices.tolist()) if kind == "gpu_pages": resources.copy_full_att_tokens(source_indices, indices) if hybrid: - resources.restore_bounded_att_runtime_from_pages(request.tolist(), page_kind="big", page_slot_ids=[source_slots[-1]]) - resources.reset_attention_tail_state(request.tolist()) + resources.restore_bounded_att_runtime_from_pages( + req_indices=request.tolist(), state_token_lens=[prefix_len], + page_kind="big", page_slot_ids=[source_slots[-1]], + ) elif kind == "cpu_pages": slots = big.alloc(len(pages)) if hybrid else None if hybrid: @@ -157,18 +156,18 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=indices, layouts=resources.get_cpu_page_layouts(), big_page_slot_ids=slots, - req_idx=int(request[0]) if getattr(resources, "cpu_tail_restore", False) else None, + req_idx=int(request[0]) if hybrid else None, + state_token_len=prefix_len if hybrid else None, ) - if hybrid: - resources.restore_bounded_att_runtime_from_pages(request.tolist(), page_kind="big", page_slot_ids=[slots[-1]]) elif kind == "dp_copy": resources.copy_from_dp( - resources, - source_indices, - indices, - mappings=resources.dp_copy_mappings(resources), - src_req_indices=source_req.tolist() if hybrid else None, - dst_req_indices=request.tolist() if hybrid else None, + remote_att_manager=resources, + src_token_indices=source_indices, + dst_token_indices=indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=int(source_req[0]), + dst_req_idx=int(request[0]), ) else: resources.read_from_pd_page( diff --git a/unit_tests/common/att_resources/conftest.py b/unit_tests/common/att_resources/conftest.py index 32bb70d93e..829c82ac3f 100644 --- a/unit_tests/common/att_resources/conftest.py +++ b/unit_tests/common/att_resources/conftest.py @@ -24,9 +24,9 @@ def allocator_environment(kv_allocator_environment: None, monkeypatch: pytest.Mo state_classes = [FullAttentionState, BoundedAttentionState] for state_class in state_classes: state_classes.extend(state_class.__subclasses__()) - getter = state_class.get_cpu_transfer_regions - if hasattr(getter, "cache_clear"): - getter.cache_clear() + for getter in (state_class.get_cpu_transfer_regions, state_class.get_pd_transfer_regions): + if hasattr(getter, "cache_clear"): + getter.cache_clear() get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index 950f7cde10..61320d1cef 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -9,17 +9,15 @@ from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator from lightllm.common.attention_state.full_att import MLAAttentionConfig, MLAFullAttentionState from lightllm.common.attention_state.bounded_att import GDNState -from lightllm.common.attention_state.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state from unit_tests.common.att_resources.test_bounded_state import make_state -from unit_tests.common.att_resources.test_tail_state import make_tail if TYPE_CHECKING: import torch class MainDraftTestResourceManager(AttentionStateManager): - def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): + def __init__(self, device: str | torch.device | int = "cpu"): super().__init__(request_capacity=3) self.register_state("main", make_state(device=device)) draft_state = GDNState( @@ -41,16 +39,6 @@ def __init__(self, device: str | torch.device | int = "cpu", with_tail: bool = T ) prepared_full_state.allocate_storage(self.main_gqa_fp32_state.capacity) self.register_state("draft0", prepared_full_state) - self.att_tail_type = f"kpool_c{tail_pool_size}_fp32" - if with_tail: - tail = make_tail(device=device) - state = KPoolAttentionTailState( - config=replace(tail.config, pool_size=tail_pool_size), - request_capacity=tail.request_capacity, - mtp_step=tail.mtp_step, - device=tail.tail_buffer.device, - ) - self.register_state("main", state) self.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) self.bounded_att_big_page_allocator = StatePageAllocator(3) self.bounded_att_small_page_allocator = StatePageAllocator(2) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 78ae9342d5..94d3daba49 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -9,6 +9,7 @@ import torch from lightllm.common.attention_state.types import TensorField +from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState from lightllm.common.attention_state.full_att import ( FullAttentionState, GQAAttentionConfig, MLAAttentionConfig, GroupQuantizedAttentionConfig, StaticFP8AttentionConfig, PackedMLAAttentionConfig, @@ -99,7 +100,7 @@ def owned_states(attributes: Mapping[str, object], state_class: type | tuple[typ def main_linear_state(resources: AttentionStateManager): - (state,) = resources.get_bounded_attention_states("main") + (state,) = (state for state in resources.get_bounded_attention_states("main") if hasattr(state, "conv_state")) return state @@ -207,7 +208,7 @@ def build_hybrid_resources( from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state - from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state from lightllm.models.glm5_next.config import Glm5NextResourceConfig manager = AttentionStateManager(request_capacity=request_capacity) @@ -241,12 +242,14 @@ def build_hybrid_resources( continue manager.register_state(role, state) if glm: - tail = create_kpool_state( - config=KPoolStateConfig( + window = create_window_compressor_state( + config=WindowCompressorStateConfig( logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), - pool_size=config.index_kpool, - head_dim=config.index_head_dim, - dtype=state.config.dtype, + window_size=config.index_kpool, + compress_stride=config.index_kpool, + candidate_dim=config.index_head_dim, + target="indexer", + dtype=config.conv_state_dtype, tp_rank=tp_rank, tp_world_size=config.tp_world_size, ), @@ -254,12 +257,25 @@ def build_hybrid_resources( mtp_step=mtp_step, device="cuda", ) - manager.register_state(role, tail) + manager.register_state(role, window) manager.full_att_allocator = full_allocator manager.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) manager.bounded_att_small_page_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None manager.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 manager.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 - manager.cpu_tail_restore = glm manager.finish_state_registration() return manager + + +def window_states( + manager: AttentionStateManager, + role: str | None = None, +) -> tuple[WindowCompressorState, ...]: + return tuple(state for state in manager.get_bounded_attention_states(role) if isinstance(state, WindowCompressorState)) + + +def linear_states( + manager: AttentionStateManager, + role: str | None = None, +) -> tuple[GDNState | KDAState, ...]: + return tuple(state for state in manager.get_bounded_attention_states(role) if isinstance(state, (GDNState, KDAState))) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index fe3d1c8037..fc2ec7faeb 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -22,7 +22,7 @@ from lightllm.common.attention_state.utils import deepcopy_result from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState -from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.common.attention_state.full_att import ( FullAttentionState, @@ -212,11 +212,11 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ) - tail = KPoolAttentionTailState( - config=KPoolStateConfig( + tail = WindowCompressorState( + config=WindowCompressorStateConfig( logical_layer_ids=(1,), - pool_size=4, - head_dim=3, + window_size=4, compress_stride=4, target="indexer", + candidate_dim=3, dtype=torch.bfloat16, ), request_capacity=3, @@ -232,7 +232,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( conv = reference.conv_state[:, 0, :, : config.conv_width] ssm = reference.ssm_state[:, 0] bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() - tail_bytes = len(tail.config.logical_layer_ids) * tail.config.pool_size * 2 * tail.config.head_dim * tail.config.dtype.itemsize + tail_bytes = len(tail.config.logical_layer_ids) * tail.config.window_size * 2 * tail.config.candidate_dim * tail.config.dtype.itemsize manager = AttentionStateManager(request_capacity=3) manager.register_state("main", reference) manager.register_state("main", tail) @@ -241,10 +241,10 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( assert manager.get_pd_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 reference.conv_state.fill_(2) reference.ssm_state.fill_(3) - tail.tail_buffer.fill_(4) + tail.buffer.fill_(4) pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") layouts = manager.get_pd_runtime_page_layouts() - assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "tail"} + assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "compressor"} manager.write_to_pd_page( pd_page_buffers=pages, pd_page_index=1, @@ -255,7 +255,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( ) reference.conv_state.zero_() reference.ssm_state.zero_() - tail.tail_buffer.zero_() + tail.buffer.zero_() manager.read_from_pd_page( pd_page_buffers=pages, pd_page_index=1, @@ -266,10 +266,10 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( ) assert (reference.conv_state[:, 0, :, : config.conv_width] == 2).all() assert (reference.ssm_state[:, 0] == 3).all() - expected_tail = torch.zeros_like(tail.tail_buffer[:, 0]) + expected_tail = torch.zeros_like(tail.buffer[:, 0]) for position in range(4, 8): - expected_tail[:, position % tail.tail_buffer.shape[2]].fill_(4) - assert torch.equal(tail.tail_buffer[:, 0], expected_tail) + expected_tail[:, position % tail.buffer.shape[2]].fill_(4) + assert torch.equal(tail.buffer[:, 0], expected_tail) assert not pages[0].any() for state in (reference, tail): monkeypatch.setattr( @@ -453,6 +453,8 @@ def test_pd_runtime_capacity_includes_field_padding( full_page_tokens: int, ): class PaddedGDNState(GDNState): + @deepcopy_result + @lru_cache(maxsize=None) def get_pd_transfer_regions(self): return tuple( replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) if region.field == "ssm" else region @@ -603,11 +605,11 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( state = cls(config, request_capacity=1, device="meta") manager.register_state(role, state) if with_tail: - state = KPoolAttentionTailState( - config=KPoolStateConfig( + state = WindowCompressorState( + config=WindowCompressorStateConfig( logical_layer_ids=(0,), - pool_size=4, - head_dim=1, + window_size=4, compress_stride=4, target="indexer", + candidate_dim=1, dtype=torch.float32, ), request_capacity=1, diff --git a/unit_tests/common/att_resources/test_kpool_transfer.py b/unit_tests/common/att_resources/test_kpool_transfer.py deleted file mode 100644 index 334d68f071..0000000000 --- a/unit_tests/common/att_resources/test_kpool_transfer.py +++ /dev/null @@ -1,332 +0,0 @@ -"""K-pool 融合传输：环形寻址、完整位模式、窗口外槽位保留及编译缓存。""" - -from dataclasses import replace - -import pytest -import torch - -from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig -from lightllm.common.attention_state.att_tail import kpool_transfer -from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("dtype,head_dim,byte_start", [ - (torch.bfloat16, 128, 16), - (torch.float16, 33, 2), - (torch.float32, 128, 4), -]) -@pytest.mark.parametrize("page_device", ["cuda", "cpu"]) -def test_kpool_single_launch_preserves_bits_and_other_slots( - dtype: torch.dtype, - head_dim: int, - byte_start: int, - page_device: str, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config = KPoolStateConfig( - logical_layer_ids=(0, 3, 7), - pool_size=4, - head_dim=head_dim, - dtype=dtype, - ) - source = KPoolAttentionTailState( - config=config, - request_capacity=3, - mtp_step=1, - device='cuda', - ) - target = KPoolAttentionTailState( - config=config, - request_capacity=3, - mtp_step=5, - device='cuda', - ) - source.tail_buffer.view(torch.uint8).random_(0, 256) - original = source.tail_buffer.clone() - target.tail_buffer.view(torch.uint8).fill_(173) - region = source.get_pd_transfer_regions()[0] - # 同时验证非默认的外部轴顺序和字段起点；不能假定 window/channel 连续排列。 - region = replace( - region, - axes=("channel", "layer", "window"), - shape=(2 * head_dim, 3, 4), - byte_start=byte_start, - byte_end=byte_start + region.byte_end, - ) - layout = StateTransferLayout( - state_name=source.state_name, - tokens_per_page=None, - byte_start=byte_start, - byte_end=region.byte_end, - regions=(region,), - ) - pages = torch.full( - size=(3, layout.byte_end + 16), - fill_value=253, - dtype=torch.uint8, - device=page_device, - pin_memory=page_device == "cpu", - ) - kernel = kpool_transfer._transfer_kpool_pages - launches = [] - - class CountLaunches: - def __getitem__(self, grid: tuple[int]): - def launch(**kwargs: object) -> None: - launches.append((grid, kwargs["WRITE"])) - kernel[grid](**kwargs) - return launch - - def reject_copy(*args: object, **kwargs: object) -> None: - pytest.fail("K-pool transfer must copy in one Triton launch") - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with monkeypatch.context() as patch: - patch.setattr(kpool_transfer, "_transfer_kpool_pages", CountLaunches()) - for method in ("copy_", "fill_", "zero_", "index_select", "index_copy_", "roll"): - patch.setattr(torch.Tensor, method, reject_copy) - with torch.cuda.stream(stream): - source.write_to_pd_page( - req_idx=1, - pd_page_buffers=pages, - pd_page_index=1, - state_token_len=11, - layout=layout, - ) - target.read_from_pd_page( - pd_page_buffers=pages, - req_idx=2, - pd_page_index=1, - state_token_len=11, - layout=layout, - ) - stream.synchronize() - assert len(launches) == 2 and [item[1] for item in launches] == [True, False] - expected = torch.empty_like(target.tail_buffer[:, 2]) - expected.view(torch.uint8).fill_(173) - for position in range(7, 11): - expected[:, position % 9] = source.tail_buffer[:, 1, position % 5] - assert torch.equal(target.tail_buffer[:, 2].view(torch.uint8), expected.view(torch.uint8)) - assert (target.tail_buffer[:, :2].view(torch.uint8) == 173).all() - assert torch.equal(source.tail_buffer.view(torch.uint8), original.view(torch.uint8)) - assert (pages[[0, 2]] == 253).all() - assert (pages[1, :byte_start] == 253).all() and (pages[1, layout.byte_end:] == 253).all() - - -@pytest.mark.parametrize("local_circular", [False, True]) -@pytest.mark.parametrize("external_circular", [False, True]) -def test_kpool_runtime_axis_overrides_are_bidirectional( - local_circular: bool, - external_circular: bool, -) -> None: - config = KPoolStateConfig( - logical_layer_ids=(0, 2), - pool_size=4, - head_dim=128, - dtype=torch.bfloat16, - ) - state = KPoolAttentionTailState( - config=config, - request_capacity=2, - mtp_step=3, - device='cuda', - ) - state.tail_buffer.view(torch.uint8).random_(0, 256) - original = state.tail_buffer[:, 1].clone() - regions = state.get_pd_transfer_regions() - layout = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=None, - byte_start=0, - byte_end=regions[0].byte_end, - regions=regions, - ) - pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - window = AxisMapping( - axis="window", - local_start=5 if local_circular else 1, - external_start=3 if external_circular else 0, - length=4, - local_circular=local_circular, - external_circular=external_circular, - ) - kpool_transfer.transfer_kpool_pages( - tail=state.tail_buffer[:, 1], - page_buffers=pages, - page_index=0, - state_token_len=11, - layout=layout, - write=True, - axis_mapping_overrides=(window,), - grid_num=1, - ) - external = pages[0].view(config.dtype).view(2, 4, 256) - for index in range(4): - assert torch.equal( - external[:, (window.external_start + index) % 4].view(torch.uint8), - original[:, (window.local_start + index) % 7].view(torch.uint8), - ) - state.tail_buffer[:, 1].fill_(37) - kpool_transfer.transfer_kpool_pages( - tail=state.tail_buffer[:, 1], - page_buffers=pages, - page_index=0, - state_token_len=11, - layout=layout, - write=False, - axis_mapping_overrides=(window,), - grid_num=1, - ) - expected = torch.full_like(original, 37) - for index in range(4): - slot = (window.local_start + index) % 7 - expected[:, slot] = original[:, slot] - assert torch.equal(state.tail_buffer[:, 1].view(torch.uint8), expected.view(torch.uint8)) - assert layout.regions[0].axis_mappings[0].local_start == 0 - - -def test_kpool_tp_replica_skips_writes_but_restores_tail() -> None: - config = KPoolStateConfig( - logical_layer_ids=(0,), - pool_size=4, - head_dim=128, - dtype=torch.bfloat16, - tp_rank=1, - tp_world_size=2, - ) - state = KPoolAttentionTailState( - config=config, - request_capacity=2, - mtp_step=2, - device='cuda', - ) - regions = state.get_pd_transfer_regions() - layout = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=None, - byte_start=0, - byte_end=regions[0].byte_end, - regions=regions, - ) - pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - pages.view(torch.bfloat16).fill_(13) - state.write_to_pd_page( - req_idx=1, - pd_page_buffers=pages, - pd_page_index=0, - state_token_len=11, - layout=layout, - ) - assert (pages.view(torch.bfloat16) == 13).all() - state.read_from_pd_page( - pd_page_buffers=pages, - req_idx=1, - pd_page_index=0, - state_token_len=11, - layout=layout, - ) - expected = torch.zeros_like(state.tail_buffer[:, 1]) - for position in range(7, 11): - expected[:, position % 6].fill_(13) - assert torch.equal(state.tail_buffer[:, 1], expected) - - -def test_kpool_window_positions_do_not_create_kernel_variants() -> None: - config = KPoolStateConfig( - logical_layer_ids=(0,), - pool_size=4, - head_dim=128, - dtype=torch.bfloat16, - ) - state = KPoolAttentionTailState( - config=config, - request_capacity=2, - mtp_step=5, - device='cuda', - ) - regions = state.get_pd_transfer_regions() - layout = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=None, - byte_start=0, - byte_end=regions[0].byte_end, - regions=regions, - ) - pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - kernel = kpool_transfer._transfer_kpool_pages - cache_size = None - for length in (*range(30), (1 << 33) + 7): - state.write_to_pd_page( - req_idx=1, - pd_page_buffers=pages, - pd_page_index=0, - state_token_len=length, - layout=layout, - ) - state.read_from_pd_page( - pd_page_buffers=pages, - req_idx=0, - pd_page_index=0, - state_token_len=length, - layout=layout, - ) - current_size = len(kernel.device_caches[torch.cuda.current_device()][0]) - if cache_size is None: - cache_size = current_size - assert current_size == cache_size - torch.cuda.synchronize() - - -@pytest.mark.parametrize("length", [0, 3, 515]) -def test_kpool_large_window_with_one_program(length: int) -> None: - config = KPoolStateConfig( - logical_layer_ids=(0, 3, 9, 12), - pool_size=128, - head_dim=128, - dtype=torch.bfloat16, - ) - source = KPoolAttentionTailState( - config=config, - request_capacity=2, - mtp_step=2, - device="cuda", - ) - target = KPoolAttentionTailState( - config=config, - request_capacity=2, - mtp_step=7, - device="cuda", - ) - source.tail_buffer.view(torch.uint8).random_(0, 256) - target.tail_buffer.fill_(37) - regions = source.get_pd_transfer_regions() - layout = StateTransferLayout( - state_name=source.state_name, - tokens_per_page=None, - byte_start=0, - byte_end=regions[0].byte_end, - regions=regions, - ) - pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - for state, write in ((source, True), (target, False)): - kpool_transfer.transfer_kpool_pages( - tail=state.tail_buffer[:, 1], - page_buffers=pages, - page_index=0, - state_token_len=length, - layout=layout, - write=write, - grid_num=1, - ) - expected = torch.full_like(target.tail_buffer[:, 1], 37) - for position in range(length - 128, length): - if position < 0: - expected[:, position % 135].zero_() - else: - expected[:, position % 135] = source.tail_buffer[:, 1, position % 130] - assert torch.equal(target.tail_buffer[:, 1].view(torch.uint8), expected.view(torch.uint8)) - assert (target.tail_buffer[:, 0] == 37).all() diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index f153ffd96e..3d7fc0c3d2 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -10,7 +10,7 @@ from lightllm.common.attention_state.full_att import GQAAttentionConfig, FullAttentionState, GQAFullAttentionState from lightllm.common.attention_state.bounded_att import KDAState, GDNState from unit_tests.common.att_resources.test_bounded_state import make_state -from unit_tests.common.att_resources.test_tail_state import make_tail +from unit_tests.common.att_resources.test_window_compressor_state import make_window from unit_tests.common.att_resources.test_capacity import definition from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.models.llama.model import LlamaTpPartModel @@ -148,7 +148,7 @@ def test_mla_selection_excludes_same_layer_tail() -> None: from dataclasses import replace manager = AttentionStateManager(request_capacity=3) - tail = make_tail() + tail = make_window() template = definition("mla") state = type(template)( layout_config=replace(template.config, logical_layer_groups=((3, 7),)), device="meta", @@ -159,7 +159,7 @@ def test_mla_selection_excludes_same_layer_tail() -> None: assert tail.contains_logical_layer(3) and tail.contains_logical_layer(7) assert not tail.contains_logical_layer(0) assert manager.get_state("main", "mla", 7) is state - assert manager.get_state("main", "kpool", 7) is tail + assert manager.get_state("main", "compressor", 7) is tail assert manager.get_state( model_role="main", state_type_filter="", @@ -180,15 +180,15 @@ def test_query_matches_complete_model_role(role: str) -> None: assert manager.get_state(role, "gqa", 0).state_name == f"{role}_gqa_bf16_state" -@pytest.mark.parametrize("algorithm", ["kda", "gdn", "kpool"]) +@pytest.mark.parametrize("algorithm", ["kda", "gdn", "compressor"]) def test_precision_variants_have_distinct_names_and_can_coexist(algorithm: str) -> None: from dataclasses import replace manager = AttentionStateManager(request_capacity=3) - template = make_tail() if algorithm == "kpool" else make_state(cls=KDAState if algorithm == "kda" else GDNState) + template = make_window() if algorithm == "compressor" else make_state(cls=KDAState if algorithm == "kda" else GDNState) states = [] for dtype in (torch.bfloat16, torch.float32): - config = replace(template.config, **({"dtype": dtype} if algorithm == "kpool" else {"conv_dtype": dtype})) + config = replace(template.config, **({"dtype": dtype} if algorithm == "compressor" else {"conv_dtype": dtype})) state = type(template)(config=config, request_capacity=3, device="cpu") manager.register_state("main", state) states.append(state) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 51a56fd05a..9da30911eb 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -14,10 +14,8 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.types import PageKind, StateTransferLayout from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager -from lightllm.common.attention_state.att_tail import KPoolAttentionTailState from unit_tests.common.att_resources.test_full_state import full_state, layout_for from unit_tests.common.att_resources.test_bounded_state import page_layout -from unit_tests.common.att_resources.test_tail_state import make_tail, tail_layout if TYPE_CHECKING: from easydict import EasyDict @@ -31,8 +29,8 @@ def manager_cpu_page_args(allocator_environment: None, monkeypatch: pytest.Monke monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 4) -def composition(device: str | torch.device | int = "cpu", with_tail: bool = True, tail_pool_size: int = 4): - return MainDraftTestResourceManager(device, with_tail=with_tail, tail_pool_size=tail_pool_size) +def composition(device: str | torch.device | int = "cpu"): + return MainDraftTestResourceManager(device) def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): @@ -41,8 +39,7 @@ def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] if kind != "kv": layouts += [page_layout(state) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)] - if kind == "att_state": - layouts += [tail_layout(state) for state in manager.get_attention_tail_states("main")] + offset = 0 result = {} for layout in layouts: @@ -71,9 +68,6 @@ def init_source(manager: AttentionStateManager): state.conv_state[:, 1].fill_(23 + i) state.ssm_state[:, :3].fill_(31 + i) state.ssm_state[:, 3:6].fill_(47 + i) - for state in manager.get_attention_tail_states("main"): - state.tail_buffer[:, 0].fill_(3) - state.tail_buffer[:, 1].fill_(9) slots = manager.bounded_att_big_page_allocator.alloc(2) manager.save_bounded_att_runtime_to_pages( req_indices=[0, 1], @@ -116,29 +110,11 @@ def fail(*args: object, **kwargs: object): assert manager.bounded_att_big_page_allocator.free_count == 3 -def test_disabled_domains(): - manager = AttentionStateManager(request_capacity=3) - manager.register_state("main", full_state()) - manager.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) - manager.finish_state_registration() - manager.reset_bounded_att_state([0]) - manager.reset_attention_tail_state([0]) - with pytest.raises(AssertionError, match="disabled"): - manager.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[4], - page_kind="big", - page_slot_ids=[0], - ) - empty_manager = AttentionStateManager(request_capacity=3) - empty_manager.finish_state_registration() - with pytest.raises(AssertionError, match="no full-attention"): - empty_manager.copy_full_att_tokens(torch.tensor([0], dtype=torch.long), torch.tensor([1], dtype=torch.long)) @pytest.mark.parametrize("operation", ["save_bounded_att_runtime_to_pages", "restore_bounded_att_runtime_from_pages"]) def test_page_operations_forward_each_request_length(operation: str, monkeypatch: pytest.MonkeyPatch): - manager = composition(with_tail=False) + manager = composition() calls = [] def record_lengths( @@ -162,7 +138,7 @@ def record_lengths( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for token indices") def test_cpu_restore_forwards_absolute_length_for_partial_prefix(monkeypatch: pytest.MonkeyPatch): - manager = composition(with_tail=False) + manager = composition() layouts = layouts_for(manager) manager.bounded_att_cpu_load_big_page_slot = 2 for state in manager.get_full_attention_states(): @@ -195,7 +171,7 @@ def record_restore( @pytest.mark.parametrize("operation", ["write_to_pd_page", "read_from_pd_page"]) def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): - manager = composition(with_tail=False) + manager = composition() layouts = layouts_for(manager, kind="att_state") lengths = [] @@ -221,106 +197,14 @@ def record_length( assert lengths == [(2, 29)] * len(manager.get_bounded_attention_states()) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") -def test_reset_tail_clears_all_tail_states_and_preserves_linear(): - first = make_tail() - second = KPoolAttentionTailState( - config=replace(first.config, pool_size=3), - request_capacity=first.request_capacity, - mtp_step=first.mtp_step, - device=first.tail_buffer.device, - ) - first.tail_buffer.fill_(5) - second.tail_buffer.fill_(7) - tail_manager = AttentionStateManager(request_capacity=3) - tail_manager.register_state("main", first) - tail_manager.register_state("main", second) - tail_manager.finish_state_registration() - tail_manager.reset_attention_tail_state([0, 1]) - assert not first.tail_buffer[:, :2].any() and not second.tail_buffer[:, :2].any() - assert (first.tail_buffer[:, 2] == 5).all() and (second.tail_buffer[:, 2] == 7).all() - manager = composition(device="cuda") - init_source(manager) - # Preserve accepted-state cursors as well as both algorithms' runtime and - # cached pages. A full reset here would silently corrupt a prefix restore. - for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state): - state.mtp_state_index[1] = 2 - tensors = [ - *full_state_test_buffers(manager.main_gqa_fp32_state).values(), - *full_state_test_buffers(manager.draft0_mla_fp32_state).values(), - ] - for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state): - tensors.extend(( - state.conv_state, state.ssm_state, state.mtp_state_index, - state.big_pages_pool.conv_state, state.big_pages_pool.ssm_state, - state.small_pages_pool.conv_state, state.small_pages_pool.ssm_state, - )) - preserved = [(tensor, tensor.contiguous().view(torch.uint8).clone()) for tensor in tensors] - manager.reset_attention_tail_state([1]) - assert not getattr(manager, f"main_{manager.att_tail_type}_state").tail_buffer[:, 1].any() - for tensor, snapshot in preserved: - assert torch.equal(tensor.contiguous().view(torch.uint8), snapshot) - - -@pytest.mark.parametrize("device", ["cuda"]) -def test_cpu_all_linear_pages_and_tail_reset_at_final_boundary(device: str | torch.device | int): - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA unavailable") - source, target = composition(device), composition(device) - slots = init_source(source) - layouts = layouts_for(source) - pages = torch.zeros( - (3, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=device == "cuda" - ) - source.write_req_to_cpu( - token_indices=torch.arange(8, dtype=torch.long, device=device), - cpu_page_buffers=pages, - cpu_page_indices=[2, 0], - layouts=layouts, - big_page_slot_ids=slots, - ) - if device == "cuda": - torch.cuda.synchronize() - tail = getattr(target, f"main_{target.att_tail_type}_state") - tail.tail_buffer.fill_(37) - assert tail.state_name not in layouts - target_slots = target.bounded_att_big_page_allocator.alloc(2) - target.read_req_from_cpu( - cpu_page_buffers=pages, - cpu_page_indices=[2, 0], - token_indices=torch.arange( - start=4, - end=12, - dtype=torch.long, - device=device, - ), - state_token_len=8, - layouts=layouts, - big_page_slot_ids=target_slots, - req_idx=2, - ) - for src, dst in zip( - (source.main_kda_conv_fp32_ssm_fp32_state, source.draft0_gdn_conv_fp32_ssm_fp32_state), (target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state) - ): - for src_slot, dst_slot in zip(slots, target_slots): - assert torch.equal(src.big_pages_pool.conv_state[src_slot], dst.big_pages_pool.conv_state[dst_slot]) - assert torch.equal(src.big_pages_pool.ssm_state[src_slot], dst.big_pages_pool.ssm_state[dst_slot]) - assert (tail.tail_buffer[:, :2] == 37).all() - assert not tail.tail_buffer[:, 2].any() - for i, state in enumerate((target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state)): - assert (state.conv_state[:, 2, :, :2] == 23 + i).all() - assert (state.ssm_state[:, 6] == 47 + i).all() - for src, dst in zip( - (source.main_gqa_fp32_state, source.draft0_mla_fp32_state), - (target.main_gqa_fp32_state, target.draft0_mla_fp32_state), - ): - assert torch.equal(read_kv_storage(src)[:, :8], read_kv_storage(dst)[:, 4:12]) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("token_count", [2, 4, 6, 8]) def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: int): - source, target = composition("cuda", with_tail=False), composition("cuda", with_tail=False) + source, target = composition("cuda"), composition("cuda") source_slots = init_source(source) layouts = layouts_for(source) page_count = (token_count + 3) // 4 @@ -381,193 +265,18 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i torch.testing.assert_close(read_kv_storage(dst)[:, 4:4 + token_count], read_kv_storage(src)[:, :token_count]) -@pytest.mark.parametrize("tail_pool_size", [4, 8]) -def test_cpu_restore_resets_selected_tail_after_all_pages(tail_pool_size: int): - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - source, target = composition(device="cuda", tail_pool_size=tail_pool_size), composition(device="cuda", tail_pool_size=tail_pool_size) - slots = init_source(source) - layouts = layouts_for(source) - pages = torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) - source.write_req_to_cpu( - token_indices=torch.arange(8, dtype=torch.long, device="cuda"), - cpu_page_buffers=pages, - cpu_page_indices=[0, 1], - layouts=layouts, - big_page_slot_ids=slots, - ) - torch.cuda.synchronize() - target_slots = target.bounded_att_big_page_allocator.alloc(2) - tail = getattr(target, f"main_{target.att_tail_type}_state") - tail.tail_buffer.fill_(37) - # 缓存流程保证最终前缀 L=8 是可恢复边界；中间页 L=4 不单独重置 tail。 - target.read_req_from_cpu( - cpu_page_buffers=pages, - cpu_page_indices=[0, 1], - token_indices=torch.arange(8, dtype=torch.long, device="cuda"), - state_token_len=8, - layouts=layouts, - big_page_slot_ids=target_slots, - req_idx=0, - ) - assert not getattr(target, f"main_{target.att_tail_type}_state").tail_buffer[:, 0].any() - assert (getattr(target, f"main_{target.att_tail_type}_state").tail_buffer[:, 1:] == 37).all() - - -@pytest.mark.parametrize("device", ["cuda"]) -def test_pd_composition_and_dp_copy_keep_resource_identity(device: str | torch.device | int, pd_args: EasyDict): - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA unavailable") - source, target, dp_target = composition(device), composition(device), composition(device) - init_source(source) - for kind in ("kv", "att_state"): - layouts = layouts_for(source, kind=kind) - page = torch.empty(max(layout.byte_end for layout in layouts.values()), dtype=torch.uint8, device=device) - src_args = {"token_indices": torch.tensor([0, 2], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 1, "state_token_len": 8} - dst_args = {"token_indices": torch.tensor([5, 7], dtype=torch.long, device=device)} if kind == "kv" else {"req_idx": 2, "state_token_len": 8} - source.write_to_pd_page( - pd_page_buffers=page.unsqueeze(0), - pd_page_index=0, - kind=kind, - layouts=layouts, - **src_args, - ) - target.read_from_pd_page( - pd_page_buffers=page.unsqueeze(0), - pd_page_index=0, - kind=kind, - layouts=layouts, - **dst_args, - ) - dp_tail = dp_target.get_attention_tail_states()[0] - dp_tail.tail_buffer.fill_(19) - dp_target.copy_from_dp( - remote_att_manager=source, - src_token_indices=torch.tensor([0, 2], dtype=torch.int32, device=device), - dst_token_indices=torch.tensor([5, 7], dtype=torch.long, device=device), - src_big_page_ids=[], - dst_big_page_ids=[], - src_req_idx=1, - dst_req_idx=2, - ) - for a, b in zip( - (target.main_gqa_fp32_state, target.draft0_mla_fp32_state), - (dp_target.main_gqa_fp32_state, dp_target.draft0_mla_fp32_state), - ): - assert torch.equal(read_kv_storage(a)[:, [5, 7]], read_kv_storage(b)[:, [5, 7]]) - for a, b in zip( - (target.main_kda_conv_fp32_ssm_fp32_state, target.draft0_gdn_conv_fp32_ssm_fp32_state), (dp_target.main_kda_conv_fp32_ssm_fp32_state, dp_target.draft0_gdn_conv_fp32_ssm_fp32_state) - ): - assert torch.equal(a.conv_state[:, 2], b.conv_state[:, 2]) - assert torch.equal(a.ssm_state[:, 6], b.ssm_state[:, 6]) - pd_tail = target.get_attention_tail_states()[0].tail_buffer[:, 2] - expected_tail = torch.zeros_like(pd_tail) - for position in range(4, 8): - expected_tail[:, position % pd_tail.shape[1]].fill_(9) - assert torch.equal(pd_tail, expected_tail) - assert not dp_tail.tail_buffer[:, 2].any() - assert (dp_tail.tail_buffer[:, :2] == 19).all() - assert (source.get_attention_tail_states()[0].tail_buffer[:, 1] == 9).all() -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") -def test_dp_copy_resets_only_selected_tail_requests(pd_args: EasyDict): - source, target = composition(device="cuda"), composition(device="cuda") - init_source(source) - tails = target.get_attention_tail_states() - for state in tails: - state.tail_buffer.fill_(19) - target.copy_from_dp( - remote_att_manager=source, - src_token_indices=torch.empty(0, dtype=torch.long), - dst_token_indices=torch.empty(0, dtype=torch.long), - src_big_page_ids=[], - dst_big_page_ids=[], - src_req_idx=1, - dst_req_idx=2, - ) - assert all(not state.tail_buffer[:, 2].any() for state in tails) - assert all((state.tail_buffer[:, :2] == 19).all() for state in tails) - - -@pytest.mark.parametrize("copy_tokens,copy_pages", [(True, True), (True, False), (False, True)]) -def test_dp_copy_includes_selected_runtime_and_main_draft_pages(copy_tokens: bool, copy_pages: bool): - if not torch.cuda.is_available(): - pytest.skip("CUDA unavailable") - device = "cuda" - source, target = composition(device), composition(device) - init_source(source) - for state in target.get_full_attention_states(): - for buffer in full_state_test_buffers(state).values(): - buffer.fill_(-1) - for index, state in enumerate(source.get_bounded_attention_states()): - for page in range(3): - state.big_pages_pool.conv_state[page].fill_(10 * index + page) - state.big_pages_pool.ssm_state[page].fill_(20 * index + page) - for state in target.get_bounded_attention_states(): - state.big_pages_pool.conv_state.fill_(-1) - state.big_pages_pool.ssm_state.fill_(-2) - state.small_pages_pool.conv_state.fill_(31) - state.small_pages_pool.ssm_state.fill_(32) - state.conv_state.fill_(41) - state.ssm_state.fill_(42) - state.mtp_state_index.fill_(1) - for state in target.get_attention_tail_states(): - state.tail_buffer.fill_(51) - - target.copy_from_dp( - remote_att_manager=source, - src_token_indices=torch.tensor([0, 2] if copy_tokens else [], dtype=torch.long, device=device), - dst_token_indices=torch.tensor([5, 7] if copy_tokens else [], dtype=torch.long, device=device), - src_big_page_ids=[2, 0] if copy_pages else [], - dst_big_page_ids=[0, 1] if copy_pages else [], - src_req_idx=1, - dst_req_idx=2, - ) - for state in target.get_full_attention_states(): - src_state = getattr(source, state.state_name) - values = read_kv_storage(state) - if copy_tokens: - assert torch.equal(values[:, [5, 7]], read_kv_storage(src_state)[:, [0, 2]]) - else: - assert (values == -1).all() - assert (values[:, :5] == -1).all() - for state in target.get_bounded_attention_states(): - src_pool = getattr(source, state.state_name).big_pages_pool - dst_pool = state.big_pages_pool - if copy_pages: - for src_page, dst_page in ((2, 0), (0, 1)): - assert torch.equal(dst_pool.conv_state[dst_page], src_pool.conv_state[src_page]) - assert torch.equal(dst_pool.ssm_state[dst_page], src_pool.ssm_state[src_page]) - else: - assert (dst_pool.conv_state == -1).all() and (dst_pool.ssm_state == -2).all() - assert (dst_pool.conv_state[2] == -1).all() and (dst_pool.ssm_state[2] == -2).all() - assert (state.small_pages_pool.conv_state == 31).all() - assert (state.small_pages_pool.ssm_state == 32).all() - source_state = getattr(source, state.state_name) - assert torch.equal( - state.conv_state[:, 2, :, :state.config.conv_width], - source_state.conv_state[:, 1, :, :source_state.config.conv_width], - ) - dst_slot = 2 * (state.mtp_step + 1) - src_slot = source_state.mtp_step + 1 - assert torch.equal(state.ssm_state[:, dst_slot], source_state.ssm_state[:, src_slot]) - assert (state.conv_state[:, :2] == 41).all() - assert (state.conv_state[:, 2, :, state.config.conv_width:] == 41).all() - assert (state.ssm_state[:, :dst_slot] == 42).all() - assert (state.ssm_state[:, dst_slot + 1:] == 42).all() - assert state.mtp_state_index[2] == 0 - assert (state.mtp_state_index[:2] == 1).all() - for state in target.get_attention_tail_states(): - assert not state.tail_buffer[:, 2].any() - assert (state.tail_buffer[:, :2] == 51).all() + + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("readies", [[True, False, False], [False, True, False], [False, False, True], [True, True, True]]) def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bool]): - source = composition("cuda", with_tail=False) + source = composition("cuda") big_slots = init_source(source) small_slots = source.bounded_att_small_page_allocator.alloc(1) source.save_bounded_att_runtime_to_pages( diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index a74b71fcc3..04c82bfdd6 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -156,7 +156,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( torch.cuda.synchronize() for rank, (resources, full) in enumerate(targets): if glm: - resources.main_kpool_c4_bf16_state.tail_buffer.fill_(13) + resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(13) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] resources.read_req_from_cpu( cpu_page_buffers=buffer, @@ -182,8 +182,9 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( dst_fields[name].tensor[:, 8:16].contiguous().view(torch.uint8), ) if glm: - assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, 1].any() - assert (resources.main_kpool_c4_bf16_state.tail_buffer[:, [0, 2]] == 13).all() + # 无重叠窗口在组边界不依赖旧历史，CPU 恢复无需修改其数据。 + assert (resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1] == 13).all() + assert (resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, [0, 2]] == 13).all() assert not buffer[0].any() diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index d2a63e1ec6..51f09bde85 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -143,7 +143,7 @@ def _snapshot( state = main_linear_state(resources) result["conv"] = state.conv_state[:, req_idx, ..., cursor : cursor + state.config.conv_width].float().cpu().numpy() result["ssm"] = state.ssm_state[:, req_idx * 3 + cursor].cpu().numpy() - tail = resources.main_kpool_c4_bf16_state.tail_buffer[:, req_idx] + tail = resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, req_idx] result["tail"] = torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() return result @@ -199,7 +199,7 @@ def _worker(connection: Connection, backend: str, sender: bool): field.tensor.copy_(values) state = main_linear_state(resources) for i, tensor in enumerate( - (state.conv_state, state.ssm_state, resources.main_kpool_c4_bf16_state.tail_buffer) + (state.conv_state, state.ssm_state, resources.main_compressor_indexer_w4_s4_bf16_state.buffer) ): tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) expected = _snapshot( @@ -286,7 +286,7 @@ def _worker(connection: Connection, backend: str, sender: bool): assert main_linear_state(resources).mtp_state_index[2] == 0 assert not main_linear_state(resources).conv_state[:, :2].any() assert not main_linear_state(resources).ssm_state[:, :6].any() - assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, :2].any() + assert not resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, :2].any() assert not read_kv_storage(full)[:, [0, 1, 2, 3]].any() assert (pool.conv_state == 41).all() assert (pool.ssm_state == 43).all() diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 013555fd3e..186565bb69 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -17,7 +17,7 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("glm", [False, True]) -def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( +def test_main_and_draft_transfers_copy_pd_and_dp_windows( glm: bool, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, ): pd_args.pd_kv_page_size = 128 @@ -41,8 +41,8 @@ def test_main_and_draft_transfers_copy_pd_tails_and_reset_dp_tails( if "token" in field.axes: field.tensor.fill_(11 + index) if glm: - getattr(source, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 0].fill_(7 + index) - getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 2].fill_(19) + getattr(source, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 0].fill_(7 + index) + getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 2].fill_(19) main_linear_state(source).conv_state.fill_(3) main_linear_state(source).ssm_state.fill_(5) page = torch.empty(source.get_pd_page_bytes(), dtype=torch.uint8, device="cuda") @@ -87,12 +87,12 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: if "token" in field.axes: assert (field.tensor[:, 8:12] == 11 + index).all() if glm: - tail = getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer + tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer expected_tail = torch.zeros_like(tail[:, 1]) for position in range(4, 8): expected_tail[:, position % tail.shape[2]].fill_(7 + index) assert torch.equal(tail[:, 1], expected_tail) # PD restores only the prefill logical window. - assert not tail[:, 2].any() # DP fetch resets the aligned prefix. + assert (tail[:, 2] == 7 + index).all() # DP 直接复制相同 MTP 容量的 ring。 slot = source.bounded_att_big_page_allocator.alloc(1)[0] source.save_bounded_att_runtime_to_pages( @@ -132,5 +132,8 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: if "token" in field.axes: assert (field.tensor[:, 8:16] == 11 + index).all() if glm: - assert not getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 1].any() - assert not getattr(target, f"{role}_kpool_c4_bf16_state").tail_buffer[:, 2].any() + # CPU 组边界恢复不修改无重叠窗口，保留先前 PD/DP 的数据。 + tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer + for position in range(4, 8): + assert (tail[:, 1, position % tail.shape[2]] == 7 + index).all() + assert (tail[:, 2] == 7 + index).all() diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 9bbbb9fec0..1fd8f3e21e 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -1,5 +1,7 @@ from __future__ import annotations +from unit_tests.common.att_resources.model_helpers import window_states, linear_states + from unit_tests.common.att_resources.model_helpers import full_state_test_fields from unit_tests.common.att_resources.model_helpers import full_state_test_buffers @@ -202,8 +204,8 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( main_linear_state(source).conv_state[:, 1, ..., :2].fill_(13) main_linear_state(source).ssm_state[:, 3].fill_(17) if glm: - source.main_kpool_c4_bf16_state.tail_buffer[:, 1].fill_(23) - target.main_kpool_c4_bf16_state.tail_buffer[:, 2].fill_(29) + source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1].fill_(23) + target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 2].fill_(29) torch.cuda.synchronize(source_device) # 同时验证运行态与指定大页复制，保留其他页及小页快照。 for pool in (main_linear_state(target).big_pages_pool, main_linear_state(target).small_pages_pool): @@ -276,9 +278,12 @@ def barrier(group: str): assert (main_linear_state(target).small_pages_pool.conv_state == 41).all() assert (main_linear_state(target).small_pages_pool.ssm_state == 43).all() if glm: - assert not target.main_kpool_c4_bf16_state.tail_buffer[:, 2].any() - assert (source.main_kpool_c4_bf16_state.tail_buffer[:, 1] == 23).all() - assert not target.main_kpool_c4_bf16_state.tail_buffer[:, :2].any() + assert torch.equal( + target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 2].cpu(), + source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1].cpu(), + ) + assert (source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1] == 23).all() + assert not target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, :2].any() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -319,8 +324,8 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( for storage_buffer in full_state_test_buffers(full).values(): storage_buffer[:, 4].zero_() main_linear_state(resources).conv_state.fill_(11) - resources.main_kpool_c4_bf16_state.tail_buffer.fill_(7) - for state in resources.get_bounded_attention_states(): + resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(7) + for state in linear_states(resources): for pool in (state.big_pages_pool, state.small_pages_pool): if pool is not None: pool.conv_state.zero_() @@ -353,7 +358,9 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( error = parent.recv() assert error is None, error assert (read_kv_storage(full)[:, 4] == 9).all() - assert not resources.main_kpool_c4_bf16_state.tail_buffer[:, 1].any() + window = resources.main_compressor_indexer_w4_s4_bf16_state + assert not window.buffer[:, 1, :, :window.config.candidate_dim].any() + assert torch.isneginf(window.buffer[:, 1, :, window.config.candidate_dim:]).all() assert (main_linear_state(resources).conv_state == 11).all() assert (main_linear_state(resources).big_pages_pool.conv_state[slot[0]] == 11).all() assert not main_linear_state(resources).big_pages_pool.ssm_state[slot[0]].any() @@ -512,12 +519,12 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: assert resources.bounded_att_big_page_allocator is None assert resources.bounded_att_small_page_allocator is None if kind == "dp": - for state in resources.get_bounded_attention_states(): + for state in linear_states(resources): assert state.big_pages_pool.conv_state.is_shared() assert not state.big_pages_pool.conv_state.is_pinned() resources.register_page_pools_pin_memory() resources.register_page_pools_pin_memory() - for state in resources.get_bounded_attention_states(): + for state in linear_states(resources): for pool in (state.big_pages_pool, state.small_pages_pool): if kind == "dp": assert pool is not None @@ -532,7 +539,8 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: assert not hasattr(resources, "req_to_token_indexs") assert mapping[1, 0].item() == 7 resources.copy_full_att_tokens(torch.tensor([0], dtype=torch.long, device="cuda"), torch.tensor([4], dtype=torch.long, device="cuda")) - resources.reset_attention_tail_state([1]) + for window in window_states(resources): + window.reset_bounded_att_state([1]) torch.cuda.synchronize() # 一类通道耗尽不会影响另一类通道的读取计数。 with pytest.raises(RuntimeError, match="exhausted"): diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index f1eb62ec87..b139b092ca 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -17,7 +17,7 @@ from lightllm.common.attention_state.bounded_att import ( BoundedAttentionState, ConvSSMStateConfig, create_gdn_state, create_kda_state, ) -from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state +from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state @pytest.mark.parametrize( @@ -156,21 +156,21 @@ def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Calla def test_kpool_creation_preserves_pool_and_mtp_dimensions(): - config = KPoolStateConfig( + config = WindowCompressorStateConfig( logical_layer_ids=(3, 7), - pool_size=4, - head_dim=128, + window_size=4, compress_stride=4, target="indexer", + candidate_dim=128, dtype=torch.bfloat16, ) - state = create_kpool_state( + state = create_window_compressor_state( config=config, request_capacity=3, mtp_step=2, device="cpu", ) assert state.config is config and state.state_name is None - assert state.tail_buffer.shape == (2, 3, 6, 256) - assert state.get_state_type() == "kpool_c4_bf16" + assert state.buffer.shape == (2, 3, 6, 256) + assert state.get_state_type() == "compressor_indexer_w4_s4_bf16" def test_create_rejects_encoding_from_another_logical_family(): diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index bc03b18597..787d3ef04f 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -27,6 +27,34 @@ KINDS = ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] +@pytest.mark.parametrize("kind", [*KINDS, "compressed", "kda", "gdn", "swa", "compressor"]) +def test_pd_region_cache_returns_independent_snapshots(kind: str, pd_args: EasyDict) -> None: + from lightllm.common.attention_state.bounded_att import GDNState, KDAState + from unit_tests.common.att_resources.test_bounded_state import make_state + from unit_tests.common.att_resources.test_dsv4_state import make_atomic_state, make_full + + pd_args.pd_kv_page_size = 256 + if kind in KINDS: + state = definition(kind) + elif kind == "compressed": + state = make_full(ratio=4) + elif kind in ("kda", "gdn"): + state = make_state(cls=KDAState if kind == "kda" else GDNState, device="meta") + else: + state = make_atomic_state("swa" if kind == "swa" else "kv_w8") + + getter = state.get_pd_transfer_regions + before = getter.cache_info() + first, second = getter(), getter() + after = getter.cache_info() + assert after.misses == before.misses + 1 and after.hits == before.hits + 1 + assert first == second and first is not second + for first_region, second_region in zip(first, second): + assert first_region is not second_region + for first_mapping, second_mapping in zip(first_region.axis_mappings, second_region.axis_mappings): + assert first_mapping is not second_mapping + + @pytest.mark.parametrize("kind", KINDS) def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_args: EasyDict): from inspect import signature @@ -91,7 +119,7 @@ def get_state_type(self) -> str: @pytest.mark.parametrize("model_name", ["main", "draft1"]) def test_register_state_preserves_object_and_storage(kind: str, model_name: str): from unit_tests.common.att_resources.test_bounded_state import make_state - from unit_tests.common.att_resources.test_tail_state import make_tail + from unit_tests.common.att_resources.test_window_compressor_state import make_window if kind == "full_meta": state = definition("gqa") @@ -104,8 +132,8 @@ def test_register_state_preserves_object_and_storage(kind: str, model_name: str) storage = state.conv_state pool = state.big_pages_pool else: - state = make_tail() - storage = state.tail_buffer + state = make_window() + storage = state.buffer config = state.config manager = AttentionStateManager(request_capacity=3) manager.register_state(model_name, state) @@ -118,7 +146,7 @@ def test_register_state_preserves_object_and_storage(kind: str, model_name: str) elif kind == "bounded": assert state.conv_state is storage and state.big_pages_pool is pool else: - assert state.tail_buffer is storage + assert state.buffer is storage other_role = "draft0" if model_name == "main" else "main" with pytest.raises(AssertionError, match="multiple model roles"): manager.register_state(other_role, state) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 28c04506ff..33c1098a94 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -32,7 +32,7 @@ def test_created_state_queries_keep_the_same_objects_through_storage_initializat if isinstance(state, FullAttentionState) and state.state_name.startswith(role + "_") )) for query in ( - manager.get_full_attention_states, manager.get_bounded_attention_states, manager.get_attention_tail_states, + manager.get_full_attention_states, manager.get_bounded_attention_states, ): with pytest.raises(AssertionError, match="State registration"): query() @@ -68,15 +68,14 @@ def check_created_before_allocation(manager: AttentionStateManager, *args: objec def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): manager = MainDraftTestResourceManager() - queries = (manager.get_full_attention_states, manager.get_bounded_attention_states, manager.get_attention_tail_states) + queries = (manager.get_full_attention_states, manager.get_bounded_attention_states) results = {(query.__name__, role): query(role) for query in queries for role in (None, "main", "draft0")} - assert manager._get_states.cache_info().currsize == 9 + assert manager._get_states.cache_info().currsize == 6 assert manager.get_full_attention_states() == ( manager.draft0_mla_fp32_state, manager.main_gqa_fp32_state, ) assert manager.get_bounded_attention_states("main") == (manager.main_kda_conv_fp32_ssm_fp32_state,) assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_conv_fp32_ssm_fp32_state,) - assert manager.get_attention_tail_states("draft0") == () hits = manager._get_states.cache_info().hits for query in queries: @@ -84,7 +83,7 @@ def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): assert query(role) is results[(query.__name__, role)] with pytest.raises(AssertionError, match="unknown model role"): query("target") - assert manager._get_states.cache_info().hits == hits + 9 + assert manager._get_states.cache_info().hits == hits + 6 manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.fill_(7) assert (manager.get_bounded_attention_states("main")[0].conv_state == 7).all() @@ -123,7 +122,7 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): assert exported.get_bounded_attention_states() == tuple( getattr(exported, state.state_name) for state in source_states ) - for state in (*source.get_full_attention_states(), *source.get_attention_tail_states()): + for state in source.get_full_attention_states(): assert getattr(exported, state.state_name) is state @@ -156,7 +155,6 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): source = MainDraftTestResourceManager() source.get_full_attention_states() source.get_bounded_attention_states() - source.get_attention_tail_states() imported = pickle.loads(pickle.dumps(source.create_pd_ipc_copy())) assert imported._state_registration_finished assert imported.full_att_allocator is None @@ -164,19 +162,18 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): assert imported.bounded_att_small_page_allocator is None misses = source._get_states.cache_info().misses for query in ( - imported.get_full_attention_states, imported.get_bounded_attention_states, imported.get_attention_tail_states, + imported.get_full_attention_states, imported.get_bounded_attention_states, ): for state in query(): name = state.state_name assert state is getattr(imported, name) and state is not getattr(source, name) - assert imported._get_states.cache_info().misses == misses + 3 + assert imported._get_states.cache_info().misses == misses + 2 def test_clearing_query_cache_releases_manager_and_states(): manager = MainDraftTestResourceManager() manager.get_full_attention_states() manager.get_bounded_attention_states() - manager.get_attention_tail_states() manager_ref = weakref.ref(manager) state_ref = weakref.ref(manager.main_kda_conv_fp32_ssm_fp32_state) del manager diff --git a/unit_tests/common/att_resources/test_tail_state.py b/unit_tests/common/att_resources/test_tail_state.py deleted file mode 100644 index 04d2d92ba3..0000000000 --- a/unit_tests/common/att_resources/test_tail_state.py +++ /dev/null @@ -1,124 +0,0 @@ -from __future__ import annotations - -import pytest -import torch - -from lightllm.common.attention_state.att_tail import KPoolAttentionTailState, KPoolStateConfig -from lightllm.common.attention_state.types import ( - AxisMapping, StateTransferLayout, TransferRegion, -) - - -def make_tail( device: str | torch.device | int = "cpu", mtp_step: int = 2): - return KPoolAttentionTailState( - config=KPoolStateConfig( - logical_layer_ids=(3, 7), - pool_size=4, - head_dim=2, - dtype=torch.float32, - ), - request_capacity=3, - mtp_step=mtp_step, - device=device, - ) - - -def tail_layout(state: KPoolAttentionTailState): - shape = (2, state.config.pool_size, 4) - nbytes = shape[0] * shape[1] * shape[2] * 4 - return StateTransferLayout( - state_name=state.state_name, - tokens_per_page=None, - byte_start=0, - byte_end=nbytes + 16, - regions=( - TransferRegion( - field="tail", - byte_start=16, - byte_end=16 + nbytes, - dtype=torch.float32, - axes=("layer", "window", "channel"), - shape=shape, - axis_mappings=( - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=shape[1], - local_circular=True, - ), - AxisMapping( - axis="channel", - local_start=0, - external_start=0, - length=4, - ), - ), - ), - ), - ) - - -def test_tail_reset_clears_selected_requests_only(): - state = make_tail() - state.tail_buffer.fill_(17) - state.reset_attention_tail_state([0, 1]) - assert not state.tail_buffer[:, :2].any() - assert (state.tail_buffer[:, 2] == 17).all() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("device", ["cuda"]) -@pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) -@pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) -def test_pd_logical_window_restores_different_mtp_rings( - device: str, - source_step: int, - target_step: int, - length: int, -): - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA unavailable") - source = make_tail(device=device, mtp_step=source_step) - target = make_tail(device=device, mtp_step=target_step) - source.tail_buffer.copy_(torch.arange(source.tail_buffer.numel(), device=device).reshape_as(source.tail_buffer)) - original = source.tail_buffer.clone() - target.tail_buffer.fill_(37) - layout = tail_layout(source) - assert layout == tail_layout(target) - assert source.get_pd_transfer_regions() == target.get_pd_transfer_regions() - pages = torch.full( - size=(3, layout.byte_end + 16), - fill_value=253, - dtype=torch.uint8, - device=device, - ) - source.write_to_pd_page( - req_idx=1, - pd_page_buffers=pages, - pd_page_index=1, - state_token_len=length, - layout=layout, - ) - target.read_from_pd_page( - pd_page_buffers=pages, - req_idx=2, - pd_page_index=1, - state_token_len=length, - layout=layout, - ) - expected = torch.full_like(target.tail_buffer[:, 2], 37) - for position in range(length - 4, length): - slot = position % (4 + target_step) - if position < 0: - expected[:, slot].zero_() - else: - expected[:, slot] = source.tail_buffer[:, 1, position % (4 + source_step)] - assert torch.equal(target.tail_buffer[:, 2], expected) - assert (target.tail_buffer[:, :2] == 37).all() - assert torch.equal(source.tail_buffer, original) - assert (pages[0] == 253).all() and (pages[2] == 253).all() - assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end:] == 253).all() - if length < 4: - external = pages[1, 16:layout.byte_end].view(torch.float32).view(2, 4, 4) - assert not external[:, :4 - length].any() diff --git a/unit_tests/common/att_resources/test_window_compressor_state.py b/unit_tests/common/att_resources/test_window_compressor_state.py new file mode 100644 index 0000000000..8aa3c6cf06 --- /dev/null +++ b/unit_tests/common/att_resources/test_window_compressor_state.py @@ -0,0 +1,192 @@ +"""统一窗口状态：K-pool 的组边界及不同 MTP 容量间的 PD 恢复。""" + +from dataclasses import replace +from functools import partial +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_window( + device: str = "cpu", + mtp_step: int = 2, + dtype: torch.dtype = torch.float32, +) -> WindowCompressorState: + return WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(3, 7), + window_size=4, + compress_stride=4, + candidate_dim=2, + target="indexer", + dtype=dtype, + ), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + + +def window_layout(state: WindowCompressorState) -> StateTransferLayout: + regions = tuple( + replace(region, byte_start=region.byte_start + 16, byte_end=region.byte_end + 16) + for region in state.get_pd_transfer_regions() + ) + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None, + byte_start=16, + byte_end=regions[-1].byte_end, + regions=regions, + ) + + +def test_nonoverlapping_window_needs_no_cpu_snapshot() -> None: + state = make_window() + assert state.buffer.shape == (2, 3, 6, 4) + state.buffer.fill_(17) + state.reset_bounded_att_state([0, 1]) + assert not state.buffer[:, :2, :, :2].any() + assert torch.isneginf(state.buffer[:, :2, :, 2:]).all() + assert (state.buffer[:, 2] == 17).all() + state.initialize_page_pools(big_page_capacity=5, small_page_capacity=2) + assert state.big_pages_pool is state.small_pages_pool is None + assert state.get_cpu_transfer_regions() == () + # 组边界不依赖此前窗口；下一组四个输入均由新数据覆盖，无需读取快照。 + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[0], + ) + with pytest.raises(AssertionError, match="stride"): + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[7], + page_kind="big", + page_slot_ids=[0], + ) + + +@pytest.mark.parametrize("window_size,compress_stride", [(4, 4), (8, 4), (128, 128)]) +@pytest.mark.parametrize("page_kind", ["small", "cpu"]) +def test_window_cache_boundaries_must_close_compressor_groups( + window_size: int, + compress_stride: int, + page_kind: str, + cpu_args: StartArgs, +) -> None: + cpu_args.linear_att_hash_page_size = compress_stride + cpu_args.cpu_cache_token_page_size = 2 * compress_stride + if page_kind == "small": + cpu_args.linear_att_hash_page_size += 1 + else: + cpu_args.cpu_cache_token_page_size += 1 + state = WindowCompressorState( + config=replace(make_window().config, window_size=window_size, compress_stride=compress_stride), + request_capacity=3, + device="meta", + ) + # 即使不占 CPU 页，也不能在未完成的压缩组内恢复，否则缺少该组历史。 + with pytest.raises(AssertionError, match="pages must align to the compressor stride"): + state.get_cpu_transfer_regions() + + +@pytest.mark.parametrize("window_size,compress_stride", [(4, 4), (8, 4), (128, 128)]) +def test_manager_allocates_page_tensors_only_for_snapshot_states( + window_size: int, + compress_stride: int, + cpu_args: StartArgs, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.linear_att_hash_page_size = 128 + cpu_args.linear_att_page_block_num = 2 + cpu_args.cpu_cache_token_page_size = 256 + state = WindowCompressorState( + config=replace(make_window().config, window_size=window_size, compress_stride=compress_stride), + request_capacity=3, + device="cpu", + ) + initialize_pools = Mock(wraps=state.initialize_page_pools) + monkeypatch.setattr(state, "initialize_page_pools", initialize_pools) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", + partial(FullAttKVAllocator, pin_memory=False), + ) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((3,),), + global_head_num=1, + head_dim=4, + dtype=torch.bfloat16, + ), + device="meta", + )) + manager.register_state("main", state) + manager.finish_state_registration() + manager.allocate_storage( + token_capacity=512, + page_size=4, + big_page_tokens=256, + small_page_capacity=3, + device="cpu", + ) + if window_size == compress_stride: + initialize_pools.assert_not_called() + assert state.big_pages_pool is state.small_pages_pool is None + assert state.state_name not in manager.get_cpu_page_layouts() + else: + initialize_pools.assert_called_once() + assert state.big_pages_pool.buffer.shape == (4, 2, window_size - compress_stride, 4) + assert state.small_pages_pool.buffer.shape == (3, 2, window_size - compress_stride, 4) + assert state.state_name in manager.get_cpu_page_layouts() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) +@pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) +def test_pd_window_restores_different_mtp_rings( + dtype: torch.dtype, + source_step: int, + target_step: int, + length: int, +) -> None: + source = make_window(device="cuda", mtp_step=source_step, dtype=dtype) + target = make_window(device="cuda", mtp_step=target_step, dtype=dtype) + source.buffer.copy_(torch.arange(source.buffer.numel(), device="cuda").reshape_as(source.buffer)) + target.buffer.fill_(37) + layout = window_layout(source) + assert layout == window_layout(target) + pages = torch.full((3, layout.byte_end + 16), 253, dtype=torch.uint8, device="cuda") + source.write_to_pd_page( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=pages, + req_idx=2, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) + expected = torch.full_like(target.buffer[:, 2], 37) + for position in range(length - 4, length): + slot = position % (4 + target_step) + expected[:, slot] = 0 if position < 0 else source.buffer[:, 1, position % (4 + source_step)] + torch.testing.assert_close(target.buffer[:, 2], expected, rtol=0, atol=0) + assert (target.buffer[:, :2] == 37).all() + assert (pages[0] == 253).all() and (pages[2] == 253).all() + assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end:] == 253).all() diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 301e0ae90a..5d42ff63a5 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -111,7 +111,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) full = main_full linear = resources.main_kda_conv_bf16_ssm_fp32_state - tail = resources.main_kpool_c4_bf16_state + tail = resources.main_compressor_indexer_w4_s4_bf16_state assert not hasattr(req, "req_to_conv_state") att_kv = full.get_layer_view(3).kv assert att_kv.shape == (17, 1, 512) @@ -130,13 +130,12 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( assert torch.equal(index_bytes, expected_index_bytes) assert tail.get_layer_view(3).shape == (4, 4 + mtp_step, 256) if mtp_step: - assert resources.draft0_kpool_c4_bf16_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) + assert resources.draft0_compressor_indexer_w4_s4_bf16_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) assert draft_full.get_layer_view(4).kv.shape == att_kv.shape kind = "small" if small_page else "big" allocator = resources.bounded_att_small_page_allocator if small_page else resources.bounded_att_big_page_allocator slot = allocator.alloc(1)[0] resources.reset_bounded_att_state([0]) - resources.reset_attention_tail_state([0]) linear.conv_state[:, 0].normal_() linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() @@ -150,31 +149,25 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( torch.cuda.synchronize() linear.conv_state[:, 0].zero_() linear.ssm_state[:, 0].zero_() - tail.tail_buffer.normal_() - unchanged = tail.tail_buffer[:, [0, 1, 3]].clone() + tail.buffer.normal_() + unchanged = tail.buffer.clone() resources.restore_bounded_att_runtime_from_pages( req_indices=[2], state_token_lens=[8], page_kind=kind, page_slot_ids=[slot], ) - resources.reset_attention_tail_state([2]) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) - # Both page sizes are aligned to complete pools. Restoring a prefix must - # discard stale tail values from a previously allocated request slot. - assert not tail.tail_buffer[:, 2].any() - assert torch.equal(tail.tail_buffer[:, [0, 1, 3]], unchanged) - tail.tail_buffer[:, 2].normal_() - preserved_tail = tail.tail_buffer.clone() + # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 + assert torch.equal(tail.buffer, unchanged) resources.reset_bounded_att_state([2]) assert not linear.conv_state[:, 2].any() assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() - assert torch.equal(tail.tail_buffer, preserved_tail) - resources.reset_attention_tail_state([2]) - assert not tail.tail_buffer[:, 2].any() - assert torch.equal(tail.tail_buffer[:, [0, 1, 3]], unchanged) + assert not tail.buffer[:, 2, :, :tail.config.candidate_dim].any() + assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim:]).all() + assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only # belong to live requests. Bytewise equality checks FP8 scale integrity. packed_bytes = full.storage.view(torch.uint8) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index d1a0eaf29f..b46a87cde6 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -592,7 +592,7 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): att_state_manager=SimpleNamespace( get_state=lambda state_type_filter, **kwargs: ( SimpleNamespace(get_layer_view=lambda logical_layer_id: tail) - if "kpool" in state_type_filter else full_att_state + if "compressor_indexer" in state_type_filter else full_att_state ), ), is_draft_model=False, @@ -624,11 +624,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): assert full[-1, :2049].unique().numel() == 2049 assert (full[-1, :2049] < tokens).all() # Restore an aligned prefix using only token KV, then decode across the - # next pool boundary. The runtime tail starts empty, as on a page restore. + # next pool boundary. Stale window data is irrelevant at a complete group boundary. prefix = 2060 storage = storage.clone() storage[ragged[prefix:].long()] = 0 - tail.zero_() + tail.fill_(19) infer.is_prefill = False infer.b1_cu_q_seq_len[1] = 1 infer.max_q_seq_len = 1 diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 657289a1f5..1440abd623 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -424,7 +424,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) main.prepare_attention_states() assert main.att_state_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.is_cuda - assert main.att_state_manager.main_kpool_c4_bf16_state.tail_buffer.is_cuda + assert main.att_state_manager.main_compressor_indexer_w4_s4_bf16_state.buffer.is_cuda for step, draft in enumerate(drafts): draft.prepare_attention_states() assert draft.att_state_manager is main.att_state_manager @@ -436,7 +436,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): assert state.config.logical_layer_groups == ((4 + step,),) assert state.state_name == f"draft{step}_mla_packed_bf16_state" assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None - assert getattr(draft.att_state_manager, f"draft{step}_kpool_c4_bf16_state").tail_buffer.is_cuda + assert getattr(draft.att_state_manager, f"draft{step}_compressor_indexer_w4_s4_bf16_state").buffer.is_cuda # NoPE state initialization must not require a model-owned RoPE cache. for model in (main, *drafts): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index bc971862a4..218095216a 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,4 +1,4 @@ -from unit_tests.common.att_resources.model_helpers import build_hybrid_resources +from unit_tests.common.att_resources.model_helpers import window_states, build_hybrid_resources from unit_tests.common.att_resources.model_helpers import full_state_test_fields import dataclasses @@ -96,8 +96,8 @@ def create(tp, target_mtp_step=None): full=main, full_states=tuple(state for state in (main, draft) if state is not None), linear=resources.main_kda_conv_bf16_ssm_fp32_state, - tail=resources.main_kpool_c4_bf16_state, - tails=resources.get_attention_tail_states(), + tail=resources.main_compressor_indexer_w4_s4_bf16_state, + tails=window_states(resources), resources=resources, ) ) @@ -160,7 +160,7 @@ def test_pd_kv_and_runtime_state_roundtrip( offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) - state.tail_buffer[:, src_req].copy_(tail[offset:end]) + state.buffer[:, src_req].copy_(tail[offset:end]) offset = end untouched = [] @@ -171,15 +171,15 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() for state in mem.tails: - state.tail_buffer.normal_() - original_tails.append(torch.cat([state.tail_buffer[:, dst_req] for state in mem.tails]).clone()) + state.buffer.normal_() + original_tails.append(torch.cat([state.buffer[:, dst_req] for state in mem.tails]).clone()) untouched.append( [ x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.linear.conv_state, 1), (mem.linear.ssm_state, dst_mtp_size), - *((state.tail_buffer, 1) for state in mem.tails), + *((state.buffer, 1) for state in mem.tails), ) ] ) @@ -230,7 +230,7 @@ def test_pd_kv_and_runtime_state_roundtrip( ) assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) - restored_tail = torch.cat([state.tail_buffer[:, dst_req] for state in mem.tails]) + restored_tail = torch.cat([state.buffer[:, dst_req] for state in mem.tails]) expected_tail = original_tails[rank].clone() for position in range(length - 4, length): expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] @@ -241,7 +241,7 @@ def test_pd_kv_and_runtime_state_roundtrip( ( (mem.linear.conv_state, 1), (mem.linear.ssm_state, dst_mtp_size), - *((state.tail_buffer, 1) for state in mem.tails), + *((state.buffer, 1) for state in mem.tails), ), untouched[rank], ): @@ -264,7 +264,7 @@ def test_pd_kv_and_runtime_state_roundtrip( all_indexes = indexes + new_indexes compress_pools( raw=raw, - tail=mem.tail.tail_buffer[0], + tail=mem.tail.buffer[0], packed_buffer=mem.full.get_layer_view(3).indexer_k, ape=ape, lengths=tensor(range(length + 1, length + count + 1)), @@ -295,23 +295,25 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, @pytest.mark.parametrize("prefix", [1, 3, 4, 5, 11, 12]) def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, prefix): """真实 prefill 导出后，在不同大小的 ring 上验证、拒绝回退并再次验证。""" - from lightllm.common.attention_state.att_tail import KPoolStateConfig, create_kpool_state + from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state from lightllm.common.attention_state.types import StateTransferLayout torch.manual_seed(53) - config = KPoolStateConfig( + config = WindowCompressorStateConfig( logical_layer_ids=(0,), - pool_size=4, - head_dim=128, + window_size=4, + compress_stride=4, + candidate_dim=128, + target="indexer", dtype=torch.bfloat16, ) - source = create_kpool_state( + source = create_window_compressor_state( config=config, request_capacity=2, mtp_step=source_step, device="cuda", ) - target = create_kpool_state( + target = create_window_compressor_state( config=config, request_capacity=2, mtp_step=target_step, @@ -373,7 +375,7 @@ def prefill(values, destination, tail): state_token_len=prefix, layout=layout, ) - target.tail_buffer.fill_(17) + target.buffer.fill_(17) target.read_from_pd_page( pd_page_buffers=pages, req_idx=0, @@ -381,7 +383,7 @@ def prefill(values, destination, tail): state_token_len=prefix, layout=layout, ) - assert (target.tail_buffer[:, 1] == 17).all() + assert (target.buffer[:, 1] == 17).all() width = target_step + 1 # 第一次验证后回访已接受位置；新增候选历史必须由 D 的本地 ring 正确保留。 for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index d75d10e776..737e5e2165 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -34,7 +34,7 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): @pytest.mark.usefixtures("kv_allocator_environment") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") -def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): +def test_bounded_cache_miss_resets_runtime_and_preserves_pages(monkeypatch): from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from unit_tests.common.att_resources.test_manager_helpers import composition, init_source @@ -44,9 +44,7 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): ) resources = composition(device="cuda") slots = init_source(resources) - tail = resources.main_kpool_c4_fp32_state resources.reset_bounded_att_state([1]) - assert (tail.tail_buffer[:, 1] == 9).all() resources.restore_bounded_att_runtime_from_pages( req_indices=[1], state_token_lens=[8], @@ -83,8 +81,6 @@ def test_bounded_reset_preserves_tail_and_cache_miss_resets_all(monkeypatch): expected_conv, expected_ssm = (23, 47) if state is resources.main_kda_conv_fp32_ssm_fp32_state else (24, 48) assert (pool.conv_state[slots[-1]] == expected_conv).all() assert (pool.ssm_state[slots[-1]] == expected_ssm).all() - assert not tail.tail_buffer[:, 1].any() - assert (tail.tail_buffer[:, 0] == 3).all() assert resources.bounded_att_big_page_allocator.free_count == 1 assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 From 43ffe867ed35426d4f3abc85858ce29abf78d689 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 01:41:14 +0000 Subject: [PATCH 31/56] fix: address attention-state regressions and synchronize NCCL receives --- docs/CN/source/models/add_new_model.md | 5 +-- docs/EN/source/models/add_new_model.md | 5 +-- docs/cache_manager_design.md | 6 +++ lightllm/common/attention_state/types.py | 2 + .../attention/nsa/dsv4_fp8_flashmla_sparse.py | 6 ++- .../layer_infer/pre_layer_infer.py | 3 +- lightllm/models/glm5_next/model.py | 6 +++ .../server/router/model_infer/infer_batch.py | 6 ++- .../dp_backend/dp_shared_kv_trans.py | 12 ++--- .../mode_backend/dp_backend/impl.py | 4 +- .../mode_backend/multi_level_kv_cache.py | 3 +- .../decode_node_impl/decode_trans_process.py | 8 ++++ .../mode_backend/pd/nccl_kv_transporter.py | 5 +++ .../prefill_trans_process.py | 6 +-- lightllm/server/tokenizer.py | 22 +++++---- test/att_resources/validate_model_cache.py | 18 ++++++-- .../att_resources/test_network_transport.py | 2 + .../basemodel/attention/test_dsv4_flashmla.py | 37 ++++++++++++++- .../common/test_deepseek4_paged_cache.py | 4 +- unit_tests/models/deepseek_v4/test_model.py | 23 ++++++++++ .../mode_backend/test_nccl_control_channel.py | 33 ++++++++++++++ .../model_infer/test_multimodal_chunking.py | 5 ++- unit_tests/server/test_tokenizer.py | 45 ++++++++++++++++--- 23 files changed, 215 insertions(+), 51 deletions(-) diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 60f276ac94..7edb0c9701 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -109,7 +109,6 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_req_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -128,9 +127,9 @@ class TpPartBaseModel: | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | -全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。仅保留资源准备和运行时初始化两个单模型公共入口，均不接收 draft 列表。 +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。资源准备和运行时初始化分别通过这两个单模型公共入口完成，均不接收 draft 列表。 -公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。三个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded/tail 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 +公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 ## 2. 添加 bloom 模型的示例说明 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index 6d814a8fe7..cd71541598 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -109,7 +109,6 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_req_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -128,9 +127,9 @@ Commonly used interfaces that need to be inherited and implemented | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | -After all main/draft models are constructed, `ModeBackend.init_model` calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization are the only two public startup entry points; both operate on one model and take no draft-model list. +After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()` returned by the models, binds it to every model through `bind_req_manager()`, then calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization each use a single-model entry point; neither takes a draft-model list. -The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The three State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded/tail also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. +The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The two State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded factories also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. ## 2. Example of Adding Bloom Model diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index b6eef47841..5efe8e3df1 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1027,6 +1027,8 @@ PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效 CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 +NCCL 接收线程在 recv 后记录本地 Xfer 的 event，Decode 在读取 PD 页前调用 event.synchronize()，不增加反向确认传输。该同步要求 record 已经提交；未记录的 event 会直接返回，P 的 done 通知本身不能保证 D 的接收线程已执行 record。当前最小实现不额外协调这两个主机线程的提交顺序。 + #### 调用方与资源实现的边界 radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存和恢复状态页时调用上述通用接口，前缀与槽位的对应关系由 radix 调用流程保证；槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 @@ -1153,6 +1155,8 @@ class StateTransferLayout: token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_mappings`。其余搬运轴通过 `AxisMapping` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。普通非连续范围使用多个连续片段表达；环形轴使用 `local_circular` / `external_circular`，按两侧 Tensor 的实际轴容量对 `start + i` 取模。标记始终按本地 State → 外部页解释，读回时反向执行；搬运长度不超过各环形轴容量。 +环形标记只描述寻址约定，必须由具体 State 的专用 kernel 实现；普通 full/bounded 的 TP 切片路径不支持将这两个标记设为 True。当前窗口压缩器通过 `state_token_len` 和两侧窗口容量直接完成环形搬运，不读取 `AxisMapping` 的环形标记。 + `shape` 表达该外部字段区域的完整范围，而 `axis_mappings` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： ```python @@ -1838,6 +1842,8 @@ DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并 V4 沿用共享 `AttentionStateManager`，不引入专用 ReqManager 或 MemoryManager。各模型通过 `create_full_att_states` / `create_bounded_att_states` 创建自己的定义，注册完成后统一分配容量。 main 与 draft0..draft35 使用相同 token 分配域，但分别持有存储；SWA prefill workspace 也按模型角色区分。 +P/D 的 decode 节点仍需在初始化时对 HOLD 请求执行一次 prefill，因此保留每层一个 SWA HOLD 页； +其普通 prefill workspace 容量为 0，不分配长输入临时区，实际请求的前缀由 prefill 节点传入。 模型组合配置保留在 `models/deepseek_v4/config.py`，不放入原子 State 目录。 组合配置的 `create_swa_config()` / `create_compressed_config()` 负责选择参与层及算法参数， diff --git a/lightllm/common/attention_state/types.py b/lightllm/common/attention_state/types.py index 9b578b2658..8e20c6d6e0 100644 --- a/lightllm/common/attention_state/types.py +++ b/lightllm/common/attention_state/types.py @@ -70,6 +70,8 @@ class AxisMapping: TransferRegion 中的方向固定为本地 State 到外部页，读取时反向应用。 环形轴按对应 Tensor 的实际轴长度取模，不要求两侧容量相同。 + 环形标记只是布局描述，须由具体 State 的专用 kernel 实现；普通 TP 切片路径不支持取模。 + 当前窗口压缩器直接按 state_token_len 搬运 ring，不依赖这两个标记。 示例：本地 head[0:2] 写到全局传输区域 head[2:4]，读回时反向复制。 >>> head_mapping = AxisMapping(axis="head", local_start=0, external_start=2, length=2) diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py index 448efcf979..1fd992aa37 100644 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py @@ -39,8 +39,6 @@ def __init__(self, model): @classmethod def get_swa_prefill_workspace(cls, model): """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" - if model.args.run_mode == "decode": - return None, 0 from lightllm.models.deepseek_v4.config import create_attention_config config = create_attention_config( @@ -57,6 +55,10 @@ def get_swa_prefill_workspace(cls, model): ) * 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 diff --git a/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py index 8a923e97ab..2e09c9d3b3 100644 --- a/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py @@ -23,6 +23,7 @@ def context_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, laye def token_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, layer_weight): input_embdings = super().token_forward(input_ids, infer_state, layer_weight) - if not infer_state.is_cuda_graph: + # 捕获时保留 graph 输入引用；普通 decode 及时释放未使用的 draft 特征。 + if not infer_state.is_cuda_graph_capture: infer_state.mtp_draft_input_hiddens = None return hc_expand(input_embdings, self.hc_mult) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 5d33a188eb..bd1b745bc0 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -112,6 +112,12 @@ def _verify_params(self): super()._verify_params() assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): + raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa": + # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 + if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: + raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") def autotune_layers(self): return 4 diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index ecd8d555e3..0c911b3799 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -636,7 +636,11 @@ def _init_all_state(self): # 仅载入调用方显式指定的 token ID 序列；字符串由 detokenization / HTTP 匹配。 self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_token_ids_list() self.multimodal_params = self.multimodal_params.to_dict() - self.mm_input_spans = get_mm_input_spans(self.multimodal_params) if self.args.disable_chunked_mm_input else [] + self.mm_input_spans = ( + get_mm_input_spans(self.shm_req.shm_prompt_ids.arr[: self.shm_req.input_len], self.multimodal_params) + if self.args.disable_chunked_mm_input + else [] + ) self.shared_kv_node: Union[TreeNode, HybridAttPagedTreeNode] = None self.finish_status = FinishStatus() diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 90fc008b86..4c693a1ac8 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -78,22 +78,22 @@ def build_shared_kv_trans_tasks( trans_tasks: List[TransTask] = [] available_lens = self.shared_req_infos.arr[0 : len(reqs), :, self._KV_LEN_INDEX] - rank_max_radix_cache_lens = np.max(available_lens, axis=1) + max_available_kv_lens = np.max(available_lens, axis=1) # 为本 DP 负责的请求选择可用前缀最长的来源。 - for req_index, req, max_req_radix_cache_len, req_dp_rank in zip( + for req_index, req, max_available_kv_len, req_dp_rank in zip( list(range(len(reqs))), reqs, - rank_max_radix_cache_lens, + max_available_kv_lens, req_dp_ranks, ): # 当前请求是本 dp_rank 负责的 - is_current_dp_handle = req_dp_rank == self.dp_rank_in_node + is_local_dp_request = req_dp_rank == self.dp_rank_in_node # 计算需要传输的 kv 长度， 不能超过 req.get_cur_total_len() - 1 - trans_size = min(max_req_radix_cache_len, req.get_cur_total_len() - 1) - req.cur_kv_len + trans_size = min(max_available_kv_len, req.get_cur_total_len() - 1) - req.cur_kv_len target_kv_len = req.cur_kv_len + trans_size alloc_token_num = req._kv_cache_alloc_need(target_kv_len) if trans_size > 0 else 0 - if is_current_dp_handle and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): + if is_local_dp_request and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): assert req.hold_kv_len == req.cur_kv_len src_dp_rank = int(available_lens[req_index].argmax()) src_info = self.shared_req_infos.arr[req_index, src_dp_rank] diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 4b72c2ec4f..67ff390a63 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -98,8 +98,8 @@ def _init_reqs(self, reqs: List[Tuple]): infer_reqs = g_infer_context.add_reqs(reqs, init_prefix_cache=True) req_dp_ranks = [req[3] for req in reqs] self.dp_kv_shared_module.fill_reqs_info(reqs=infer_reqs) - trans_taskes = self.dp_kv_shared_module.build_shared_kv_trans_tasks(reqs=infer_reqs, req_dp_ranks=req_dp_ranks) - self.dp_kv_shared_module.kv_trans(trans_tasks=trans_taskes) + transfer_tasks = self.dp_kv_shared_module.build_shared_kv_trans_tasks(reqs=infer_reqs, req_dp_ranks=req_dp_ranks) + self.dp_kv_shared_module.kv_trans(trans_tasks=transfer_tasks) # other_dp_reqs 只是为本 DP 做完 prefix cache / KV 拉取后的临时本地对象， # 真正推理仍在其归属 DP 上进行。这里仅清理本 DP 的 InferReq 与 KV 引用， diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 8723445951..10e540430e 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -255,8 +255,7 @@ def write_req_to_cpu( producer = torch.cuda.current_stream() # 阶段一：等待当前生产流已提交的数据写入，不阻塞主机线程。 cpu_kv_cache_stream.wait_stream(producer) - # TODO：将 bounded CPU→CPU 快照搬运改为当前流上的 kernel。 - # 现有主机 copy 不受 wait_stream 约束，kernel 落地前源快照仍须预先就绪。 + # bounded CPU→CPU 快照也通过当前搬运流上的 Triton kernel 复制，遵循相同流依赖。 with torch.cuda.stream(cpu_kv_cache_stream): # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index bffd16ea24..98cf480995 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -397,6 +397,14 @@ def read_page_to_mems_loop(self): torch.cuda.set_device(self.device_id) while True: trans_task: PDChunckedTransTask = self.ready_page_task_queue.get() + if trans_task.xfer_handle is not None: + # P 的 done 只表示其发送流已完成，不能替代 D 本地接收流的完成确认。 + # 后续会在独立的 copy_cuda_stream 上读取 PD 页；这里先等待 recv 后记录的 + # event，避免搬运流在接收尚未完成时读到旧数据或只写入一部分的页面。 + # 前提：接收线程已提交 event.record(stream)；未记录的 event 会直接返回， + # 此处 synchronize 不会等待未来的 record，也不解决两个主机线程的提交竞态。 + # NIXL 接收端没有本地 Xfer，沿用其远端 WRITE 完成通知流程，不进入此分支。 + trans_task.xfer_handle.event.synchronize() copy_start_event = torch.cuda.Event(enable_timing=True) copy_end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.stream(stream=self.copy_cuda_stream): diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index 54b06bab00..84d94d34e1 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -308,6 +308,10 @@ def send_page(self, trans_task: PDChunckedTransTask) -> _NcclXferHandle: return _NcclXferHandle(peer_name=self.peer_name, event=event) def start_recv(self, trans_task: PDChunckedTransTask): + trans_task.xfer_handle = _NcclXferHandle( + peer_name=self.peer_name, + event=torch.cuda.Event(), + ) self._get_recv_queue().put(copy.copy(trans_task)) return @@ -355,6 +359,7 @@ def _recv_page(self, trans_task: PDChunckedTransTask): comm = self._ensure_comm(is_server=False) stream = self._get_stream() comm.recv(page_tensor, src=0, stream=stream) + trans_task.xfer_handle.event.record(stream) if self.transporter.args.pd_detail_log: logger.debug( f"NCCL recv page done request_id={trans_task.request_id} " f"dst_page={trans_task.dst_page_index}" diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index 1f2544b8fe..b3c1dcc609 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -410,14 +410,14 @@ def _update_task_status_once(self): try: self.transporter.send_write_done_task_to_decode_node(completed_task) except BaseException as e: - logger.exception(f"send WRITE done nixl notify failed: {completed_task.to_str()}") - completed_task.error_info = f"send WRITE done nixl notify failed: {str(e)}" + logger.exception(f"send WRITE done PD notify failed: {completed_task.to_str()}") + completed_task.error_info = f"send WRITE done PD notify failed: {str(e)}" self.failed_queue.put(completed_task) continue if self.args.pd_detail_log: logger.debug( - f"send WRITE done nixl notify " + f"send WRITE done PD notify " f"request_id={completed_task.request_id} " f"kv=[{completed_task.start_kv_index},{completed_task.end_kv_index}) " f"src_page={completed_task.src_page_index} dst_page={completed_task.dst_page_index}" diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 4b5ae89530..12ddbb1c68 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -55,8 +55,17 @@ def get_tokenizer( model_cfg, _ = PretrainedConfig.get_config_dict(tokenizer_name) model_type = model_cfg.get("model_type", "") if model_type == "deepseek_v32": - # Register the custom config before AutoTokenizer tries to resolve it. + # Register the custom config for HuggingFace metadata loading. from ..models import deepseek3_2 # noqa: F401 + from ..models.deepseek3_2.tokenizer import DeepSeekV32Tokenizer + + if tokenizer_mode == "slow": + raise ValueError("DeepSeek-V3.2 requires a fast tokenizer loaded from tokenizer.json") + # checkpoint 声明 LlamaTokenizerFast，但实际使用 ByteLevel；部分 Transformers + # 版本会按 Llama 类重建 SentencePiece 规则，丢失空格和换行，故直接保留 JSON 中的规则。 + tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_name, *args, **kwargs) + logger.info("Using DeepSeek-V3.2 tokenizer mode with Python-based chat template encoding.") + return DeepSeekV32Tokenizer(tokenizer) if model_type == "deepseek_v4": from ..models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer @@ -81,17 +90,6 @@ def get_tokenizer( "slowdown. Consider using a fast tokenizer instead." ) - # DeepSeek-V3.2 custom tokenizer mode: wraps the HF tokenizer with - # a Python-based apply_chat_template that uses encoding_dsv32.py. - if model_type == "deepseek_v32": - from ..models.deepseek3_2.tokenizer import DeepSeekV32Tokenizer - - hf_tokenizer = AutoTokenizer.from_pretrained( - tokenizer_name, trust_remote_code=trust_remote_code, *args, **kwargs - ) - logger.info("Using DeepSeek-V3.2 tokenizer mode with Python-based chat template encoding.") - return DeepSeekV32Tokenizer(hf_tokenizer) - if model_cfg["architectures"][0] == "TarsierForConditionalGeneration": from ..models.tarsier2.model import Tarsier2Tokenizer from ..models.qwen2_vl.vision_process import Qwen2VLImageProcessor diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 69379e4150..52577d4ce1 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -95,7 +95,12 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP first_next_ids=executor._argmax_ids(output.logits), ) if hybrid: - resources.save_bounded_att_runtime_to_pages(source_req.tolist(), page_kind="big", page_slot_ids=[source_slots[page]]) + resources.save_bounded_att_runtime_to_pages( + req_indices=source_req.tolist(), + state_token_lens=[length], + page_kind="big", + page_slot_ids=[source_slots[page]], + ) torch.cuda.synchronize() next_token = output.logits.argmax(dim=-1).cpu() source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() @@ -218,9 +223,14 @@ def decode(request: torch.Tensor): output = draft.forward(batch) logits.append(output.logits.clone()) if draft_models and hybrid: - ( - resources.main_kda_conv_bf16_ssm_fp32_state if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") else resources.main_gdn_conv_bf16_ssm_fp32_state - ).commit_bounded_att_state(request, [accepted_width]) + cuda_req_indices = request.cuda(non_blocking=True) + accepted_steps = torch.full_like(cuda_req_indices, accepted_width, dtype=torch.int32) + for item in (model, *draft_models): + resources.commit_bounded_att_state( + req_indices=cuda_req_indices, + accepted_steps=accepted_steps, + role=item.att_state_role, + ) return logits reference = decode(source_req) diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 51f09bde85..867b6fa8e7 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -261,6 +261,8 @@ def _worker(connection: Connection, backend: str, sender: bool): key=request.get_key(), ) assert done.get_key() == request.get_key() + if backend == "nccl": + request.xfer_handle.event.synchronize() # Match production: unpack on a separate stream after the done notification. with torch.cuda.stream(copy_stream): copy_page( diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index 464f07c558..6fb3770dfb 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -45,6 +45,40 @@ def find_spec(self, fullname, path=None, target=None): assert result.returncode == 0, result.stdout + result.stderr +@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) +def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.MonkeyPatch, run_mode: str) -> None: + backend_class = attention.DeepseekV4FlashMlaFp8SparseAttBackend + model = SimpleNamespace( + args=SimpleNamespace(run_mode=run_mode, enable_prefill_microbatch_overlap=False), + config=dict( + compress_ratios=[0, 4, 128], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + ), + layers_infer=[SimpleNamespace(layer_num_=layer) for layer in range(3)], + tp_rank_=0, + tp_world_size_=4, + batch_max_tokens=4096, + max_req_num=8, + att_state_role="main", + ) + monkeypatch.setattr( + backend_class, "get_gpu_workspace_buffer", + lambda key_name, workspace_size, dtype: torch.empty(workspace_size, dtype=dtype), + ) + workspace, capacity = backend_class.get_swa_prefill_workspace(model=model) + # D 节点仍执行 HOLD 请求的单 token prefill，但无需保留普通请求的长 prefill 临时区。 + assert workspace is not None + if run_mode == "decode": + assert capacity == 0 + assert workspace.shape == (3, 128, 585) + else: + assert capacity >= model.batch_max_tokens + assert workspace.shape == (3, capacity + 128, 585) + + def _pool(page_size, pages): stride = ((page_size * 584 + 575) // 576) * 576 buffer = torch.zeros((pages, stride), dtype=torch.uint8, device="cuda") @@ -266,11 +300,12 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): packed_kv = (packed_kv[0], extra) params["extra_indices"].copy_(indices) + monkeypatch.setattr(backend, "get_att_input_params", lambda infer_state, logical_layer_id: packed_kv) + def new_request(length): infer_state = DeepseekV4InferStateInfo() infer_state.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) infer_state.hidden_collector = NoopHiddenCollector() - infer_state.req_manager = SimpleNamespace(get_att_input_params=lambda layer_index: packed_kv) infer_state.attn_params = dict(params, swa_lengths=torch.full_like(params["swa_lengths"], length)) if ratio: infer_state.attn_params["extra_lengths"] = torch.full_like(params["extra_lengths"], length) diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index 71a348012e..a4549dc078 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -946,9 +946,9 @@ def build(): "Model", (), { - "config": dict(num_attention_heads=64, compress_ratios=[0], head_dim=512, + "config": dict(num_attention_heads=64, compress_ratios=[0, 0, 0], head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128), - "layers_infer": [SimpleNamespace(layer_num_=0)], + "layers_infer": [SimpleNamespace(layer_num_=2)], "tp_rank_": 0, "tp_world_size_": 1, "req_manager": requests, diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 6e830f2a2b..8394d18b48 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -7,6 +7,29 @@ from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "capture"]) +@pytest.mark.parametrize("draft", [False, True], ids=["main", "dspark"]) +def test_pre_layer_decode_uses_current_graph_flag(monkeypatch: pytest.MonkeyPatch, capture: bool, draft: bool) -> None: + from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo + from lightllm.models.deepseek_v4.layer_infer import pre_layer_infer + from lightllm.models.deepseek_v4_dspark.layer_infer.pre_layer_infer import DeepseekV4DSparkPreLayerInfer + + layer_class = DeepseekV4DSparkPreLayerInfer if draft else pre_layer_infer.DeepseekV4PreLayerInfer + layer = layer_class.__new__(layer_class) + layer.hc_mult = 4 + state = DeepseekV4InferStateInfo() + state.is_cuda_graph_capture = capture + hidden = torch.ones((1, 8)) + state.mtp_draft_input_hiddens = hidden + monkeypatch.setattr(pre_layer_infer.LlamaMultimodalPreLayerInfer, "token_forward", lambda *args: hidden) + monkeypatch.setattr(pre_layer_infer, "hc_expand", lambda embeddings, streams: embeddings.repeat(1, streams)) + + output = layer.token_forward(input_ids=None, infer_state=state, layer_weight=None) + assert output.shape == (1, 32) + # 捕获时保留输入引用供后续回放使用；普通 decode 可以释放不再使用的 draft 特征。 + assert state.mtp_draft_input_hiddens is (hidden if capture else None) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_rotary_cache_keeps_unscaled_swa_and_yarn_compression(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py index 135dd4fb4a..173dd8f7fe 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -4,6 +4,8 @@ import concurrent.futures import socket import threading +from types import SimpleNamespace +from unittest.mock import Mock import rpyc @@ -13,6 +15,37 @@ import pytest +def test_page_events_are_recorded_after_send_and_receive(monkeypatch: pytest.MonkeyPatch): + page_buffers = module.torch.empty(size=(2, 8), dtype=module.torch.uint8, device="cpu") + transporter = SimpleNamespace(kv_move_buffer=page_buffers, args=SimpleNamespace(pd_detail_log=False)) + peer = module._NcclPeer(transporter, "remote") + stream = object() + operations = [] + comm = SimpleNamespace( + send=Mock(side_effect=lambda tensor, dst, stream: operations.append(("send", tensor.numel(), dst, stream))), + recv=Mock(side_effect=lambda tensor, src, stream: operations.append(("recv", tensor.numel(), src, stream))), + ) + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: comm) + monkeypatch.setattr(peer, "_get_stream", lambda: stream) + event = Mock() + event.record.side_effect = lambda stream: operations.append(("record", stream)) + monkeypatch.setattr(module.torch.cuda, "Event", lambda: event) + task = SimpleNamespace(src_page_index=1, dst_page_index=0) + + handle = peer.send_page(task) + assert operations == [("send", 8, 1, stream), ("record", stream)] + assert handle.event is event + + operations.clear() + receive_queue = Mock() + monkeypatch.setattr(peer, "_get_recv_queue", lambda: receive_queue) + peer.start_recv(task) + queued_task = receive_queue.put.call_args.args[0] + assert queued_task.xfer_handle is task.xfer_handle + peer._recv_page(queued_task) + assert operations == [("recv", 8, 0, stream), ("record", stream)] + + def test_concurrent_control_channels_retry_port_before_publishing_metadata(monkeypatch: pytest.MonkeyPatch): with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 737e5e2165..24714e9b71 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -255,14 +255,15 @@ def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disabl req.multimodal_params = SimpleNamespace( to_dict=lambda: { "images": [ - {"start_idx": 450, "token_num": 150}, - {"start_idx": 700, "token_num": 200}, + {"start_idx": 450, "token_id": 450, "token_num": 150}, + {"start_idx": 700, "token_id": 700, "token_num": 200}, ], "audios": [], } ) req._init_all_state() assert req.hybrid_cache_len == 768 + assert req.mm_input_spans == ([(450, 600), (700, 900)] if disable_chunked_mm_input else []) @pytest.mark.parametrize("hybrid", [False, True]) diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py index 2accf0ce89..88e695fac7 100644 --- a/unit_tests/server/test_tokenizer.py +++ b/unit_tests/server/test_tokenizer.py @@ -1,5 +1,6 @@ """Tokenizer entry points must not eagerly load unrelated model implementations.""" +import json import os from pathlib import Path import subprocess @@ -13,6 +14,42 @@ ROOT = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize("model_type", ["deepseek_v32", "deepseek_v3"]) +def test_deepseek_bytelevel_checkpoint_preserves_whitespace(tmp_path: Path, model_type: str) -> None: + from tokenizers import Tokenizer, decoders, pre_tokenizers, trainers + from tokenizers.models import BPE + from transformers import PreTrainedTokenizerFast + from lightllm.server.tokenizer import get_tokenizer + + text = "Question: What is 2 + 3?\nAnswer: #### 5" + backend = Tokenizer(BPE(unk_token="")) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) + backend.decoder = decoders.ByteLevel() + backend.train_from_iterator( + [text], + trainers.BpeTrainer( + vocab_size=300, + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), + special_tokens=[""], + ), + ) + fast = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="") + fast.save_pretrained(tmp_path) + # 真实 DeepSeek checkpoint 声明 LlamaTokenizerFast，但实际编码和解码均为 ByteLevel。 + config_path = tmp_path / "tokenizer_config.json" + config = json.loads(config_path.read_text()) + config.update(tokenizer_class="LlamaTokenizerFast", legacy=True) + config_path.write_text(json.dumps(config)) + (tmp_path / "config.json").write_text(json.dumps({ + "model_type": model_type, "architectures": ["DeepseekV3ForCausalLM"], + })) + + tokenizer = get_tokenizer(str(tmp_path)) + expected_ids = backend.encode(text).ids + assert tokenizer.encode(text, add_special_tokens=False) == expected_ids + assert tokenizer.decode(expected_ids, skip_special_tokens=True) == text + + @pytest.mark.parametrize("model_type", ["llama", "deepseek_v32", "deepseek_v4"]) def test_tokenizer_imports_are_isolated(model_type): result = subprocess.run( @@ -48,13 +85,6 @@ def find_spec(self, fullname, path=None, target=None): from lightllm.server import tokenizer as module assert not any(name.startswith("lightllm.models.") for name in sys.modules) -if model_type == "deepseek_v32": - original = module.AutoTokenizer.from_pretrained - def load(*args, **kwargs): - assert AutoConfig.for_model("deepseek_v32").model_type == "deepseek_v32" - return original(*args, **kwargs) - module.AutoTokenizer.from_pretrained = load - with tempfile.TemporaryDirectory() as directory: fast = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]"))) fast.save_pretrained(directory) @@ -64,6 +94,7 @@ def load(*args, **kwargs): tokenizer = module.get_tokenizer(directory) assert tokenizer.encode("hello") == [0] if model_type == "deepseek_v32": + assert AutoConfig.for_model("deepseek_v32").model_type == "deepseek_v32" assert tokenizer.__class__.__module__ == "lightllm.models.deepseek3_2.tokenizer" prompt = tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}]) assert "hello" in prompt From 472a689ac7372eed1e6d275dc6f09ee1f827421d Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 01:45:19 +0000 Subject: [PATCH 32/56] test: align PD and V4 fixtures with runtime initialization --- unit_tests/common/att_resources/test_dsv4_integration.py | 5 ++++- .../router/model_infer/mode_backend/test_pd_token_indices.py | 1 + 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py index f620b230ef..355704a32b 100644 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -243,10 +243,12 @@ def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): dspark_markov_rank=4, ) args = SimpleNamespace( - llm_kv_type="None", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", run_mode="decode" + llm_kv_type="None", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", run_mode="decode", + enable_prefill_microbatch_overlap=False, ) main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) main.config, main.args, main.max_req_num, main.tp_rank_, main.tp_world_size_ = config, args, 2, 0, 1 + main.batch_max_tokens = 2048 main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] main.torch_memory_saver = SimpleNamespace(region=lambda **kwargs: nullcontext()) main.prepare_attention_states() @@ -258,6 +260,7 @@ def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): draft.config = cls.normalize_model_config(config) assert draft.config == cls.normalize_model_config(draft.config) draft.args, draft.max_req_num, draft.tp_rank_, draft.tp_world_size_ = args, 2, 0, 1 + draft.batch_max_tokens = main.batch_max_tokens draft.torch_memory_saver = main.torch_memory_saver first_layer = 3 + sum(len(model.layers_infer) for model in drafts) draft.layers_infer = [ diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py index 01addc80f9..c281d95899 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -69,6 +69,7 @@ def capture_copy(**kwargs: object) -> None: decode_dp_index=0, req_idx=1, end_kv_index=7, + xfer_handle=None, ) for page, kind, indices in [(0, "kv", [5, 1, 7]), (1, "kv", [2]), (0, "att_state", None)] ] From 26b1f67bef766d409315f37cbf875192ff8a295d Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 02:14:30 +0000 Subject: [PATCH 33/56] perf: remove redundant CPU cache offload stream wait --- .../model_infer/mode_backend/multi_level_kv_cache.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 10e540430e..54839a5e80 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -252,17 +252,14 @@ def write_req_to_cpu( """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers - producer = torch.cuda.current_stream() - # 阶段一：等待当前生产流已提交的数据写入，不阻塞主机线程。 - cpu_kv_cache_stream.wait_stream(producer) - # bounded CPU→CPU 快照也通过当前搬运流上的 Triton kernel 复制，遵循相同流依赖。 + # KV 和 bounded 快照均在搬运流上复制，由后续完成事件统一确认就绪。 with torch.cuda.stream(cpu_kv_cache_stream): # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() page_len_list = req.shm_req.token_hash_page_len_list.get_all() assert len(token_hash_list) == len(page_len_list) - # 阶段二：确定本次可写出范围，申请 CPU 页。 + # 阶段一：确定本次可写出范围，申请 CPU 页。 if self.backend.is_master_in_dp: find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) @@ -319,7 +316,7 @@ def write_req_to_cpu( assert req.cur_kv_len >= move_token_num token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num] - # 阶段三：只查找实际分配的 CPU 页所需快照；完整大页必须存在，不做缺失回退。 + # 阶段二：只查找实际分配的 CPU 页所需快照；完整大页必须存在，不做缺失回退。 resources = self.backend.model.att_state_manager big_page_slot_ids, small_page_slot_ids = None, None if self.backend.is_hybrid_att_model: From 7c4d82fa4d832625adfbde6a629617908490b916 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 02:48:04 +0000 Subject: [PATCH 34/56] fix: unify NCCL send and receive status polling --- docs/cache_manager_design.md | 2 +- .../decode_node_impl/decode_trans_process.py | 22 ++++-- .../mode_backend/pd/nccl_kv_transporter.py | 47 +++++++++--- .../att_resources/test_network_transport.py | 6 +- .../mode_backend/test_nccl_control_channel.py | 71 +++++++++++++++++-- .../mode_backend/test_pd_token_indices.py | 21 ++++++ 6 files changed, 145 insertions(+), 24 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 5efe8e3df1..212c1a4706 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1027,7 +1027,7 @@ PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效 CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 -NCCL 接收线程在 recv 后记录本地 Xfer 的 event，Decode 在读取 PD 页前调用 event.synchronize()，不增加反向确认传输。该同步要求 record 已经提交；未记录的 event 会直接返回，P 的 done 通知本身不能保证 D 的接收线程已执行 record。当前最小实现不额外协调这两个主机线程的提交顺序。 +NCCL Xfer 用 mode 区分 send/recv 句柄，二者共用 finish_cuda_event。接收线程在 recv 后创建并记录 event，record 返回后才发布；在此之前 recv 句柄的 check_status() 保持 PROC，避免未记录的 CUDA event 被误判为完成。Decode 与 NIXL 统一轮询 check_task_status()，直到状态为 DONE 后才读取 PD 页。接收提交失败时返回 ERR，由读页线程转入已有失败队列，不读取无效页面。 #### 调用方与资源实现的边界 diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index 98cf480995..ca4b51e9b5 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -398,13 +398,21 @@ def read_page_to_mems_loop(self): while True: trans_task: PDChunckedTransTask = self.ready_page_task_queue.get() if trans_task.xfer_handle is not None: - # P 的 done 只表示其发送流已完成，不能替代 D 本地接收流的完成确认。 - # 后续会在独立的 copy_cuda_stream 上读取 PD 页；这里先等待 recv 后记录的 - # event，避免搬运流在接收尚未完成时读到旧数据或只写入一部分的页面。 - # 前提：接收线程已提交 event.record(stream)；未记录的 event 会直接返回， - # 此处 synchronize 不会等待未来的 record，也不解决两个主机线程的提交竞态。 - # NIXL 接收端没有本地 Xfer，沿用其远端 WRITE 完成通知流程，不进入此分支。 - trans_task.xfer_handle.event.synchronize() + # 只有 NCCL 接收端会创建本地 Xfer 句柄；NIXL 的远端 WRITE 完成通知 + # 已代表其写入完成，不经过这里的本地 CUDA event 轮询。 + # P 的 done 不代表 D 本地接收已完成；统一轮询传输状态， + # 直到接收 event 完成后，独立搬运流才读取 PD 页。 + while True: + status = self.transporter.check_task_status(trans_task=trans_task) + if status == "PROC": + time.sleep(0.001) + continue + if status != "DONE": + trans_task.error_info = trans_task.xfer_handle.error_info + self.failed_queue.put(trans_task) + break + if status != "DONE": + continue copy_start_event = torch.cuda.Event(enable_timing=True) copy_end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.stream(stream=self.copy_cuda_stream): diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index 84d94d34e1..efacb98531 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -264,16 +264,37 @@ def _get_notify_source_agent_name(self, notify: bytes) -> str: @dataclass class _NcclXferHandle: peer_name: str - event: torch.cuda.Event + mode: str + finish_cuda_event: Optional[torch.cuda.Event] = None status: str = "PROC" error_info: Optional[str] = None + def __post_init__(self) -> None: + assert self.mode in ("send", "recv"), f"invalid NCCL handle mode: {self.mode}" + def check_status(self) -> str: if self.status != "PROC": return self.status + if self.mode == "send": + return self._check_send_status() + return self._check_recv_status() + + def _check_send_status(self) -> str: + if self.finish_cuda_event is None: + return "PROC" + try: + if self.finish_cuda_event.query(): + self.status = "DONE" + except BaseException as e: + self.status = "ERR" + self.error_info = str(e) + return self.status + def _check_recv_status(self) -> str: + if self.finish_cuda_event is None: + return "PROC" try: - if self.event.query(): + if self.finish_cuda_event.query(): self.status = "DONE" except BaseException as e: self.status = "ERR" @@ -297,21 +318,22 @@ def send_page(self, trans_task: PDChunckedTransTask) -> _NcclXferHandle: stream = self._get_stream() comm.send(page_tensor, dst=1, stream=stream) - event = torch.cuda.Event() - event.record(stream) + send_finish_cuda_event = torch.cuda.Event() + send_finish_cuda_event.record(stream) if self.transporter.args.pd_detail_log: logger.debug( f"NCCL send page posted request_id={trans_task.request_id} " f"src_page={trans_task.src_page_index} dst_agent={self.peer_name}" ) - return _NcclXferHandle(peer_name=self.peer_name, event=event) - - def start_recv(self, trans_task: PDChunckedTransTask): - trans_task.xfer_handle = _NcclXferHandle( + return _NcclXferHandle( peer_name=self.peer_name, - event=torch.cuda.Event(), + mode="send", + finish_cuda_event=send_finish_cuda_event, ) + + def start_recv(self, trans_task: PDChunckedTransTask): + trans_task.xfer_handle = _NcclXferHandle(peer_name=self.peer_name, mode="recv") self._get_recv_queue().put(copy.copy(trans_task)) return @@ -359,12 +381,17 @@ def _recv_page(self, trans_task: PDChunckedTransTask): comm = self._ensure_comm(is_server=False) stream = self._get_stream() comm.recv(page_tensor, src=0, stream=stream) - trans_task.xfer_handle.event.record(stream) + finish_cuda_event = torch.cuda.Event() + finish_cuda_event.record(stream) + # 必须在 record 返回后才发布；CUDA event 自身不会等待未来的 record。 + trans_task.xfer_handle.finish_cuda_event = finish_cuda_event if self.transporter.args.pd_detail_log: logger.debug( f"NCCL recv page done request_id={trans_task.request_id} " f"dst_page={trans_task.dst_page_index}" ) except BaseException as e: + trans_task.xfer_handle.status = "ERR" + trans_task.xfer_handle.error_info = str(e) trans_task.error_info = str(e) logger.exception(str(e)) self._drop_comm() diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 867b6fa8e7..1f11415b1c 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -262,7 +262,11 @@ def _worker(connection: Connection, backend: str, sender: bool): ) assert done.get_key() == request.get_key() if backend == "nccl": - request.xfer_handle.event.synchronize() + deadline = time.monotonic() + 60 + while transporter.check_task_status(request) == "PROC": + assert time.monotonic() < deadline, "PD receive timed out" + time.sleep(0.002) + assert transporter.check_task_status(request) == "DONE" # Match production: unpack on a separate stream after the done notification. with torch.cuda.stream(copy_stream): copy_page( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py index 173dd8f7fe..54bb5db743 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -1,6 +1,5 @@ from __future__ import annotations -from typing import TYPE_CHECKING import concurrent.futures import socket import threading @@ -8,12 +7,10 @@ from unittest.mock import Mock import rpyc +import pytest from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as module -if TYPE_CHECKING: - import pytest - def test_page_events_are_recorded_after_send_and_receive(monkeypatch: pytest.MonkeyPatch): page_buffers = module.torch.empty(size=(2, 8), dtype=module.torch.uint8, device="cpu") @@ -34,7 +31,10 @@ def test_page_events_are_recorded_after_send_and_receive(monkeypatch: pytest.Mon handle = peer.send_page(task) assert operations == [("send", 8, 1, stream), ("record", stream)] - assert handle.event is event + assert handle.mode == "send" + assert handle.finish_cuda_event is event + event.query.return_value = True + assert handle.check_status() == "DONE" operations.clear() receive_queue = Mock() @@ -42,8 +42,69 @@ def test_page_events_are_recorded_after_send_and_receive(monkeypatch: pytest.Mon peer.start_recv(task) queued_task = receive_queue.put.call_args.args[0] assert queued_task.xfer_handle is task.xfer_handle + assert task.xfer_handle.mode == "recv" + assert task.xfer_handle.finish_cuda_event is None peer._recv_page(queued_task) assert operations == [("recv", 8, 0, stream), ("record", stream)] + assert task.xfer_handle.finish_cuda_event is event + assert task.xfer_handle.check_status() == "DONE" + assert event.query.call_count == 2 + + +def test_receive_waits_for_record_to_return(monkeypatch: pytest.MonkeyPatch): + record_started = threading.Event() + finish_record = threading.Event() + event = Mock() + + def record(stream: object) -> None: + record_started.set() + assert finish_record.wait(timeout=5) + + event.record.side_effect = record + monkeypatch.setattr(module.torch.cuda, "Event", lambda: event) + transporter = SimpleNamespace( + kv_move_buffer=module.torch.empty(size=(1, 8), dtype=module.torch.uint8), + args=SimpleNamespace(pd_detail_log=False), + ) + peer = module._NcclPeer(transporter, "remote") + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: SimpleNamespace(recv=Mock())) + monkeypatch.setattr(peer, "_get_stream", object) + handle = module._NcclXferHandle(peer_name="remote", mode="recv") + task = SimpleNamespace(dst_page_index=0, xfer_handle=handle) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + receive = executor.submit(peer._recv_page, task) + try: + assert record_started.wait(timeout=5) + assert handle.check_status() == "PROC" + assert handle.finish_cuda_event is None + event.query.assert_not_called() + finally: + finish_record.set() + receive.result(timeout=5) + assert handle.check_status() == "DONE" + assert handle.finish_cuda_event is event + event.query.assert_called_once() + + +def test_receive_submission_failure_wakes_waiter(monkeypatch: pytest.MonkeyPatch): + transporter = SimpleNamespace( + kv_move_buffer=module.torch.empty(size=(1, 8), dtype=module.torch.uint8), + send_error_info_to_prefill_node=Mock(), + ) + peer = module._NcclPeer(transporter, "remote") + comm = SimpleNamespace(recv=Mock(side_effect=RuntimeError("recv failed"))) + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: comm) + monkeypatch.setattr(peer, "_get_stream", object) + monkeypatch.setattr(peer, "_drop_comm", Mock()) + handle = module._NcclXferHandle(peer_name="remote", mode="recv") + task = SimpleNamespace(dst_page_index=0, xfer_handle=handle) + + peer._recv_page(task) + + assert handle.check_status() == "ERR" + assert handle.error_info == "recv failed" + transporter.send_error_info_to_prefill_node.assert_called_once_with(task) def test_concurrent_control_channels_retry_port_before_publishing_metadata(monkeypatch: pytest.MonkeyPatch): diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py index c281d95899..a6f9c623e3 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -2,6 +2,7 @@ from types import SimpleNamespace from unittest.mock import Mock +from queue import Queue import pytest import torch @@ -101,3 +102,23 @@ def capture_copy(**kwargs: object) -> None: class _StopLoop(Exception): pass + + +def test_failed_receive_is_not_read_or_copied(monkeypatch: pytest.MonkeyPatch) -> None: + module = decode_trans_process._DecodeTransModule.__new__(decode_trans_process._DecodeTransModule) + module.device_id = 0 + module.failed_queue = Queue() + handle = SimpleNamespace(check_status=Mock(return_value="ERR"), error_info="recv failed") + task = SimpleNamespace(xfer_handle=handle) + module.transporter = SimpleNamespace(check_task_status=Mock(return_value="ERR")) + module.ready_page_task_queue = SimpleNamespace(get=Mock(side_effect=[task, _StopLoop()])) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + copy = Mock() + monkeypatch.setattr(decode_trans_process, "copy_page", copy) + + with pytest.raises(_StopLoop): + module.read_page_to_mems_loop() + + assert module.failed_queue.get_nowait() is task + assert task.error_info == "recv failed" + copy.assert_not_called() From 30650a3a455aea12170cf0e4b50705f8893846ed Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 03:37:34 +0000 Subject: [PATCH 35/56] refactor MLA logical configuration and factories --- docs/cache_manager_design.md | 15 +- .../attention_state/full_att/__init__.py | 14 +- .../full_att/create_mla_state.py | 56 +------ .../full_att/create_packed_mla_state.py | 38 +++++ .../attention_state/full_att/mla_state.py | 20 +-- .../full_att/packed_mla_state.py | 151 ++++++++++++------ lightllm/models/deepseek2/model.py | 8 +- lightllm/models/deepseek3_2/model.py | 12 +- lightllm/models/glm5_next/model.py | 13 +- lightllm/models/glm5_next_mtp/model.py | 14 +- .../common/att_resources/manager_fixtures.py | 6 +- .../common/att_resources/model_helpers.py | 4 +- .../common/att_resources/test_capacity.py | 34 ++-- .../test_fp8_e8m0_packed_mla_state.py | 34 ++-- .../common/att_resources/test_full_state.py | 6 +- .../common/att_resources/test_layer_groups.py | 3 +- .../att_resources/test_model_resources.py | 14 +- .../test_quantized_full_state.py | 11 +- .../att_resources/test_state_creation.py | 57 +++---- .../common/att_resources/test_tp_config.py | 6 +- unit_tests/models/glm5_next/test_cache.py | 12 +- unit_tests/models/glm5_next/test_fp8_kv.py | 10 +- unit_tests/models/glm5_next/test_pd_cache.py | 14 +- 23 files changed, 317 insertions(+), 235 deletions(-) create mode 100644 lightllm/common/attention_state/full_att/create_packed_mla_state.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 212c1a4706..90a3b76e0d 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -402,7 +402,7 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交直 | 逻辑类型 | 创建入口 | 参数决定的实现 | | --- | --- | --- | | GQA（含 MHA/MQA） | `full_att.create_gqa_state(config, *, kv_type, page_size, device, ...)` | 普通浮点、INT4/INT8 分组量化、per-tensor/per-head 静态 FP8 | -| MLA | `full_att.create_mla_state(config, kv_type, packed, page_size, device, fp8_format, ...)` | 普通 MLA、packed MLA、group128/FP32 或 group64/E8M0 packed FP8 MLA | +| MLA | `full_att.create_mla_state(config, page_size, device)` / `create_packed_mla_state(config, page_size, device)` | 普通 latent MLA 与带 indexer 的 packed MLA 分开创建 | | 序列压缩 attention | `full_att.create_compressed_att_state(config, page_size, device)` | FP8/E8M0 packed 压缩 KV 与可选 indexer；V4 CSA/C4、HCA/C128 由模型分别配置 | | KDA | `bounded_att.create_kda_state(config, request_capacity, mtp_step=0, *, device)` | KDA Conv/SSM 运行态 | | GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | @@ -411,7 +411,7 @@ token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交直 每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器。K-pool 复用窗口压缩器入口。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 -GQA/MLA 的创建入口分别接收 `GQAAttentionConfig`、`MLAAttentionConfig`，每种具体 State 使用独立配置，不继承通用配置。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。MLA 的 `config.latent_dim` 始终传逻辑 latent 维度，`packed=True` 时由 packed 配置推导包含 indexer、scale 和对齐的物理行宽；`"fp8kv_dsa"` 仅用于 packed MLA。 +GQA/MLA 的创建入口分别接收 `GQAAttentionConfig`、`MLAConfig` 或 `PackedMLALogicalConfig`，每种具体 State 使用独立配置，不继承通用配置。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。普通 MLA 只描述 latent；packed MLA 只描述带 indexer 的逻辑 format，工厂内部再转换为具体物理布局。 MLA 的 packed FP8 编码由 `fp8_format` 显式指定：默认 `"group128_fp32"` 使用原有 `PackedMLAFullAttentionState`；`"group64_e8m0"` 使用独立的 `FP8E8M0PackedMLAState`， @@ -543,7 +543,7 @@ FullAttentionState └── CompressedAttentionState # V4 的 C4/C128 压缩块 ``` -所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。配置也不建立继承关系：普通 GQA 使用 `GQAAttentionConfig`；分组量化使用 `GroupQuantizedAttentionConfig`，以 `scale_dtype` 明确 scale 精度；静态 FP8 使用 `StaticFP8AttentionConfig`，不要求调用方传入固定的物理 dtype。普通 MLA 使用 `MLAAttentionConfig`；packed MLA 使用 `PackedMLAAttentionConfig`，以 `latent_dim` 表达逻辑宽度，`dtype` 直接记录 FP16、BF16 或 `torch.float8_e4m3fn` 编码。FP8 行的 uint8 承载类型只在存储分配函数中确定，不提供 `storage_dtype` 配置或属性；物理行宽及 indexer 偏移由编码推导。MLA 固定单份 latent 并在 TP rank 间复制，不提供 GQA 的 head 数或 `replicated` 开关。 +所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。配置也不建立继承关系：普通 GQA 使用 `GQAAttentionConfig`；分组量化使用 `GroupQuantizedAttentionConfig`，以 `scale_dtype` 明确 scale 精度；静态 FP8 使用 `StaticFP8AttentionConfig`，不要求调用方传入固定的物理 dtype。普通 MLA 使用 `MLAConfig`，packed MLA 使用带 indexer 的 `PackedMLALogicalConfig`；工厂再转换为具体 State 的物理配置。具体 packed 配置使用 `kv_encoding` 显式选择 `fp16`、`bf16` 或 `fp8_group128`，`dtype` 由编码提供只读 property。FP8 行的 uint8 承载类型只在存储分配函数中确定，不提供 `storage_dtype` 配置或属性；物理行宽及 indexer 偏移由编码推导。MLA 固定单份 latent 并在 TP rank 间复制，不提供 GQA 的 head 数或 `replicated` 开关。 `gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 @@ -1704,7 +1704,7 @@ BaseModel 提供两个默认返回空列表的 `classmethod` 扩展点，具体 需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_attention_states()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 -`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig` 或 `MLAAttentionConfig` 并调用 `create_gqa_state()` 或 `create_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded 工厂同样直接调用算法 create；GLM main/draft 分别选择池化窗口的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 +`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig`、`MLAConfig` 或 `PackedMLALogicalConfig`，并调用对应的 `create_gqa_state()`、`create_mla_state()` 或 `create_packed_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded 工厂同样直接调用算法 create；GLM main/draft 分别选择池化窗口的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 `AttentionStateManager.initialize_storage()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `allocate_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 @@ -1738,8 +1738,9 @@ State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_si | `GQAAttentionConfig` | 逻辑层分组、全局 KV head 数、`head_dim`、KV 浮点 `dtype`、TP 坐标与 `replicated` | `head_num` 返回当前 rank 的单份 K 或 V 的 head 数 | | `GroupQuantizedAttentionConfig` | GQA 几何与 TP 信息、`bits`、`group_size`、动态 `scale_dtype` | `head_num`；KV 编码固定为 INT4/INT8，scale 精度独立描述 | | `StaticFP8AttentionConfig` | GQA 几何与 TP 信息、`scale_mode`，不传物理 dtype 或标定 Tensor | `head_num`；FP8 编码固定，标定通过独立流程绑定 | -| `MLAAttentionConfig` | 逻辑层分组、`latent_dim`、浮点 `dtype` 和 TP 坐标 | 固定复制单份 latent，无 head 数配置 | -| `PackedMLAAttentionConfig` | 逻辑层分组、`latent_dim`、KV 编码 `dtype`（FP16/BF16/FP8）、indexer 与编码参数、TP 坐标 | `kv_bytes`、`token_bytes`、`indexer_offset` 从编码推导；FP8 行使用内部 uint8 buffer 承载混合字段，不另设 `storage_dtype` 属性 | +| `MLAConfig` | 普通 MLA 的逻辑层分组、`latent_dim`、未压缩 latent 的 `latent_dtype` 和 TP 坐标 | `create_mla_state()` 只创建普通 latent State | +| `PackedMLALogicalConfig` | 带 indexer 的逻辑层分组、`latent_dim`、packed `format`、indexer 宽度和 TP 坐标 | `create_packed_mla_state()` 再转换为具体物理布局；浮点精度由 `format` 决定 | +| `PackedMLAAttentionConfig` | 持有 `PackedMLALogicalConfig`，只保留量化分组等物理布局参数 | `kv_bytes`、`token_bytes`、`indexer_offset` 从逻辑格式、NoPE/RoPE 宽度和 scale 数量推导；FP8 行使用内部 uint8 buffer 承载混合字段，不另设 `storage_dtype` 属性 | | `ConvSSMStateConfig` | 层、`global_k_heads`、`global_v_heads`、TP 坐标及 Conv/SSM 维度/dtype 信息；不保存资源身份 | `num_k_heads`、`num_v_heads` 分别返回本地 head 数，`conv_channels` 返回本地卷积通道数 | GQA 的 `replicated=True` 表示每个 rank 保存全部 head；MLA 固定复制单份 latent，不设置该参数。GQA 的 `replicated=False` 时按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 @@ -1761,7 +1762,7 @@ flowchart LR ``` -GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()` / `create_mla_state()` 继续负责量化版本、存储布局及算法专有的 TP 副本规则。 +GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()`、`create_mla_state()` 和 `create_packed_mla_state()` 分别负责具体存储布局及算法专有的 TP 副本规则。 CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_page_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index d3e1f56a9d..88a3f85e2f 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -1,16 +1,18 @@ from .create_mla_state import create_mla_state +from .create_packed_mla_state import create_packed_mla_state from .create_gqa_state import create_gqa_state from .base_state import FullAttentionState from .gqa_state import GQAAttentionConfig, GQAFullAttentionState -from .mla_state import MLAAttentionConfig, MLAFullAttentionState +from .mla_state import MLAConfig, MLAFullAttentionState __all__ = [ "create_gqa_state", "create_mla_state", + "create_packed_mla_state", "FullAttentionState", "GQAFullAttentionState", "GQAAttentionConfig", - "MLAAttentionConfig", + "MLAConfig", "MLAFullAttentionState", ] @@ -19,7 +21,12 @@ GroupQuantizedAttentionConfig, ) from .static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig -from .packed_mla_state import PackedMLAFullAttentionState, PackedMLAAttentionConfig, PackedMLALayerView +from .packed_mla_state import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, + PackedMLALayerView, +) __all__ += [ "GroupQuantizedFullAttentionState", @@ -28,6 +35,7 @@ "StaticFP8AttentionConfig", "PackedMLAFullAttentionState", "PackedMLAAttentionConfig", + "PackedMLALogicalConfig", "PackedMLALayerView", ] diff --git a/lightllm/common/attention_state/full_att/create_mla_state.py b/lightllm/common/attention_state/full_att/create_mla_state.py index bd2e5a6ecb..3c36883a89 100644 --- a/lightllm/common/attention_state/full_att/create_mla_state.py +++ b/lightllm/common/attention_state/full_att/create_mla_state.py @@ -1,62 +1,18 @@ -"""MLA 的统一创建入口，选择普通、packed 或 packed FP8 State。""" +"""Independent creation entry points for normal and packed MLA States.""" from __future__ import annotations -from typing import Literal - import torch from .base_state import FullAttentionState -from .mla_state import MLAAttentionConfig, MLAFullAttentionState -from .packed_mla_state import PackedMLAAttentionConfig, PackedMLAFullAttentionState -from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState +from .mla_state import MLAConfig, MLAFullAttentionState def create_mla_state( - config: MLAAttentionConfig, - kv_type: Literal["None", "fp8kv_dsa"] = "None", - packed: bool = False, - indexer_dim: int = 128, + config: MLAConfig, page_size: int = 1, device: str | torch.device = "cuda", - fp8_format: Literal["group128_fp32", "group64_e8m0"] = "group128_fp32", ) -> FullAttentionState: - """创建普通或 packed MLA State，统一处理量化编码、TP 副本和物理行宽。 - - config.latent_dim 始终是逻辑 KV latent 维度，尚未包含 indexer、scale 和 padding。 - packed=True 时把 indexer 字段存入同一行；fp8kv_dsa 仅适用于 packed MLA。 - fp8_format 由模型显式选择，仅在 packed FP8 时生效；不根据 latent 维度推断编码。 - group128_fp32 与 group64_e8m0 使用独立 State。 - """ - if not packed: - assert kv_type == "None", f"unsupported MLA encoding: {kv_type}" - return MLAFullAttentionState( - layout_config=config, - page_size=page_size, - device=device, - ) - assert kv_type in ("None", "fp8kv_dsa"), f"unsupported packed MLA encoding: {kv_type}" - fp8 = kv_type == "fp8kv_dsa" - assert not fp8 or fp8_format in ("group128_fp32", "group64_e8m0"), f"unknown packed MLA FP8 format: {fp8_format}" - if fp8 and fp8_format == "group64_e8m0": - assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" - return FP8E8M0PackedMLAState( - layout_config=FP8E8M0PackedMLAConfig( - logical_layer_groups=config.logical_layer_groups, - tp_rank=config.tp_rank, - tp_world_size=config.tp_world_size, - indexer_dim=indexer_dim, - ), - page_size=page_size, - device=device, - ) - packed_config = PackedMLAAttentionConfig( - logical_layer_groups=config.logical_layer_groups, - latent_dim=config.latent_dim, - dtype=torch.float8_e4m3fn if fp8 else config.dtype, - tp_rank=config.tp_rank, - tp_world_size=config.tp_world_size, - indexer_dim=indexer_dim, - fp8_nope_dim=512 if fp8 else 0, - ) - return PackedMLAFullAttentionState(layout_config=packed_config, page_size=page_size, device=device) + """Create a normal latent MLA State without indexer fields.""" + assert type(config) is MLAConfig, "normal MLA creation requires MLAConfig" + return MLAFullAttentionState(layout_config=config, page_size=page_size, device=device) diff --git a/lightllm/common/attention_state/full_att/create_packed_mla_state.py b/lightllm/common/attention_state/full_att/create_packed_mla_state.py new file mode 100644 index 0000000000..00c72638a1 --- /dev/null +++ b/lightllm/common/attention_state/full_att/create_packed_mla_state.py @@ -0,0 +1,38 @@ +"""Creation entry point for packed MLA States.""" + +from __future__ import annotations + +import torch + +from .base_state import FullAttentionState +from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + + +def create_packed_mla_state( + config: PackedMLALogicalConfig, + page_size: int = 1, + device: str | torch.device = "cuda", +) -> FullAttentionState: + """Create a packed MLA State; packed formats always include an indexer field.""" + assert type(config) is PackedMLALogicalConfig, "packed MLA creation requires PackedMLALogicalConfig" + if config.format == "packed_group64_e8m0": + assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, ( + "FP8/E8M0 MLA requires a 512-dimensional latent" + ) + return FP8E8M0PackedMLAState( + layout_config=FP8E8M0PackedMLAConfig( + logical_layer_groups=config.logical_layer_groups, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + indexer_dim=config.indexer_dim, + ), + page_size=page_size, + device=device, + ) + + return PackedMLAFullAttentionState( + layout_config=PackedMLAAttentionConfig(logical_config=config), + page_size=page_size, + device=device, + ) diff --git a/lightllm/common/attention_state/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla_state.py index da82e0cddd..cb49d5efcd 100644 --- a/lightllm/common/attention_state/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla_state.py @@ -20,8 +20,8 @@ @dataclass(frozen=True) -class MLAAttentionConfig: - """普通 MLA 的独立 latent 布局；单份 latent 固定在各 TP rank 完整复制。""" +class MLAConfig: + """普通 latent MLA 的逻辑配置，不包含 indexer 或 packed 物理字段。""" # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 @@ -29,8 +29,8 @@ class MLAAttentionConfig: logical_layer_groups: tuple[tuple[int, ...], ...] # 编码前的完整 latent 宽度，包含模型需要的 RoPE 部分，不含 indexer 或 padding。 latent_dim: int - # 未量化 latent 的浮点精度，也作为 MLA 工厂选择 packed 编码前的原始精度。 - dtype: torch.dtype + # 未压缩 latent 在运行态和存储中的浮点精度。 + latent_dtype: torch.dtype # 所属模型 TP 组内的 rank；各 rank 保存相同 latent，仅 rank 0 写出共享外部页。 tp_rank: int = 0 # 保存完整 latent 副本的 TP rank 数，不改变 latent 的维度或物理行宽。 @@ -55,7 +55,7 @@ class MLAFullAttentionState(FullAttentionState): def __init__( self, - layout_config: MLAAttentionConfig, + layout_config: MLAConfig, page_size: int = 1, device: str | torch.device = "cuda", ) -> None: @@ -75,7 +75,7 @@ def __init__( def get_state_type(self) -> str: """返回 MLA 压缩 KV 精度对应的存储方案标记。""" - return "mla_" + dtype_name(dtype=self.config.dtype) + return "mla_" + dtype_name(dtype=self.config.latent_dtype) def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" @@ -87,14 +87,14 @@ def allocate_storage(self, token_capacity: int) -> None: 1, self.config.latent_dim, ), - dtype=self.config.dtype, + dtype=self.config.latent_dtype, device=self.device, ) def get_token_bytes(self) -> int: """由配置计算当前 TP rank 的单 token latent 字节数，无需分配存储。""" config = self.config - return len(config.logical_layer_groups) * config.latent_dim * config.dtype.itemsize + return len(config.logical_layer_groups) * config.latent_dim * config.latent_dtype.itemsize # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 @@ -280,8 +280,8 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: TransferRegion( field="kv", byte_start=0, - byte_end=prod(shape) * config.dtype.itemsize, - dtype=config.dtype, + byte_end=prod(shape) * config.latent_dtype.itemsize, + dtype=config.latent_dtype, axes=("layer", "token", "head", "dim"), shape=shape, axis_mappings=( diff --git a/lightllm/common/attention_state/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla_state.py index a596e5a97c..39c54c8963 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla_state.py @@ -1,7 +1,7 @@ from __future__ import annotations from functools import lru_cache -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Literal from dataclasses import dataclass from math import prod @@ -24,33 +24,92 @@ @dataclass(frozen=True) -class PackedMLAAttentionConfig: - """group128/FP32 scale 或浮点 packed MLA，固定单份 latent 并在各 TP rank 复制。""" +class PackedMLALogicalConfig: + """带 indexer 的 packed MLA 逻辑配置；不描述实际字节布局。""" - # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 - # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 - # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 logical_layer_groups: tuple[tuple[int, ...], ...] - # 完整逻辑 latent 宽度，不含 indexer、scale 和物理 padding；FP8 为 576。 + # 编码前的完整 latent 宽度；具体分块、scale 和 padding 由工厂转换。 latent_dim: int - # KV 编码精度：FP16、BF16 或 float8_e4m3fn；FP8 仅量化非 RoPE 部分。 - # packed FP8 的底层 buffer 使用 uint8 承载 KV、FP32 scale 和 BF16 RoPE，不作为配置输入。 - dtype: torch.dtype - # indexer key 的 FP8 维度，另保存一个 FP32 scale，不计入此维度。 + # packed MLA 的逻辑格式；所有选项都包含 indexer。 + format: Literal["packed_fp16", "packed_bf16", "packed_group128_fp8", "packed_group64_e8m0"] + # indexer key 的逻辑宽度。 indexer_dim: int = 128 - # 0 为未量化；512 表示前 512 维按 group128 编码，其余 64 维以 BF16 保存。 - fp8_nope_dim: int = 0 - # FP8 非 RoPE 特征的量化分组宽度，当前编码固定为 128。 - quant_group_size: int = 128 - # 所属模型 TP 组内的 rank；各 rank 保存相同 latent，仅 rank 0 写出共享外部页。 + # latent 中保留原始 RoPE 精度的宽度；NoPE 宽度由 latent_dim - rope_dim 得到。 + rope_dim: int = 64 tp_rank: int = 0 - # 保存完整 latent 副本的 TP rank 数，不改变 latent 的维度或物理行宽。 tp_world_size: int = 1 + def __post_init__(self) -> None: + assert self.logical_layer_groups and all(self.logical_layer_groups) + layers = [layer for group in self.logical_layer_groups for layer in group] + assert len(layers) == len(set(layers)), "logical layers cannot repeat across storage groups" + assert ( + self.latent_dim > 0 + and 0 <= self.rope_dim <= self.latent_dim + and self.indexer_dim > 0 + and 0 <= self.tp_rank < self.tp_world_size + ) + assert self.format in ( + "packed_fp16", "packed_bf16", "packed_group128_fp8", "packed_group64_e8m0" + ), "unknown packed MLA logical format" + + +@dataclass(frozen=True) +class PackedMLAAttentionConfig: + """由 packed 逻辑配置解析出的物理布局，不重复保存逻辑字段。""" + + logical_config: PackedMLALogicalConfig + quant_group_size: int = 128 + + @property + def logical_layer_groups(self): + return self.logical_config.logical_layer_groups + + @property + def latent_dim(self): + return self.logical_config.latent_dim + + @property + def kv_encoding(self): + return { + "packed_fp16": "fp16", + "packed_bf16": "bf16", + "packed_group128_fp8": "fp8_group128", + }[self.logical_config.format] + + @property + def indexer_dim(self): + return self.logical_config.indexer_dim + + @property + def tp_rank(self): + return self.logical_config.tp_rank + + @property + def tp_world_size(self): + return self.logical_config.tp_world_size + + @property + def rope_dim(self): + return self.logical_config.rope_dim + + @property + def dtype(self) -> torch.dtype: + """返回编码对应的逻辑存储 dtype；FP8 的物理承载仍由 State 内部使用 uint8。""" + if self.kv_encoding == "fp16": + return torch.float16 + if self.kv_encoding == "bf16": + return torch.bfloat16 + return torch.float8_e4m3fn + @property def kv_bytes(self) -> int: - """单行 KV 区域字节数，FP8 包含 512 字节数据、16 字节 scales 和 64 维 BF16。""" - return 656 if self.dtype == torch.float8_e4m3fn else self.latent_dim * self.dtype.itemsize + """按逻辑 NoPE/RoPE 宽度和 scale 数量计算单行 KV 字节数。""" + if self.kv_encoding == "fp8_group128": + nope_dim = self.latent_dim - self.rope_dim + scale_count = nope_dim // self.quant_group_size + return nope_dim + scale_count * torch.float32.itemsize + self.rope_dim * torch.bfloat16.itemsize + return self.latent_dim * self.dtype.itemsize @property def token_bytes(self) -> int: @@ -60,24 +119,19 @@ def token_bytes(self) -> int: @property def indexer_offset(self) -> int: """保持既有格式：FP8 紧接 KV，浮点格式将对齐间隙放在 KV 与 indexer 之间。""" - return self.kv_bytes if self.dtype == torch.float8_e4m3fn else self.token_bytes - self.indexer_dim - 4 + return self.kv_bytes if self.kv_encoding == "fp8_group128" else self.token_bytes - self.indexer_dim - 4 def __post_init__(self) -> None: - assert isinstance(self.logical_layer_groups, tuple) and all( - isinstance(group, tuple) and group for group in self.logical_layer_groups - ), "layer groups must be a tuple of nonempty tuples" - layers = [layer for group in self.logical_layer_groups for layer in group] - assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" - assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" - assert self.latent_dim > 0 and self.indexer_dim > 0 and 0 <= self.tp_rank < self.tp_world_size - if self.dtype == torch.float8_e4m3fn: - assert (self.fp8_nope_dim, self.latent_dim, self.quant_group_size) == (512, 576, 128), ( - "FP8 MLA kernel requires 512 latent dimensions, 64 RoPE dimensions and groups of 128" - ) + assert self.kv_encoding in ("fp16", "bf16", "fp8_group128"), "unknown packed MLA KV encoding" + if self.kv_encoding == "fp8_group128": + assert ( + self.dtype == torch.float8_e4m3fn + and self.rope_dim == 64 + and self.latent_dim == 576 + and self.quant_group_size == 128 + ), "group128 packed MLA requires latent 576, rope 64 and group size 128" else: - assert self.dtype in (torch.float16, torch.bfloat16) and self.fp8_nope_dim == 0, ( - "unquantized packed MLA requires FP16 or BF16 without FP8 dimensions" - ) + assert self.kv_encoding in ("fp16", "bf16") @dataclass(frozen=True) @@ -118,7 +172,7 @@ def get_state_type(self) -> str: """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" encoding = ( f"fp8_group{self.config.quant_group_size}" - if self.config.dtype == torch.float8_e4m3fn else dtype_name(dtype=self.config.dtype) + if self.config.kv_encoding == "fp8_group128" else dtype_name(dtype=self.config.dtype) ) return "mla_packed_" + encoding @@ -126,7 +180,7 @@ def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity # FP8 行混合保存多种字段，用字节 buffer 承载，字段视图再解释各自的 dtype。 - storage_dtype = torch.uint8 if self.config.dtype == torch.float8_e4m3fn else self.config.dtype + storage_dtype = torch.uint8 if self.config.kv_encoding == "fp8_group128" else self.config.dtype self.storage = torch.empty( size=( len(self.config.logical_layer_groups), @@ -160,7 +214,7 @@ def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: raw = row.view(torch.uint8) kv = ( raw[..., : self.config.kv_bytes] - if self.config.dtype == torch.float8_e4m3fn else row[..., : self.config.latent_dim] + if self.config.kv_encoding == "fp8_group128" else row[..., : self.config.latent_dim] ) indexer = raw[..., self.config.indexer_offset : self.config.indexer_offset + self.config.indexer_dim + 4] return PackedMLALayerView(kv=kv, indexer_k=indexer) @@ -172,19 +226,24 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, if not len(token_indices): return assert self.storage.is_cuda, "KV writes require CUDA storage" - if self.config.dtype == torch.float8_e4m3fn: + if self.config.kv_encoding == "fp8_group128": from lightllm.common.basemodel.triton_kernel.kv_copy.destindex_copy_kv_flashmla_fp8 import ( destindex_copy_kv_flashmla_fp8, ) raw = self.storage[storage_layer_index] + nope_dim = self.config.latent_dim - self.config.rope_dim + scale_bytes = nope_dim // self.config.quant_group_size * torch.float32.itemsize + rope_start = nope_dim + scale_bytes destindex_copy_kv_flashmla_fp8( - KV_nope=kv[..., :512], - KV_rope=kv[..., 512:], + KV_nope=kv[..., :nope_dim], + KV_rope=kv[..., nope_dim:], DestLoc=token_indices, - O_nope=raw[..., :512].view(torch.float8_e4m3fn), - O_scale=raw[..., 512:528].view(torch.float32), - O_rope=raw[..., 528:656].view(torch.bfloat16), + O_nope=raw[..., :nope_dim].view(torch.float8_e4m3fn), + O_scale=raw[..., nope_dim:rope_start].view(torch.float32), + O_rope=raw[..., rope_start : rope_start + self.config.rope_dim * torch.bfloat16.itemsize].view( + torch.bfloat16 + ), ) else: from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv @@ -284,11 +343,11 @@ def copy_tokens_from_dp( ) -> None: """校验 MLA 编码与区域尺寸，再复制完整 KV 和 indexer 数据。""" assert ( - remote_dp_full_att_state.config.fp8_nope_dim, + remote_dp_full_att_state.config.kv_encoding, remote_dp_full_att_state.config.latent_dim, remote_dp_full_att_state.config.indexer_dim, ) == ( - self.config.fp8_nope_dim, + self.config.kv_encoding, self.config.latent_dim, self.config.indexer_dim, ), "incompatible packed MLA format" @@ -374,7 +433,7 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: page_tokens, ) # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 - if config.dtype == torch.float8_e4m3fn: + if config.kv_encoding == "fp8_group128": field_specs = [ ("kv_nope", torch.uint8, 512, "dim"), ("kv_scale", torch.float32, 4, "group"), diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 047c619dfa..2521054b4a 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -69,7 +69,7 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" - from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import MLAConfig, create_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -77,13 +77,13 @@ def create_full_att_states( if not logical_layer_groups: return [] state = create_mla_state( - config=MLAAttentionConfig( + config=MLAConfig( logical_layer_groups=logical_layer_groups, latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], - dtype=get_torch_dtype(args.data_type), + latent_dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + page_size=args.page_size, device="meta", ) return [state] diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 105d735816..ca050e911b 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -1,4 +1,5 @@ from typing import TYPE_CHECKING +import torch from lightllm.models.deepseek2.model import Deepseek2TpPartModel from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer @@ -25,21 +26,22 @@ def create_full_att_states( model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" - from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] - state = create_mla_state( - config=MLAAttentionConfig( + state = create_packed_mla_state( + config=PackedMLALogicalConfig( logical_layer_groups=logical_layer_groups, latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + format="packed_group128_fp8" if args.llm_kv_type == "fp8kv_dsa" else "packed_bf16", + indexer_dim=config.get("indexer_dim", 128), ), - kv_type=args.llm_kv_type, packed=True, page_size=args.page_size, device="meta", + page_size=args.page_size, device="meta", ) return [state] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index bd1b745bc0..67ade58148 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -145,7 +145,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -153,17 +153,16 @@ def create_full_att_states( logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) if not logical_layer_groups: return [] - state = create_mla_state( - config=MLAAttentionConfig( + state = create_packed_mla_state( + config=PackedMLALogicalConfig( logical_layer_groups=logical_layer_groups, latent_dim=cfg.full_att_head_dim, - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + format="packed_group64_e8m0" if args.llm_kv_type == "fp8kv_dsa" else ( + "packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16" + ), ), - kv_type=args.llm_kv_type, - packed=True, - fp8_format="group64_e8m0", page_size=args.page_size, device="meta", ) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index fe18d2a888..37663e4fb9 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING from copy import deepcopy +import torch from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend @@ -81,24 +82,23 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" - from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) if not logical_layer_groups: return [] - state = create_mla_state( - config=MLAAttentionConfig( + state = create_packed_mla_state( + config=PackedMLALogicalConfig( logical_layer_groups=logical_layer_groups, latent_dim=config["kv_lora_rank"], - dtype=get_torch_dtype(args.data_type), tp_rank=tp_rank, tp_world_size=tp_world_size, + format="packed_group64_e8m0" if args.llm_kv_type == "fp8kv_dsa" else ( + "packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16" + ), ), - kv_type=args.llm_kv_type, - packed=True, - fp8_format="group64_e8m0", page_size=args.page_size, device="meta", ) diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index 61320d1cef..409f27fe80 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -7,7 +7,7 @@ from functools import lru_cache from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.full_att import MLAAttentionConfig, MLAFullAttentionState +from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState from lightllm.common.attention_state.bounded_att import GDNState from unit_tests.common.att_resources.test_full_state import full_state from unit_tests.common.att_resources.test_bounded_state import make_state @@ -30,10 +30,10 @@ def __init__(self, device: str | torch.device | int = "cpu"): self.register_state("draft0", draft_state) self.register_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( - layout_config=MLAAttentionConfig( + layout_config=MLAConfig( logical_layer_groups=self.main_gqa_fp32_state.config.logical_layer_groups, latent_dim=self.main_gqa_fp32_state.config.head_dim, - dtype=self.main_gqa_fp32_state.config.dtype, + latent_dtype=self.main_gqa_fp32_state.config.dtype, ), page_size=self.main_gqa_fp32_state.page_size, device=device, ) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 94d3daba49..ec2cf24cb0 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -11,7 +11,7 @@ from lightllm.common.attention_state.types import TensorField from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState from lightllm.common.attention_state.full_att import ( - FullAttentionState, GQAAttentionConfig, MLAAttentionConfig, + FullAttentionState, GQAAttentionConfig, MLAConfig, GroupQuantizedAttentionConfig, StaticFP8AttentionConfig, PackedMLAAttentionConfig, CompressedAttentionConfig, FP8E8M0PackedMLAConfig, StaticFP8FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, @@ -56,7 +56,7 @@ def bind_test_fp8_calibration( def rebuild_full_state( template: FullAttentionState, config: ( - GQAAttentionConfig | MLAAttentionConfig | GroupQuantizedAttentionConfig + GQAAttentionConfig | MLAConfig | GroupQuantizedAttentionConfig | StaticFP8AttentionConfig | PackedMLAAttentionConfig | CompressedAttentionConfig | FP8E8M0PackedMLAConfig | None ) = None, diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index fc2ec7faeb..92f70a2ab6 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -28,7 +28,7 @@ FullAttentionState, GQAFullAttentionState, GQAAttentionConfig, - MLAAttentionConfig, + MLAConfig, create_gqa_state, create_mla_state, GroupQuantizedAttentionConfig, @@ -43,15 +43,21 @@ def definition(kind: str): - if kind == "mla" or kind.startswith("packed"): + if kind == "mla": return create_mla_state( - config=MLAAttentionConfig( + config=MLAConfig(logical_layer_groups=((0,), (2,)), latent_dim=576, latent_dtype=torch.bfloat16), + page_size=4, + device="meta", + ) + if kind.startswith("packed"): + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + + return create_packed_mla_state( + config=PackedMLALogicalConfig( logical_layer_groups=((0,), (2,)), latent_dim=512 if kind == "packed_glm" else 576, - dtype=torch.bfloat16, + format="packed_group128_fp8" if kind == "packed_fp8" else "packed_bf16", ), - kv_type="fp8kv_dsa" if kind == "packed_fp8" else "None", - packed=kind.startswith("packed"), page_size=4, device="meta", ) @@ -291,8 +297,8 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( (GQAFullAttentionState, "main", 1, torch.float32), ): if state_cls is MLAFullAttentionState: - config = MLAAttentionConfig( - logical_layer_groups=((0,),), latent_dim=dim, dtype=dtype, + config = MLAConfig( + logical_layer_groups=((0,),), latent_dim=dim, latent_dtype=dtype, ) else: config = GQAAttentionConfig( @@ -642,7 +648,17 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( ): cpu_args.cpu_cache_token_page_size = page_tokens template = definition(kind) - config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) + if hasattr(template.config, "logical_config"): + config = replace( + template.config, + logical_config=replace( + template.config.logical_config, + tp_world_size=tp_world_size, + tp_rank=tp_world_size - 1, + ), + ) + else: + config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) scales = torch.ones((len(config.logical_layer_groups), 2 * config.head_num)) if kind == "static_head" else None q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None manager = AttentionStateManager(request_capacity=3) diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py index 36210e397d..745c9fcfa7 100644 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -12,10 +12,10 @@ from lightllm.common.attention_state.full_att import ( FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState, - MLAAttentionConfig, + PackedMLALogicalConfig, FullAttentionState, PackedMLAFullAttentionState, - create_mla_state, + create_packed_mla_state, ) from lightllm.common.attention_state.types import StateTransferLayout @@ -25,19 +25,16 @@ def make_state(device: str = "meta", capacity: int = 16, tp_rank: int = 0) -> FP8E8M0PackedMLAState: - state = create_mla_state( - config=MLAAttentionConfig( + state = create_packed_mla_state( + config=PackedMLALogicalConfig( logical_layer_groups=((0, 4), (8,)), latent_dim=512, - dtype=torch.bfloat16, + format="packed_group64_e8m0", tp_rank=tp_rank, tp_world_size=2, ), - kv_type="fp8kv_dsa", - packed=True, page_size=4, device=device, - fp8_format="group64_e8m0", ) state.allocate_storage(token_capacity=capacity) return state @@ -62,23 +59,20 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic assert region.shape == (2, 4, 1, 716) and region.byte_end == 2 * 4 * 716 # 默认仍是原来的 group128/FP32 编码，不能再通过 512 维自动猜测 GLM 格式。 - with pytest.raises(AssertionError, match="FP8 MLA kernel requires"): - create_mla_state( - config=MLAAttentionConfig( - logical_layer_groups=((0,),), latent_dim=512, dtype=torch.bfloat16 + with pytest.raises(AssertionError, match="latent 576"): + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), latent_dim=512, + format="packed_group128_fp8", ), - kv_type="fp8kv_dsa", - packed=True, device="meta", ) with pytest.raises(AssertionError, match="512-dimensional"): - create_mla_state( - config=MLAAttentionConfig( - logical_layer_groups=((0,),), latent_dim=576, dtype=torch.bfloat16 + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), latent_dim=576, + format="packed_group64_e8m0", ), - kv_type="fp8kv_dsa", - packed=True, - fp8_format="group64_e8m0", device="meta", ) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 5e5f768f7d..97105484be 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -17,7 +17,7 @@ GQAFullAttentionState, MLAFullAttentionState, GQAAttentionConfig, - MLAAttentionConfig, + MLAConfig, ) from lightllm.common.attention_state.types import ( AxisMapping, StateTransferLayout, TransferRegion, @@ -26,8 +26,8 @@ def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): if cls is MLAFullAttentionState: - config = MLAAttentionConfig( - logical_layer_groups=((2,), (5,)), latent_dim=4, dtype=torch.float32, + config = MLAConfig( + logical_layer_groups=((2,), (5,)), latent_dim=4, latent_dtype=torch.float32, ) else: config = GQAAttentionConfig( diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 6046e1475f..596682785d 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -81,7 +81,8 @@ def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "E ) # 共享消费者 14/15 分别使用最后一个同类型生产层 12/13，保持四份物理层存储。 assert state.config.logical_layer_groups == ((10,), (11,), (12, 14), (13, 15)) - assert state.get_token_bytes() == 4 * 2 * 2 * 128 * state.config.dtype.itemsize + latent_dtype = state.config.latent_dtype if hasattr(state.config, "latent_dtype") else state.config.dtype + assert state.get_token_bytes() == 4 * 2 * 2 * 128 * latent_dtype.itemsize @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 04c82bfdd6..8c0dd81779 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -19,7 +19,7 @@ GQAFullAttentionState, GQAAttentionConfig, ) -from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig +from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig, PackedMLALogicalConfig from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -58,11 +58,13 @@ def make_resources( if glm: full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), - dtype=torch.bfloat16, - tp_rank=rank, - tp_world_size=tp, - latent_dim=512, + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), + format="packed_bf16", + tp_rank=rank, + tp_world_size=tp, + latent_dim=512, + ) ), page_size=4, ) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 9345913fc7..664db98fbb 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -20,7 +20,7 @@ GroupQuantizedFullAttentionState, StaticFP8AttentionConfig, StaticFP8FullAttentionState, - PackedMLAAttentionConfig, + PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState, ) from lightllm.common.attention_state.types import ( @@ -70,10 +70,11 @@ def make_quant_state(kind: str, device: str | torch.device | int): return prepared_full_state fp8 = kind == "packed_fp8" config = PackedMLAAttentionConfig( - latent_dim=576, - logical_layer_groups=((1,), (3,)), - dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, - fp8_nope_dim=512 if fp8 else 0, + logical_config=PackedMLALogicalConfig( + latent_dim=576, + logical_layer_groups=((1,), (3,)), + format="packed_group128_fp8" if fp8 else "packed_bf16", + ) ) prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) prepared_full_state.allocate_storage(8) diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index b139b092ca..3339a71343 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -10,9 +10,9 @@ import torch from lightllm.common.attention_state.full_att import ( - GQAAttentionConfig, MLAAttentionConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, + GQAAttentionConfig, MLAConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, StaticFP8FullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, - create_gqa_state, create_mla_state, + PackedMLALogicalConfig, create_gqa_state, create_mla_state, create_packed_mla_state, ) from lightllm.common.attention_state.bounded_att import ( BoundedAttentionState, ConvSSMStateConfig, create_gdn_state, create_kda_state, @@ -97,23 +97,28 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( row_width: int, expected_type: str, ): - config = MLAAttentionConfig( - logical_layer_groups=((2,), (6,)), - latent_dim=kv_dim, - dtype=torch.bfloat16, - tp_rank=1, - tp_world_size=4, - ) - state = create_mla_state( - config=config, - packed=packed, - kv_type=kv_type, - page_size=4, - device="meta", - ) + if packed: + config = PackedMLALogicalConfig( + logical_layer_groups=((2,), (6,)), + latent_dim=kv_dim, + format=("packed_group128_fp8" if kv_type == "fp8kv_dsa" else "packed_bf16"), + tp_rank=1, + tp_world_size=4, + ) + state = create_packed_mla_state(config=config, page_size=4, device="meta") + else: + config = MLAConfig( + logical_layer_groups=((2,), (6,)), latent_dim=kv_dim, latent_dtype=torch.bfloat16, + tp_rank=1, tp_world_size=4, + ) + state = create_mla_state(config=config, page_size=4, device="meta") assert state.get_state_type() == expected_type assert state.config.latent_dim == config.latent_dim == kv_dim - assert state.config.dtype == (torch.float8_e4m3fn if kv_type == "fp8kv_dsa" else config.dtype) + expected_dtype = torch.float8_e4m3fn if kv_type == "fp8kv_dsa" else torch.bfloat16 + if packed: + assert state.config.dtype == expected_dtype + else: + assert state.config.latent_dtype == expected_dtype assert not hasattr(state.config, "storage_dtype") assert not hasattr(state.config, "replicated") assert state.capacity == 0 and read_kv_storage(state) is None @@ -122,7 +127,7 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( state.device = torch.device("cpu") state.allocate_storage(8) assert read_kv_storage(state).shape == (2, 12, 1, row_width) - assert read_kv_storage(state).dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else config.dtype) + assert read_kv_storage(state).dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else torch.bfloat16) if packed: # 保持旧算子读取的物理布局：FP8 indexer 紧接 KV，浮点格式保留中间的对齐间隙。 expected_indexer_offset = 656 if kv_type == "fp8kv_dsa" else kv_dim * 2 + 12 @@ -182,18 +187,16 @@ def test_create_rejects_encoding_from_another_logical_family(): ) with pytest.raises(AssertionError, match="unsupported GQA encoding"): create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") - config = MLAAttentionConfig( - logical_layer_groups=((0,),), latent_dim=576, dtype=torch.bfloat16, + packed_config = PackedMLALogicalConfig( + logical_layer_groups=((0,),), latent_dim=576, format="packed_group128_fp8", ) - with pytest.raises(AssertionError, match="unsupported packed MLA encoding"): - create_mla_state( - config=config, - packed=True, - kv_type="int4kv", + with pytest.raises(AssertionError, match="normal MLA creation"): + create_mla_state(config=packed_config, device="meta") + with pytest.raises(AssertionError, match="packed MLA creation"): + create_packed_mla_state( + config=MLAConfig(logical_layer_groups=((0,),), latent_dim=576, latent_dtype=torch.bfloat16), device="meta", ) - with pytest.raises(AssertionError, match="unsupported MLA encoding"): - create_mla_state(config=config, kv_type="fp8kv_dsa", device="meta") @pytest.mark.parametrize( diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index ebf7ef10c1..1db1627b07 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -19,7 +19,7 @@ from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.attention_state.full_att import ( GQAAttentionConfig, - MLAAttentionConfig, + MLAConfig, GQAFullAttentionState, MLAFullAttentionState, ) @@ -59,8 +59,8 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( pytest.skip("CUDA required") cpu_args.cpu_cache_token_page_size = 4 if state_cls is MLAFullAttentionState: - config = MLAAttentionConfig( - logical_layer_groups=((0,),), latent_dim=4, dtype=torch.float32, + config = MLAConfig( + logical_layer_groups=((0,),), latent_dim=4, latent_dtype=torch.float32, ) else: config = GQAAttentionConfig( diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 5d42ff63a5..94e4ab5e91 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -9,7 +9,7 @@ from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, - PackedMLAAttentionConfig, + PackedMLAAttentionConfig, PackedMLALogicalConfig, ) from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -87,10 +87,12 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( full_layer_ids += (config.all_layer_num,) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in full_layer_ids), - dtype=torch.bfloat16, - tp_world_size=tp_world_size, - latent_dim=512, + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), + format="packed_bf16", + tp_world_size=tp_world_size, + latent_dim=512, + ) ), page_size=args.page_size, ) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 2d1c6fca03..91fbf63576 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -1279,17 +1279,15 @@ def fake_context(*args): def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): from lightllm.common.attention_state import AttentionStateManager - from lightllm.common.attention_state.full_att import MLAAttentionConfig, create_mla_state + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) - config = MLAAttentionConfig( - logical_layer_groups=((3,),), latent_dim=512, dtype=torch.bfloat16, - ) - state = create_mla_state( - config=config, kv_type="fp8kv_dsa", packed=True, fp8_format="group64_e8m0", page_size=4, + config = PackedMLALogicalConfig( + logical_layer_groups=((3,),), latent_dim=512, format="packed_group64_e8m0", ) + state = create_packed_mla_state(config=config, page_size=4) state.allocate_storage(token_capacity=16) state.storage.zero_() manager = AttentionStateManager(request_capacity=1) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 218095216a..b163d0d972 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -12,7 +12,7 @@ from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, - PackedMLAAttentionConfig, + PackedMLAAttentionConfig, PackedMLALogicalConfig, ) from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -70,11 +70,13 @@ def create(tp, target_mtp_step=None): monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) full = PackedMLAFullAttentionState( PackedMLAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in full_layer_ids), - dtype=torch.bfloat16, - tp_rank=rank, - tp_world_size=tp, - latent_dim=512, + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), + format="packed_bf16", + tp_rank=rank, + tp_world_size=tp, + latent_dim=512, + ) ), page_size=args.page_size, ) From 038c4db6f6f19f0c6c74c970c40cc94fb7044c3a Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 04:40:16 +0000 Subject: [PATCH 36/56] refactor: organize attention states by type --- docs/cache_manager_design.md | 60 ++++++------------- .../attention_state/bounded_att/__init__.py | 20 +++---- .../bounded_att/gdn/__init__.py | 5 ++ .../bounded_att/{ => gdn}/create_gdn_state.py | 2 +- .../bounded_att/{ => gdn}/gdn_pages_pool.py | 4 +- .../bounded_att/{ => gdn}/gdn_state.py | 12 ++-- .../bounded_att/kda/__init__.py | 5 ++ .../bounded_att/{ => kda}/create_kda_state.py | 2 +- .../bounded_att/{ => kda}/kda_pages_pool.py | 4 +- .../bounded_att/{ => kda}/kda_state.py | 12 ++-- .../bounded_att/sliding_window/__init__.py | 4 ++ .../create_sliding_window_state.py | 0 .../fp8_e8m0_swa_state.py | 12 ++-- .../fp8_e8m0_swa_transfer.py | 0 .../bounded_att/window_compressor/__init__.py | 4 ++ .../create_window_compressor_state.py | 0 .../window_compressor_state.py | 12 ++-- .../window_compressor_transfer.py | 2 +- .../attention_state/full_att/__init__.py | 23 +++---- .../full_att/compressed_att/__init__.py | 9 +++ .../compressed_att_state.py | 10 ++-- .../create_compressed_att_state.py | 0 .../full_att/fp8_e8m0_packed_mla/__init__.py | 3 + .../fp8_e8m0_packed_mla_state.py | 14 ++--- .../fp8_e8m0_packed_mla_write.py | 0 .../attention_state/full_att/gqa/__init__.py | 4 ++ .../full_att/{ => gqa}/create_gqa_state.py | 4 +- .../full_att/{ => gqa}/gqa_state.py | 8 +-- .../full_att/{ => gqa}/gqa_transfer.py | 4 +- .../full_att/group_quantized/__init__.py | 3 + .../group_quantized_state.py | 8 +-- .../group_quantized_transfer.py | 4 +- .../attention_state/full_att/mla/__init__.py | 4 ++ .../full_att/{ => mla}/create_mla_state.py | 2 +- .../full_att/{ => mla}/mla_state.py | 10 ++-- .../full_att/packed_mla/__init__.py | 9 +++ .../create_packed_mla_state.py | 4 +- .../{ => packed_mla}/packed_mla_state.py | 8 +-- .../{ => packed_mla}/packed_mla_transfer.py | 4 +- .../{ => packed_mla}/packed_transfer.py | 4 +- .../full_att/static_fp8/__init__.py | 3 + .../{ => static_fp8}/static_fp8_state.py | 8 +-- .../{ => static_fp8}/static_fp8_transfer.py | 4 +- lightllm/common/attention_state/manager.py | 2 +- lightllm/models/deepseek_v4/config.py | 4 +- lightllm/models/deepseek_v4/infer_struct.py | 2 +- lightllm/models/deepseek_v4/model.py | 2 +- .../common/att_resources/test_dsv4_state.py | 2 +- .../test_fp8_e8m0_packed_mla_state.py | 2 +- .../common/att_resources/test_gqa_transfer.py | 2 +- .../att_resources/test_network_transport.py | 2 +- unit_tests/common/dsv4_test_utils.py | 2 +- unit_tests/models/glm5_next/test_fp8_kv.py | 4 +- 53 files changed, 185 insertions(+), 153 deletions(-) create mode 100644 lightllm/common/attention_state/bounded_att/gdn/__init__.py rename lightllm/common/attention_state/bounded_att/{ => gdn}/create_gdn_state.py (91%) rename lightllm/common/attention_state/bounded_att/{ => gdn}/gdn_pages_pool.py (92%) rename lightllm/common/attention_state/bounded_att/{ => gdn}/gdn_state.py (98%) create mode 100644 lightllm/common/attention_state/bounded_att/kda/__init__.py rename lightllm/common/attention_state/bounded_att/{ => kda}/create_kda_state.py (91%) rename lightllm/common/attention_state/bounded_att/{ => kda}/kda_pages_pool.py (93%) rename lightllm/common/attention_state/bounded_att/{ => kda}/kda_state.py (98%) create mode 100644 lightllm/common/attention_state/bounded_att/sliding_window/__init__.py rename lightllm/common/attention_state/bounded_att/{ => sliding_window}/create_sliding_window_state.py (100%) rename lightllm/common/attention_state/bounded_att/{ => sliding_window}/fp8_e8m0_swa_state.py (97%) rename lightllm/common/attention_state/bounded_att/{ => sliding_window}/fp8_e8m0_swa_transfer.py (100%) create mode 100644 lightllm/common/attention_state/bounded_att/window_compressor/__init__.py rename lightllm/common/attention_state/bounded_att/{ => window_compressor}/create_window_compressor_state.py (100%) rename lightllm/common/attention_state/bounded_att/{ => window_compressor}/window_compressor_state.py (98%) rename lightllm/common/attention_state/bounded_att/{ => window_compressor}/window_compressor_transfer.py (99%) create mode 100644 lightllm/common/attention_state/full_att/compressed_att/__init__.py rename lightllm/common/attention_state/full_att/{ => compressed_att}/compressed_att_state.py (98%) rename lightllm/common/attention_state/full_att/{ => compressed_att}/create_compressed_att_state.py (100%) create mode 100644 lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py rename lightllm/common/attention_state/full_att/{ => fp8_e8m0_packed_mla}/fp8_e8m0_packed_mla_state.py (97%) rename lightllm/common/attention_state/full_att/{ => fp8_e8m0_packed_mla}/fp8_e8m0_packed_mla_write.py (100%) create mode 100644 lightllm/common/attention_state/full_att/gqa/__init__.py rename lightllm/common/attention_state/full_att/{ => gqa}/create_gqa_state.py (91%) rename lightllm/common/attention_state/full_att/{ => gqa}/gqa_state.py (99%) rename lightllm/common/attention_state/full_att/{ => gqa}/gqa_transfer.py (99%) create mode 100644 lightllm/common/attention_state/full_att/group_quantized/__init__.py rename lightllm/common/attention_state/full_att/{ => group_quantized}/group_quantized_state.py (99%) rename lightllm/common/attention_state/full_att/{ => group_quantized}/group_quantized_transfer.py (99%) create mode 100644 lightllm/common/attention_state/full_att/mla/__init__.py rename lightllm/common/attention_state/full_att/{ => mla}/create_mla_state.py (92%) rename lightllm/common/attention_state/full_att/{ => mla}/mla_state.py (98%) create mode 100644 lightllm/common/attention_state/full_att/packed_mla/__init__.py rename lightllm/common/attention_state/full_att/{ => packed_mla}/create_packed_mla_state.py (89%) rename lightllm/common/attention_state/full_att/{ => packed_mla}/packed_mla_state.py (99%) rename lightllm/common/attention_state/full_att/{ => packed_mla}/packed_mla_transfer.py (99%) rename lightllm/common/attention_state/full_att/{ => packed_mla}/packed_transfer.py (99%) create mode 100644 lightllm/common/attention_state/full_att/static_fp8/__init__.py rename lightllm/common/attention_state/full_att/{ => static_fp8}/static_fp8_state.py (99%) rename lightllm/common/attention_state/full_att/{ => static_fp8}/static_fp8_transfer.py (99%) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 90a3b76e0d..9684bb5bfc 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -310,23 +310,17 @@ lightllm/common/ │ ├── full_att/ │ │ ├── __init__.py │ │ ├── base_state.py # FullAttentionState：抽象行为接口 -│ │ ├── create_gqa_state.py # GQA 统一创建入口：选择普通精度、INT4/INT8 或静态 FP8 -│ │ ├── gqa_state.py # GQAFullAttentionState：兼容 MHA / GQA / MQA -│ │ ├── gqa_transfer.py # GQA CPU/PD 共用的 K/V 页传输入口与 kernel +│ │ ├── gqa/ # GQA 配置、State、创建入口和传输实现 +│ │ │ ├── create_gqa_state.py +│ │ │ ├── gqa_state.py +│ │ │ └── gqa_transfer.py │ │ ├── dp_transfer.py # 各 full-attention 方案的直接索引 DP 聚合复制 -│ │ ├── group_quantized_state.py # GroupQuantizedFullAttentionState：INT4 / INT8 等 -│ │ ├── group_quantized_transfer.py # INT4/INT8 KV 与动态 scale 的定制传输 -│ │ ├── static_fp8_state.py # StaticFP8FullAttentionState -│ │ ├── static_fp8_transfer.py # FP8 K/V 原始编码的定制传输 -│ │ ├── create_mla_state.py # MLA 统一创建入口：选择普通、packed 或 packed FP8 -│ │ ├── mla_state.py # MLAFullAttentionState：压缩 latent 存储 -│ │ ├── create_compressed_att_state.py # 序列压缩历史 KV 的独立创建入口 -│ │ ├── compressed_att_state.py # 独立配置及 V4 C4/C128 packed 压缩块存储 -│ │ ├── packed_transfer.py # 单字段 packed KV 的定制传输 -│ │ ├── packed_mla_state.py # 普通精度或 group128/FP32-scale packed MLA -│ │ ├── fp8_e8m0_packed_mla_state.py # 独立 FP8E8M0PackedMLAConfig / State：group64/E8M0 编码 -│ │ ├── fp8_e8m0_packed_mla_write.py # group64/E8M0 定制量化写入，不依赖具体模型 -│ │ └── packed_mla_transfer.py # packed KV、RoPE 与 indexer 字段的定制传输 +│ │ ├── group_quantized/ # INT4/INT8 配置、State 和动态 scale 传输 +│ │ ├── static_fp8/ # 静态 FP8 配置、State 和传输 +│ │ ├── mla/ # 普通 MLA latent 配置、State 和创建入口 +│ │ ├── compressed_att/ # V4 C4/C128 压缩块配置、State 和创建入口 +│ │ ├── packed_mla/ # packed MLA 配置、State、创建入口和传输 +│ │ └── fp8_e8m0_packed_mla/ # group64/E8M0 packed MLA 配置、State 和写入 kernel │ ├── bounded_att/ │ │ ├── __init__.py │ │ ├── base_state.py # BoundedAttentionState：有界状态的抽象操作契约 @@ -334,32 +328,16 @@ lightllm/common/ │ │ ├── transfer_utils.py # bounded_att_buffer_copies：生成快照字段的复制视图 │ │ ├── copy_kernel.py # copy_bounded_att_tensors：聚合多页、多字段的 Triton 拷贝 │ │ ├── conv_ssm_types.py # KDA/GDN 的 Conv/SSM 布局配置 -│ │ ├── create_sliding_window_state.py # 按 SWA 算法命名的创建入口，当前返回 FP8/E8M0 实现 -│ │ ├── fp8_e8m0_swa_state.py # FP8E8M0SlidingWindowConfig / State：仅管理近期 KV -│ │ ├── create_window_compressor_state.py # 窗口压缩器原子运行态创建入口 -│ │ ├── window_compressor_state.py # WindowCompressorStateConfig / State：候选与 gate -│ │ ├── window_pages_pool.py # SWA 与压缩器的时间有序窗口快照池 -│ │ ├── fp8_e8m0_swa_transfer.py # FP8/E8M0 packed SWA 的 GPU 窗口搬运 -│ │ ├── window_compressor_transfer.py # 候选/gate ring 与时间有序快照的 GPU 搬运 -│ │ ├── create_kda_state.py # KDA 统一创建入口 -│ │ ├── kda_state.py # KDAState：Conv / SSM / MTP 与页操作 -│ │ ├── kda_pages_pool.py # KDAPagesPool：KDA Conv 历史和方阵递推快照 -│ │ ├── create_gdn_state.py # GDN 统一创建入口 -│ │ ├── gdn_state.py # GDNState：独立算法状态与页操作 -│ │ └── gdn_pages_pool.py # GDNPagesPool：支持不同 K/V head 数和维度的快照 -└── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 - -lightllm/models// -├── model.py # 两类 State 创建接口，决定本模型的资源组合 -└── config.py # 模型专有配置与派生参数 - -lightllm/utils/ -└── cpu_cache_layout.py # build_cpu_cache_layouts：构建 CPU cache 联合页布局 +│ │ ├── gdn/ # GDN 配置、State、页池和创建入口 +│ │ ├── kda/ # KDA 配置、State、页池和创建入口 +│ │ ├── sliding_window/ # SWA 配置、State、创建入口和传输 +│ │ ├── window_compressor/ # 窗口压缩器配置、State、创建入口和传输 +│ │ └── window_pages_pool.py # SWA 与压缩器共用的时间有序窗口快照池 ``` `allocators/` 与两类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 -`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口，各具体配置与对应 State 放在同一文件中；`bounded_att/base_state.py` 定义有界运行态的公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 +`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口；每种 full-attention 子 State 使用独立目录，目录内放置配置、State、创建入口及专用传输实现。`bounded_att/base_state.py` 定义有界运行态的公共接口；每种 bounded 子 State 同样使用独立目录，公共的页池工具和聚合 kernel 留在分类目录根部。布局或算法差异由各子目录自己的 `*_state.py` 实现，不设置另一套独立的 storage 公共继承体系。 `KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。K-pool 使用 W4/S4 的 `WindowCompressorState`，在整组边界无重叠历史，不创建大小页 pool。 @@ -385,7 +363,7 @@ KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态 CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；无重叠的窗口 State 在 CPU 对齐边界不需要快照；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`bounded_att/window_compressor_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运外部快照所需的逻辑窗口；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 的传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；窗口压缩器直接提交环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在各方案中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。压缩器使用固定 `window_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；窗口搬运检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 +`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；无重叠的窗口 State 在 CPU 对齐边界不需要快照；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`bounded_att/window_compressor/window_compressor_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运外部快照所需的逻辑窗口；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 的传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；窗口压缩器直接提交环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在各方案中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。压缩器使用固定 `window_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；窗口搬运检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 @@ -545,7 +523,7 @@ FullAttentionState 所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。配置也不建立继承关系：普通 GQA 使用 `GQAAttentionConfig`；分组量化使用 `GroupQuantizedAttentionConfig`，以 `scale_dtype` 明确 scale 精度；静态 FP8 使用 `StaticFP8AttentionConfig`，不要求调用方传入固定的物理 dtype。普通 MLA 使用 `MLAConfig`，packed MLA 使用带 indexer 的 `PackedMLALogicalConfig`；工厂再转换为具体 State 的物理配置。具体 packed 配置使用 `kv_encoding` 显式选择 `fp16`、`bf16` 或 `fp8_group128`，`dtype` 由编码提供只读 property。FP8 行的 uint8 承载类型只在存储分配函数中确定，不提供 `storage_dtype` 配置或属性；物理行宽及 indexer 偏移由编码推导。MLA 固定单份 latent 并在 TP rank 间复制，不提供 GQA 的 head 数或 `replicated` 开关。 -`gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 +`full_att/gqa/gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 该 State 只保存 KV，`config.head_num` 表示当前 TP rank 切分或复制后持有的 KV 头数。K 和 V 各自独立分配一个 `[layer, token_slot, local_kv_heads, head_dim]` buffer，不再保存 `heads_per_kv`。写入接口的输入仍按 K heads、V heads 前后排列，State 将两部分分别写入各自存储；计算 backend 直接取得 K/V 视图。Query 头到 KV 头的分组关系由 attention backend 处理，因此三者共用同一套 State 接口与存储实现。 @@ -1666,7 +1644,7 @@ RoPE 和量化仍由各模型自己的计算算子完成。资源层只负责窗 PD 在任意前缀长度导出完整 `window_size` 逻辑窗口，例如 K-pool 的长度 6 前缀仍需 携带未完成池化组的原始 key/gate。P/D 的 MTP 容量不必相同：搬运按绝对前缀长度映射 到各自的 ring，不传输或清理逻辑窗口外的候选槽位。环形搬运 kernel 位于 -`bounded_att/window_compressor_transfer.py`，一次 launch 搬运全部层的候选与 gate。 +`bounded_att/window_compressor/window_compressor_transfer.py`，一次 launch 搬运全部层的候选与 gate。 MTP 回退依赖保留的历史和绝对位置，不调用缓存恢复入口。 ## 7. 资源初始化与调用约定 diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py index 8185379fd5..55190bea6d 100644 --- a/lightllm/common/attention_state/bounded_att/__init__.py +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -1,16 +1,16 @@ -from .create_gdn_state import create_gdn_state -from .create_kda_state import create_kda_state +from .gdn.create_gdn_state import create_gdn_state +from .kda.create_kda_state import create_kda_state from .conv_ssm_types import ConvSSMStateConfig -from .gdn_pages_pool import GDNPagesPool -from .gdn_state import GDNState -from .kda_pages_pool import KDAPagesPool -from .kda_state import KDAState +from .gdn.gdn_pages_pool import GDNPagesPool +from .gdn.gdn_state import GDNState +from .kda.kda_pages_pool import KDAPagesPool +from .kda.kda_state import KDAState from .base_state import BoundedAttentionState -from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState -from .window_compressor_state import WindowCompressorStateConfig, WindowCompressorState -from .create_sliding_window_state import create_sliding_window_state -from .create_window_compressor_state import create_window_compressor_state +from .sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState +from .window_compressor.window_compressor_state import WindowCompressorStateConfig, WindowCompressorState +from .sliding_window.create_sliding_window_state import create_sliding_window_state +from .window_compressor.create_window_compressor_state import create_window_compressor_state __all__ = [ "BoundedAttentionState", diff --git a/lightllm/common/attention_state/bounded_att/gdn/__init__.py b/lightllm/common/attention_state/bounded_att/gdn/__init__.py new file mode 100644 index 0000000000..62dc94c11d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/__init__.py @@ -0,0 +1,5 @@ +from .create_gdn_state import create_gdn_state +from .gdn_pages_pool import GDNPagesPool +from .gdn_state import GDNState + +__all__ = ["create_gdn_state", "GDNPagesPool", "GDNState"] diff --git a/lightllm/common/attention_state/bounded_att/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py similarity index 91% rename from lightllm/common/attention_state/bounded_att/create_gdn_state.py rename to lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 33375c7372..3fe6f2317e 100644 --- a/lightllm/common/attention_state/bounded_att/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -4,7 +4,7 @@ import torch -from .conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_types import ConvSSMStateConfig from .gdn_state import GDNState diff --git a/lightllm/common/attention_state/bounded_att/gdn_pages_pool.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py similarity index 92% rename from lightllm/common/attention_state/bounded_att/gdn_pages_pool.py rename to lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py index 12095f12fc..a54332eb84 100644 --- a/lightllm/common/attention_state/bounded_att/gdn_pages_pool.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py @@ -5,8 +5,8 @@ import torch -from .conv_ssm_types import ConvSSMStateConfig -from .utils import create_page_pool_tensor, register_page_pool_tensor +from ..conv_ssm_types import ConvSSMStateConfig +from ..utils import create_page_pool_tensor, register_page_pool_tensor class GDNPagesPool: diff --git a/lightllm/common/attention_state/bounded_att/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py similarity index 98% rename from lightllm/common/attention_state/bounded_att/gdn_state.py rename to lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index c850067c87..2cb26f0aed 100644 --- a/lightllm/common/attention_state/bounded_att/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -8,13 +8,13 @@ from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from .base_state import BoundedAttentionState -from .conv_ssm_types import ConvSSMStateConfig +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig from .gdn_pages_pool import GDNPagesPool -from .transfer_utils import bounded_att_buffer_copies -from .copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: from lightllm.common.attention_state.types import PageKind diff --git a/lightllm/common/attention_state/bounded_att/kda/__init__.py b/lightllm/common/attention_state/bounded_att/kda/__init__.py new file mode 100644 index 0000000000..7154892202 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/__init__.py @@ -0,0 +1,5 @@ +from .create_kda_state import create_kda_state +from .kda_pages_pool import KDAPagesPool +from .kda_state import KDAState + +__all__ = ["create_kda_state", "KDAPagesPool", "KDAState"] diff --git a/lightllm/common/attention_state/bounded_att/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py similarity index 91% rename from lightllm/common/attention_state/bounded_att/create_kda_state.py rename to lightllm/common/attention_state/bounded_att/kda/create_kda_state.py index c1934c666b..49c9ad0ce5 100644 --- a/lightllm/common/attention_state/bounded_att/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -4,7 +4,7 @@ import torch -from .conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_types import ConvSSMStateConfig from .kda_state import KDAState diff --git a/lightllm/common/attention_state/bounded_att/kda_pages_pool.py b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py similarity index 93% rename from lightllm/common/attention_state/bounded_att/kda_pages_pool.py rename to lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py index 9ac96b1661..0ec20ae4fe 100644 --- a/lightllm/common/attention_state/bounded_att/kda_pages_pool.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py @@ -5,8 +5,8 @@ import torch -from .conv_ssm_types import ConvSSMStateConfig -from .utils import create_page_pool_tensor, register_page_pool_tensor +from ..conv_ssm_types import ConvSSMStateConfig +from ..utils import create_page_pool_tensor, register_page_pool_tensor class KDAPagesPool: diff --git a/lightllm/common/attention_state/bounded_att/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py similarity index 98% rename from lightllm/common/attention_state/bounded_att/kda_state.py rename to lightllm/common/attention_state/bounded_att/kda/kda_state.py index b34ffdd3f4..9bb105dd1f 100644 --- a/lightllm/common/attention_state/bounded_att/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -8,13 +8,13 @@ from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from .base_state import BoundedAttentionState -from .conv_ssm_types import ConvSSMStateConfig +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig from .kda_pages_pool import KDAPagesPool -from .transfer_utils import bounded_att_buffer_copies -from .copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors if TYPE_CHECKING: from lightllm.common.attention_state.types import PageKind diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py b/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py new file mode 100644 index 0000000000..06a759d916 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py @@ -0,0 +1,4 @@ +from .create_sliding_window_state import create_sliding_window_state +from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState + +__all__ = ["create_sliding_window_state", "FP8E8M0SlidingWindowConfig", "FP8E8M0SlidingWindowState"] diff --git a/lightllm/common/attention_state/bounded_att/create_sliding_window_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/create_sliding_window_state.py similarity index 100% rename from lightllm/common/attention_state/bounded_att/create_sliding_window_state.py rename to lightllm/common/attention_state/bounded_att/sliding_window/create_sliding_window_state.py diff --git a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py similarity index 97% rename from lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py rename to lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index 4400aa1f5b..a1d72102ac 100644 --- a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_state.py +++ b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py @@ -9,12 +9,12 @@ import torch -from ..utils import deepcopy_result -from ..types import PageKind, StateTransferLayout, TensorField, TransferRegion -from .base_state import BoundedAttentionState -from .copy_kernel import copy_bounded_att_tensors -from .transfer_utils import bounded_att_buffer_copies -from .window_pages_pool import WindowPagesPool +from ...utils import deepcopy_result +from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..window_pages_pool import WindowPagesPool from .fp8_e8m0_swa_transfer import copy_swa_window diff --git a/lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_transfer.py similarity index 100% rename from lightllm/common/attention_state/bounded_att/fp8_e8m0_swa_transfer.py rename to lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_transfer.py diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py b/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py new file mode 100644 index 0000000000..41294a3751 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py @@ -0,0 +1,4 @@ +from .create_window_compressor_state import create_window_compressor_state +from .window_compressor_state import WindowCompressorState, WindowCompressorStateConfig + +__all__ = ["create_window_compressor_state", "WindowCompressorState", "WindowCompressorStateConfig"] diff --git a/lightllm/common/attention_state/bounded_att/create_window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/create_window_compressor_state.py similarity index 100% rename from lightllm/common/attention_state/bounded_att/create_window_compressor_state.py rename to lightllm/common/attention_state/bounded_att/window_compressor/create_window_compressor_state.py diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py similarity index 98% rename from lightllm/common/attention_state/bounded_att/window_compressor_state.py rename to lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index 91b540760e..9fe226d456 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -12,12 +12,12 @@ from lightllm.utils.envs_utils import get_env_start_args -from ..utils import deepcopy_result -from ..types import PageKind, StateTransferLayout, TensorField, TransferRegion -from .base_state import BoundedAttentionState -from .copy_kernel import copy_bounded_att_tensors -from .transfer_utils import bounded_att_buffer_copies -from .window_pages_pool import WindowPagesPool +from ...utils import deepcopy_result +from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..window_pages_pool import WindowPagesPool from .window_compressor_transfer import copy_compressor_window diff --git a/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py similarity index 99% rename from lightllm/common/attention_state/bounded_att/window_compressor_transfer.py rename to lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py index 49e88896b0..dc985d68e6 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor_transfer.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from ..utils import view_max_copy_dtype +from ...utils import view_max_copy_dtype def copy_compressor_window( diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index 88a3f85e2f..78abd26ca4 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -1,9 +1,10 @@ -from .create_mla_state import create_mla_state -from .create_packed_mla_state import create_packed_mla_state -from .create_gqa_state import create_gqa_state +from .mla.create_mla_state import create_mla_state +from .packed_mla.create_packed_mla_state import create_packed_mla_state +from .gqa.create_gqa_state import create_gqa_state from .base_state import FullAttentionState -from .gqa_state import GQAAttentionConfig, GQAFullAttentionState -from .mla_state import MLAConfig, MLAFullAttentionState +from .gqa.gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .gqa import gqa_transfer +from .mla.mla_state import MLAConfig, MLAFullAttentionState __all__ = [ "create_gqa_state", @@ -16,12 +17,12 @@ "MLAFullAttentionState", ] -from .group_quantized_state import ( +from .group_quantized.group_quantized_state import ( GroupQuantizedFullAttentionState, GroupQuantizedAttentionConfig, ) -from .static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig -from .packed_mla_state import ( +from .static_fp8.static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig +from .packed_mla.packed_mla_state import ( PackedMLAFullAttentionState, PackedMLAAttentionConfig, PackedMLALogicalConfig, @@ -39,8 +40,8 @@ "PackedMLALayerView", ] -from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionState, CompressedAttentionLayerView -from .create_compressed_att_state import create_compressed_att_state +from .compressed_att.compressed_att_state import CompressedAttentionConfig, CompressedAttentionState, CompressedAttentionLayerView +from .compressed_att.create_compressed_att_state import create_compressed_att_state __all__ += [ "create_compressed_att_state", @@ -49,6 +50,6 @@ "CompressedAttentionLayerView", ] -from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState +from .fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState __all__ += ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/compressed_att/__init__.py b/lightllm/common/attention_state/full_att/compressed_att/__init__.py new file mode 100644 index 0000000000..4cbe7773f0 --- /dev/null +++ b/lightllm/common/attention_state/full_att/compressed_att/__init__.py @@ -0,0 +1,9 @@ +from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionLayerView, CompressedAttentionState +from .create_compressed_att_state import create_compressed_att_state + +__all__ = [ + "CompressedAttentionConfig", + "CompressedAttentionLayerView", + "CompressedAttentionState", + "create_compressed_att_state", +] diff --git a/lightllm/common/attention_state/full_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py similarity index 98% rename from lightllm/common/attention_state/full_att/compressed_att_state.py rename to lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py index f1203b41c3..21ce083697 100644 --- a/lightllm/common/attention_state/full_att/compressed_att_state.py +++ b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py @@ -12,11 +12,11 @@ from lightllm.utils.envs_utils import get_env_start_args -from ..utils import deepcopy_result -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState -from .packed_transfer import transfer_packed_pages -from .dp_transfer import copy_packed_tokens +from ...utils import deepcopy_result +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..dp_transfer import copy_packed_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/create_compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/create_compressed_att_state.py similarity index 100% rename from lightllm/common/attention_state/full_att/create_compressed_att_state.py rename to lightllm/common/attention_state/full_att/compressed_att/create_compressed_att_state.py diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py new file mode 100644 index 0000000000..a5bf7bf54f --- /dev/null +++ b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py @@ -0,0 +1,3 @@ +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + +__all__ = ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py similarity index 97% rename from lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py rename to lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py index 71a71f2433..1e7f58ce91 100644 --- a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py @@ -11,12 +11,12 @@ from lightllm.utils.envs_utils import get_env_start_args -from ..utils import deepcopy_result -from ..types import AxisMapping, StateTransferLayout, TransferRegion -from .base_state import FullAttentionState -from .dp_transfer import copy_packed_tokens -from .packed_transfer import transfer_packed_pages -from .packed_mla_state import PackedMLALayerView +from ...utils import deepcopy_result +from ...types import AxisMapping, StateTransferLayout, TransferRegion +from ..base_state import FullAttentionState +from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..packed_mla.packed_mla_state import PackedMLALayerView @dataclass(frozen=True) @@ -125,7 +125,7 @@ def write_full_att_kv( kv: torch.Tensor, ) -> None: """使用固定 group64/E8M0 算子量化并写入 KV，保留 indexer 和行尾 padding。""" - from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import ( + from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import ( write_fp8_e8m0_packed_kv, ) diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_write.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_write.py similarity index 100% rename from lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla_write.py rename to lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_write.py diff --git a/lightllm/common/attention_state/full_att/gqa/__init__.py b/lightllm/common/attention_state/full_att/gqa/__init__.py new file mode 100644 index 0000000000..1b28b4ad19 --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/__init__.py @@ -0,0 +1,4 @@ +from .create_gqa_state import create_gqa_state +from .gqa_state import GQAAttentionConfig, GQAFullAttentionState + +__all__ = ["create_gqa_state", "GQAAttentionConfig", "GQAFullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/create_gqa_state.py b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py similarity index 91% rename from lightllm/common/attention_state/full_att/create_gqa_state.py rename to lightllm/common/attention_state/full_att/gqa/create_gqa_state.py index d25a7c37fc..a9c0e5ac4d 100644 --- a/lightllm/common/attention_state/full_att/create_gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py @@ -7,8 +7,8 @@ import torch from .gqa_state import GQAAttentionConfig, GQAFullAttentionState -from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState -from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState +from ..group_quantized.group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from ..static_fp8.static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState def create_gqa_state( diff --git a/lightllm/common/attention_state/full_att/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py similarity index 99% rename from lightllm/common/attention_state/full_att/gqa_state.py rename to lightllm/common/attention_state/full_att/gqa/gqa_state.py index 5dea8622d2..a9d389a91d 100644 --- a/lightllm/common/attention_state/full_att/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -10,11 +10,11 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState from .gqa_transfer import transfer_gqa_pages -from .dp_transfer import copy_gqa_tokens +from ..dp_transfer import copy_gqa_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py similarity index 99% rename from lightllm/common/attention_state/full_att/gqa_transfer.py rename to lightllm/common/attention_state/full_att/gqa/gqa_transfer.py index ccd7edc990..994cd3904a 100644 --- a/lightllm/common/attention_state/full_att/gqa_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py @@ -4,8 +4,8 @@ import triton import triton.language as tl -from ..types import StateTransferLayout -from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/group_quantized/__init__.py b/lightllm/common/attention_state/full_att/group_quantized/__init__.py new file mode 100644 index 0000000000..d276016930 --- /dev/null +++ b/lightllm/common/attention_state/full_att/group_quantized/__init__.py @@ -0,0 +1,3 @@ +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState + +__all__ = ["GroupQuantizedAttentionConfig", "GroupQuantizedFullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/group_quantized_state.py b/lightllm/common/attention_state/full_att/group_quantized/group_quantized_state.py similarity index 99% rename from lightllm/common/attention_state/full_att/group_quantized_state.py rename to lightllm/common/attention_state/full_att/group_quantized/group_quantized_state.py index 5d8e6532e3..a35a1aad9b 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/group_quantized/group_quantized_state.py @@ -10,11 +10,11 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState from .group_quantized_transfer import transfer_group_quantized_pages -from .dp_transfer import copy_group_quantized_tokens +from ..dp_transfer import copy_group_quantized_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/group_quantized/group_quantized_transfer.py similarity index 99% rename from lightllm/common/attention_state/full_att/group_quantized_transfer.py rename to lightllm/common/attention_state/full_att/group_quantized/group_quantized_transfer.py index a92176982d..1662896ff0 100644 --- a/lightllm/common/attention_state/full_att/group_quantized_transfer.py +++ b/lightllm/common/attention_state/full_att/group_quantized/group_quantized_transfer.py @@ -4,8 +4,8 @@ import triton import triton.language as tl -from ..types import StateTransferLayout -from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/mla/__init__.py b/lightllm/common/attention_state/full_att/mla/__init__.py new file mode 100644 index 0000000000..557ad284a2 --- /dev/null +++ b/lightllm/common/attention_state/full_att/mla/__init__.py @@ -0,0 +1,4 @@ +from .create_mla_state import create_mla_state +from .mla_state import MLAConfig, MLAFullAttentionState + +__all__ = ["create_mla_state", "MLAConfig", "MLAFullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/create_mla_state.py b/lightllm/common/attention_state/full_att/mla/create_mla_state.py similarity index 92% rename from lightllm/common/attention_state/full_att/create_mla_state.py rename to lightllm/common/attention_state/full_att/mla/create_mla_state.py index 3c36883a89..f93c9775af 100644 --- a/lightllm/common/attention_state/full_att/create_mla_state.py +++ b/lightllm/common/attention_state/full_att/mla/create_mla_state.py @@ -4,7 +4,7 @@ import torch -from .base_state import FullAttentionState +from ..base_state import FullAttentionState from .mla_state import MLAConfig, MLAFullAttentionState diff --git a/lightllm/common/attention_state/full_att/mla_state.py b/lightllm/common/attention_state/full_att/mla/mla_state.py similarity index 98% rename from lightllm/common/attention_state/full_att/mla_state.py rename to lightllm/common/attention_state/full_att/mla/mla_state.py index cb49d5efcd..e10dd953a1 100644 --- a/lightllm/common/attention_state/full_att/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla/mla_state.py @@ -9,11 +9,11 @@ from lightllm.utils.envs_utils import get_env_start_args -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState -from .packed_transfer import transfer_packed_pages -from .dp_transfer import copy_packed_tokens +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..dp_transfer import copy_packed_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/packed_mla/__init__.py b/lightllm/common/attention_state/full_att/packed_mla/__init__.py new file mode 100644 index 0000000000..40b9e49e25 --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/__init__.py @@ -0,0 +1,9 @@ +from .create_packed_mla_state import create_packed_mla_state +from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState + +__all__ = [ + "create_packed_mla_state", + "PackedMLAAttentionConfig", + "PackedMLALogicalConfig", + "PackedMLAFullAttentionState", +] diff --git a/lightllm/common/attention_state/full_att/create_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py similarity index 89% rename from lightllm/common/attention_state/full_att/create_packed_mla_state.py rename to lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py index 00c72638a1..d60ecca7dd 100644 --- a/lightllm/common/attention_state/full_att/create_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py @@ -4,9 +4,9 @@ import torch -from .base_state import FullAttentionState +from ..base_state import FullAttentionState from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState -from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState +from ..fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState def create_packed_mla_state( diff --git a/lightllm/common/attention_state/full_att/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py similarity index 99% rename from lightllm/common/attention_state/full_att/packed_mla_state.py rename to lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py index 39c54c8963..fb6281fd7c 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py @@ -10,11 +10,11 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result, dtype_name -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState from .packed_mla_transfer import transfer_packed_mla_pages -from .dp_transfer import copy_packed_tokens +from ..dp_transfer import copy_packed_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py similarity index 99% rename from lightllm/common/attention_state/full_att/packed_mla_transfer.py rename to lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py index 5d437d78cd..5b5bb4534f 100644 --- a/lightllm/common/attention_state/full_att/packed_mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py @@ -4,8 +4,8 @@ import triton import triton.language as tl -from ..types import StateTransferLayout -from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/packed_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py similarity index 99% rename from lightllm/common/attention_state/full_att/packed_transfer.py rename to lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py index e7bdb50f92..bd04051a98 100644 --- a/lightllm/common/attention_state/full_att/packed_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py @@ -4,8 +4,8 @@ import triton import triton.language as tl -from ..types import StateTransferLayout -from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/static_fp8/__init__.py b/lightllm/common/attention_state/full_att/static_fp8/__init__.py new file mode 100644 index 0000000000..f11df1de2d --- /dev/null +++ b/lightllm/common/attention_state/full_att/static_fp8/__init__.py @@ -0,0 +1,3 @@ +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState + +__all__ = ["StaticFP8AttentionConfig", "StaticFP8FullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/static_fp8_state.py b/lightllm/common/attention_state/full_att/static_fp8/static_fp8_state.py similarity index 99% rename from lightllm/common/attention_state/full_att/static_fp8_state.py rename to lightllm/common/attention_state/full_att/static_fp8/static_fp8_state.py index b7df6f2cfe..d907d7813d 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/static_fp8/static_fp8_state.py @@ -12,11 +12,11 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ..utils import deepcopy_result -from ..types import AxisMapping, TransferRegion -from .base_state import FullAttentionState +from ...utils import deepcopy_result +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState from .static_fp8_transfer import transfer_static_fp8_pages -from .dp_transfer import copy_gqa_tokens +from ..dp_transfer import copy_gqa_tokens if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout diff --git a/lightllm/common/attention_state/full_att/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/static_fp8/static_fp8_transfer.py similarity index 99% rename from lightllm/common/attention_state/full_att/static_fp8_transfer.py rename to lightllm/common/attention_state/full_att/static_fp8/static_fp8_transfer.py index e8bd29d220..9c70016acd 100644 --- a/lightllm/common/attention_state/full_att/static_fp8_transfer.py +++ b/lightllm/common/attention_state/full_att/static_fp8/static_fp8_transfer.py @@ -4,8 +4,8 @@ import triton import triton.language as tl -from ..types import StateTransferLayout -from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype # ======================================================================== diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 7f425aae07..06ab80a947 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -117,7 +117,7 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 """ import json - from .full_att.static_fp8_state import StaticFP8FullAttentionState + from .full_att.static_fp8.static_fp8_state import StaticFP8FullAttentionState full_states = self.get_full_attention_states() calibration_states = [state for state in full_states if isinstance(state, StaticFP8FullAttentionState)] diff --git a/lightllm/models/deepseek_v4/config.py b/lightllm/models/deepseek_v4/config.py index a57e89bf4e..a5dafbf108 100644 --- a/lightllm/models/deepseek_v4/config.py +++ b/lightllm/models/deepseek_v4/config.py @@ -4,8 +4,8 @@ from dataclasses import dataclass from typing import ClassVar -from lightllm.common.attention_state.bounded_att.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig -from lightllm.common.attention_state.full_att.compressed_att_state import CompressedAttentionConfig +from lightllm.common.attention_state.bounded_att.sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import CompressedAttentionConfig @dataclass(frozen=True) diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index b0abc2f24b..6f90fa1207 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -4,7 +4,7 @@ import torch from lightllm.common.basemodel import InferStateInfo from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState -from lightllm.common.attention_state.full_att.compressed_att_state import ( +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( CompressedAttentionState, CompressedAttentionLayerView, ) diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index 5cd25f2466..9a296328ac 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -105,7 +105,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list[FullAttentionState]: """分别创建 C4/C128 定义；逻辑索引保留原始 token 域，不提前分配存储。""" - from lightllm.common.attention_state.full_att.create_compressed_att_state import create_compressed_att_state + from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import create_compressed_att_state from .config import create_attention_config assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py index 169ba44485..5e5a8a236c 100644 --- a/unit_tests/common/att_resources/test_dsv4_state.py +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -6,7 +6,7 @@ import torch from lightllm.models.deepseek_v4.config import DSV4AttentionConfig -from lightllm.common.attention_state.full_att.compressed_att_state import ( +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( CompressedAttentionConfig, CompressedAttentionState, ) diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py index 745c9fcfa7..a505ed932c 100644 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -112,7 +112,7 @@ def test_glm_main_and_draft_explicitly_select_encoding(draft: bool, kv_type: str @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_fixed_quantizer_preserves_indexer_and_other_rows() -> None: - from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv + from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv state = make_state(device="cuda") state.storage.fill_(173) diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index 0df056ad4e..bb6e213a45 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -6,7 +6,7 @@ import torch from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState -from lightllm.common.attention_state.full_att.gqa_transfer import transfer_gqa_pages +from lightllm.common.attention_state.full_att.gqa.gqa_transfer import transfer_gqa_pages from lightllm.common.attention_state.types import StateTransferLayout if TYPE_CHECKING: diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 1f11415b1c..908167e7cb 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from collections.abc import Sequence from lightllm.common.attention_state.manager import AttentionStateManager - from lightllm.common.attention_state.full_att.packed_mla_state import PackedMLAFullAttentionState + from lightllm.common.attention_state.full_att.packed_mla.packed_mla_state import PackedMLAFullAttentionState from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter from multiprocessing.connection import Connection diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py index f1a5b9a4bd..18311d021d 100644 --- a/unit_tests/common/dsv4_test_utils.py +++ b/unit_tests/common/dsv4_test_utils.py @@ -17,7 +17,7 @@ from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState from lightllm.common.attention_state.bounded_att.window_pages_pool import WindowPagesPool from lightllm.common.attention_state.bounded_att.utils import create_page_pool_tensor -from lightllm.common.attention_state.full_att.compressed_att_state import ( +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( CompressedAttentionConfig, CompressedAttentionState, ) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 91fbf63576..cf25691ad6 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -21,7 +21,7 @@ fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv +from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( _build_prefill_row_table_kernel, @@ -1283,7 +1283,7 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) config = PackedMLALogicalConfig( logical_layer_groups=((3,),), latent_dim=512, format="packed_group64_e8m0", ) From 161c3b1d56abc0f9aeca8c1619ab66188e124c23 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 04:49:46 +0000 Subject: [PATCH 37/56] refactor: group quantized states by attention type --- docs/cache_manager_design.md | 11 ++++++----- lightllm/common/attention_state/full_att/__init__.py | 6 +++--- .../full_att/fp8_e8m0_packed_mla/__init__.py | 3 --- .../common/attention_state/full_att/gqa/__init__.py | 12 +++++++++++- .../attention_state/full_att/gqa/create_gqa_state.py | 4 ++-- .../group_quantized_state.py | 0 .../group_quantized_transfer.py | 0 .../full_att/{static_fp8 => gqa}/static_fp8_state.py | 0 .../{static_fp8 => gqa}/static_fp8_transfer.py | 0 .../full_att/group_quantized/__init__.py | 3 --- .../attention_state/full_att/packed_mla/__init__.py | 3 +++ .../full_att/packed_mla/create_packed_mla_state.py | 2 +- .../fp8_e8m0_packed_mla_state.py | 2 +- .../fp8_e8m0_packed_mla_write.py | 0 .../attention_state/full_att/static_fp8/__init__.py | 3 --- lightllm/common/attention_state/manager.py | 2 +- .../att_resources/test_fp8_e8m0_packed_mla_state.py | 2 +- unit_tests/models/glm5_next/test_fp8_kv.py | 4 ++-- 18 files changed, 31 insertions(+), 26 deletions(-) delete mode 100644 lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py rename lightllm/common/attention_state/full_att/{group_quantized => gqa}/group_quantized_state.py (100%) rename lightllm/common/attention_state/full_att/{group_quantized => gqa}/group_quantized_transfer.py (100%) rename lightllm/common/attention_state/full_att/{static_fp8 => gqa}/static_fp8_state.py (100%) rename lightllm/common/attention_state/full_att/{static_fp8 => gqa}/static_fp8_transfer.py (100%) delete mode 100644 lightllm/common/attention_state/full_att/group_quantized/__init__.py rename lightllm/common/attention_state/full_att/{fp8_e8m0_packed_mla => packed_mla}/fp8_e8m0_packed_mla_state.py (99%) rename lightllm/common/attention_state/full_att/{fp8_e8m0_packed_mla => packed_mla}/fp8_e8m0_packed_mla_write.py (100%) delete mode 100644 lightllm/common/attention_state/full_att/static_fp8/__init__.py diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index 9684bb5bfc..abe8f9e616 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -310,17 +310,18 @@ lightllm/common/ │ ├── full_att/ │ │ ├── __init__.py │ │ ├── base_state.py # FullAttentionState：抽象行为接口 -│ │ ├── gqa/ # GQA 配置、State、创建入口和传输实现 +│ │ ├── gqa/ # GQA 及其量化变体的配置、State、创建入口和传输实现 │ │ │ ├── create_gqa_state.py │ │ │ ├── gqa_state.py +│ │ │ ├── group_quantized_state.py +│ │ │ ├── group_quantized_transfer.py +│ │ │ ├── static_fp8_state.py +│ │ │ ├── static_fp8_transfer.py │ │ │ └── gqa_transfer.py │ │ ├── dp_transfer.py # 各 full-attention 方案的直接索引 DP 聚合复制 -│ │ ├── group_quantized/ # INT4/INT8 配置、State 和动态 scale 传输 -│ │ ├── static_fp8/ # 静态 FP8 配置、State 和传输 │ │ ├── mla/ # 普通 MLA latent 配置、State 和创建入口 │ │ ├── compressed_att/ # V4 C4/C128 压缩块配置、State 和创建入口 -│ │ ├── packed_mla/ # packed MLA 配置、State、创建入口和传输 -│ │ └── fp8_e8m0_packed_mla/ # group64/E8M0 packed MLA 配置、State 和写入 kernel +│ │ └── packed_mla/ # packed MLA 及其 FP8/E8M0 变体的配置、State、创建入口和传输 │ ├── bounded_att/ │ │ ├── __init__.py │ │ ├── base_state.py # BoundedAttentionState：有界状态的抽象操作契约 diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index 78abd26ca4..4c57c49c1e 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -17,11 +17,11 @@ "MLAFullAttentionState", ] -from .group_quantized.group_quantized_state import ( +from .gqa.group_quantized_state import ( GroupQuantizedFullAttentionState, GroupQuantizedAttentionConfig, ) -from .static_fp8.static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig +from .gqa.static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig from .packed_mla.packed_mla_state import ( PackedMLAFullAttentionState, PackedMLAAttentionConfig, @@ -50,6 +50,6 @@ "CompressedAttentionLayerView", ] -from .fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState +from .packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState __all__ += ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py b/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py deleted file mode 100644 index a5bf7bf54f..0000000000 --- a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState - -__all__ = ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/gqa/__init__.py b/lightllm/common/attention_state/full_att/gqa/__init__.py index 1b28b4ad19..5ef32afe8e 100644 --- a/lightllm/common/attention_state/full_att/gqa/__init__.py +++ b/lightllm/common/attention_state/full_att/gqa/__init__.py @@ -1,4 +1,14 @@ from .create_gqa_state import create_gqa_state from .gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState -__all__ = ["create_gqa_state", "GQAAttentionConfig", "GQAFullAttentionState"] +__all__ = [ + "create_gqa_state", + "GQAAttentionConfig", + "GQAFullAttentionState", + "GroupQuantizedAttentionConfig", + "GroupQuantizedFullAttentionState", + "StaticFP8AttentionConfig", + "StaticFP8FullAttentionState", +] diff --git a/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py index a9c0e5ac4d..d25a7c37fc 100644 --- a/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py @@ -7,8 +7,8 @@ import torch from .gqa_state import GQAAttentionConfig, GQAFullAttentionState -from ..group_quantized.group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState -from ..static_fp8.static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState def create_gqa_state( diff --git a/lightllm/common/attention_state/full_att/group_quantized/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py similarity index 100% rename from lightllm/common/attention_state/full_att/group_quantized/group_quantized_state.py rename to lightllm/common/attention_state/full_att/gqa/group_quantized_state.py diff --git a/lightllm/common/attention_state/full_att/group_quantized/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py similarity index 100% rename from lightllm/common/attention_state/full_att/group_quantized/group_quantized_transfer.py rename to lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py diff --git a/lightllm/common/attention_state/full_att/static_fp8/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py similarity index 100% rename from lightllm/common/attention_state/full_att/static_fp8/static_fp8_state.py rename to lightllm/common/attention_state/full_att/gqa/static_fp8_state.py diff --git a/lightllm/common/attention_state/full_att/static_fp8/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py similarity index 100% rename from lightllm/common/attention_state/full_att/static_fp8/static_fp8_transfer.py rename to lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py diff --git a/lightllm/common/attention_state/full_att/group_quantized/__init__.py b/lightllm/common/attention_state/full_att/group_quantized/__init__.py deleted file mode 100644 index d276016930..0000000000 --- a/lightllm/common/attention_state/full_att/group_quantized/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState - -__all__ = ["GroupQuantizedAttentionConfig", "GroupQuantizedFullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/packed_mla/__init__.py b/lightllm/common/attention_state/full_att/packed_mla/__init__.py index 40b9e49e25..8447751eae 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/__init__.py +++ b/lightllm/common/attention_state/full_att/packed_mla/__init__.py @@ -1,9 +1,12 @@ from .create_packed_mla_state import create_packed_mla_state from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState __all__ = [ "create_packed_mla_state", "PackedMLAAttentionConfig", "PackedMLALogicalConfig", "PackedMLAFullAttentionState", + "FP8E8M0PackedMLAConfig", + "FP8E8M0PackedMLAState", ] diff --git a/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py index d60ecca7dd..afb237eb7a 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py @@ -6,7 +6,7 @@ from ..base_state import FullAttentionState from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState -from ..fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState def create_packed_mla_state( diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py similarity index 99% rename from lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py rename to lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py index 1e7f58ce91..4f018ef215 100644 --- a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py @@ -125,7 +125,7 @@ def write_full_att_kv( kv: torch.Tensor, ) -> None: """使用固定 group64/E8M0 算子量化并写入 KV，保留 indexer 和行尾 padding。""" - from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import ( + from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import ( write_fp8_e8m0_packed_kv, ) diff --git a/lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_write.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_write.py similarity index 100% rename from lightllm/common/attention_state/full_att/fp8_e8m0_packed_mla/fp8_e8m0_packed_mla_write.py rename to lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_write.py diff --git a/lightllm/common/attention_state/full_att/static_fp8/__init__.py b/lightllm/common/attention_state/full_att/static_fp8/__init__.py deleted file mode 100644 index f11df1de2d..0000000000 --- a/lightllm/common/attention_state/full_att/static_fp8/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState - -__all__ = ["StaticFP8AttentionConfig", "StaticFP8FullAttentionState"] diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 06ab80a947..ffd539d67e 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -117,7 +117,7 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 """ import json - from .full_att.static_fp8.static_fp8_state import StaticFP8FullAttentionState + from .full_att.gqa.static_fp8_state import StaticFP8FullAttentionState full_states = self.get_full_attention_states() calibration_states = [state for state in full_states if isinstance(state, StaticFP8FullAttentionState)] diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py index a505ed932c..db871d24c0 100644 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -112,7 +112,7 @@ def test_glm_main_and_draft_explicitly_select_encoding(draft: bool, kv_type: str @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_fixed_quantizer_preserves_indexer_and_other_rows() -> None: - from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv + from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv state = make_state(device="cuda") state.storage.fill_(173) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index cf25691ad6..6e9241a55b 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -21,7 +21,7 @@ fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( _build_prefill_row_table_kernel, @@ -1283,7 +1283,7 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.attention_state.full_att.fp8_e8m0_packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) config = PackedMLALogicalConfig( logical_layer_groups=((3,),), latent_dim=512, format="packed_group64_e8m0", ) From 22cae3c2a052a7575f00d6c282efdfa38d1e6619 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 05:01:37 +0000 Subject: [PATCH 38/56] docs: simplify attention resource design --- docs/cache_manager_design.md | 2055 +++++----------------------------- 1 file changed, 267 insertions(+), 1788 deletions(-) diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md index abe8f9e616..45e4112849 100644 --- a/docs/cache_manager_design.md +++ b/docs/cache_manager_design.md @@ -1,1504 +1,224 @@ -# Attention Resources 架构设计草案 +# Attention Resources 架构设计 -本文定义 attention 资源的架构设计，包括 `AttentionStateManager` 的职责、共享 allocator、各类 State 的接口、目录组织及 CPU/DP/PD 传输约定。文中的结构与接口为设计契约，待细化事项集中列于第 9 节。 +本文只定义 attention resource 的架构、公共接口和实现约束。具体模型参数、Tensor 形状、页大小、kernel 调优值和历史实现过程不在本文维护；这些内容属于代码和测试的实现细节。 -## 1. 设计约定 +## 1. 设计目标 -统一在 `lightllm/common/attention_state/` 下组织可复用的资源实现，由 `AttentionStateManager` 提供资源初始化、内部状态页聚合操作及 CPU/DP/PD 数据传输的默认实现。所有模型直接使用该类，资源组合由模型的两类 State 工厂决定。radix cache、CPU cache 和传输调用方依赖模型持有的管理器，资源层不反向依赖这些调用方。 +`lightllm/common/attention_state/` 提供可组合的 attention 状态资源。资源层负责状态存储、布局描述、快照和跨设备数据操作；请求调度、radix cache、CPU cache、通信队列和页面生命周期由调用方负责。 -公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： +设计必须满足： -```python -class AttentionStateManager: - full_att_allocator: FullAttKVAllocator | None - bounded_att_big_page_allocator: StatePageAllocator | None - bounded_att_small_page_allocator: StatePageAllocator | None -``` - -这三个入口表示资源管理器对分配域的管理责任，聚合操作与传输接口见第 6.2 节。具体初始化必须保证所需 allocator 已创建或绑定，无需的 allocator 为 `None`。例如纯 full-attention 模型只初始化 `full_att_allocator`；启用 bounded 状态页的组合再按配置初始化大小页 allocator。GPU 运行态的存在本身不要求启用页池。 - -- 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att` 成员。具体模型按需组织资源组合与实现类型，拥有的 State 成员统一遵守第 3.5 节的角色、类型和存储方案命名。 -- `full_att` 保存全局历史 KV；`bounded_att` 管理绑定请求的有界运行态，包括 KDA/GDN、SWA，以及用于 KV 压缩和 K-pool 索引池化的窗口候选/gate。各 State 决定自身的快照需求和恢复边界。 -- main/draft 可以采用不同方案；同一共享分配域的参与方使用同一个 allocator 实例，遵守统一的索引和回收约定。资源对象持有 buffer 和布局信息；管理器创建或绑定共享 allocator，调用方直接使用对应 allocator 分配与归还索引。 -- `AttentionStateManager` 提供运行态清零、token 复制、状态页的聚合保存、恢复和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 -- `req_to_token_indexs` 和请求槽位仍由 `ReqManager` 持有。资源管理器接收显式传入的物理索引、请求槽位及恢复长度，提供资源操作接口。 -- 内部调用契约使用 `assert`，例如页类型、物理索引范围、Tensor 形状及共享分配域的一致性；用户配置、外部页面内容及格式兼容性继续显式校验，资源不足等运行失败保留相应异常。GPU 上的 MTP 状态选择保留异步断言，避免引入主机同步。 -- 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 - -本文区分已确定的职责边界和接口草案。main/draft 的资源由共享的 `AttentionStateManager` 统一持有，模型只持有管理器，计算上下文和 backend 从管理器绑定各自角色的 State；后续细化事项包括新模型的具体存储方案与功能接口参数的数据形式。 - -### 1.1 接口命名、排列与注释规范 - -模块统一命名为 `attention_state`，公共管理器为 `AttentionStateManager`，模型和计算上下文通过 -`att_state_manager` 持有同一实例。管理器实现集中在 `attention_state/manager.py`，由包根目录导出。 - -| 命名 | 明确含义 | -| --- | --- | -| `register_state(model_role, state)` / `get_state(...)` | 按 `main`、`draft0..draft35` 注册或选择 State；`state_type_filter` 为算法/方案标记的包含匹配条件 | -| `finish_state_registration()` | 全部模型结束注册，允许缓存查询；不代表实际存储已分配 | -| `initialize_storage(...)` / `allocate_storage(...)` | manager 的前者规划容量并初始化存储，后者按已确定容量分配；FullAttentionState 直接提供 `allocate_storage(token_capacity)` | -| `save_bounded_att_runtime_to_pages` / `restore_bounded_att_runtime_from_pages` | 请求运行态与本地大小页快照之间的数据操作，不使用外部页 Region | -| `write_state_pages_to_cpu` / `read_state_pages_from_cpu` | 本地快照页与 CPU 联合页之间的数据操作，按布局处理字段位置和 TP 分片 | -| `write_to_pd_page` / `read_from_pd_page` | 操作 PD 页集合中 `pd_page_index` 指定的单个页面 | -| `get_cpu_page_layouts`、`get_pd_token_page_layouts`、`get_pd_runtime_page_layouts` | 分别返回 CPU 联合页、PD token 页和 PD 运行态页的布局 | -| `state_name`、`page_slot_ids`、`cpu_page_ready_flags` | 分别表示规范 State 成员名、本地状态页槽位列表和 CPU 页就绪标记 | - -管理器的 `write_req_to_cpu` / `read_req_from_cpu` 保留原名；绝对状态长度仍统一使用 -`state_token_len` / `state_token_lens`，不改为 prefix 命名。`AxisMapping` 的 -`local_start` / `external_start` 和 `local_circular` / `external_circular` 始终表示本地与外部两侧， -读回时仅反转数据方向；映射集合统一命名为 `axis_mappings`，请求窗口起点由 State 根据本次绝对前缀长度确定。 - -新增 Python 文件、类和接口按 **公共函数在前、私有辅助函数在后** 排列。已有基础流程文件以保留功能、减少 review 差异为优先，保留原有函数相对顺序，不为满足分组规范整体重排。 - -新增和改动较多的文件不使用函数声明中的独立 `*` 分隔符。函数参数不少于 4 个时（不计 `self`、`cls`），声明每行一个参数；调用处不少于 4 个实参时，使用 `参数名=值` 并逐行排列。外部 API 的位置参数约束及 `*args`、`**kwargs` 保持其原有语义；改动较少的存量文件不作整文件格式调整。 - -- 模块中先定义公共函数，再定义以下划线开头的私有辅助函数（双下划线协议函数除外）。文件名以 `_` 开头不改变文件内函数的排列规则。脚本的 `if __name__ == "__main__"` 执行入口仍放在全部定义之后。 -- 类的方法区依次排列：`__new__`、`__init__`、其他公共方法与属性、私有辅助方法。`__getstate__`、`__post_init__` 等 Python 双下划线协议方法属于公共协议，不作为私有辅助方法处理。 -- `async def`、`classmethod`、`staticmethod`、`property`、抽象方法均按方法名分类。属性 getter/setter 等存在定义依赖的方法保持必要顺序，装饰器和关联注释随函数一起移动。 -- 同一组内保留原有相对顺序，便于按接口职责阅读；不按字母排序，不为了调整位置修改函数名称或增加转发层。 -- 模块说明、导入、常量、类型定义和类级数据声明保留必要的依赖顺序。重排须检查装饰器、默认参数、类型注解及模块初始化语句的定义时依赖，并保持函数体与业务语义不变。 - -本规范适用于资源管理器、State、allocator、helper、相关模型与 server 调用代码，以及对应的测试和验证脚本。新增文件和新增类必须遵守上述完整分组规则；在存量文件中增加接口时，新增公共接口放在入口区域，新增私有辅助函数放在相关实现区域，新增代码之间仍按公共在前、私有在后排列，不牵动已有函数的位置。 - -模型初始化、推理执行、调度、缓存和传输等已有基础流程，只保留实现功能所需的改动，不混入纯函数排序、接口改名或无关格式整理。函数名称及 public/private 属性是否调整由功能需要决定，不因排列规范修改；功能重构与存量代码的风格整理分别进行。 - -State 的基类与具体实现同样使用中文大块分隔注释，依次组织基础定义、该类资源的数据操作、CPU Cache、DP、PD、私有辅助实现。KDA/GDN 直接继承 `BoundedAttentionState`；所有具体 full-attention State 直接继承 `FullAttentionState`。具体 State 之间不继承实现，分别在自己的文件中按上述分块组织完整操作。 - -- 分类数据操作统一带 `full_att` 或 `bounded_att` 标识，例如 `write_full_att_kv`、`copy_full_att_tokens`、`reset_bounded_att_state`、`save_bounded_att_runtime_to_pages`。管理器与 State 的同类操作同名，区别在于聚合多个 State 或只操作自身存储。 -- `get_layer_view` 用于计算访问，CPU/DP/PD 保持按传输方向命名。State 不增加仅用于统一名字的转发方法。 -- 所有 attention 配置字段必须提供中文注释，说明业务含义、单位或对应 Tensor 轴、全局/TP 本地范围，以及适用方案和默认值的特殊语义；派生属性还需说明计算关系。资源身份只保存在 State，不放入 config。 -- 已知的调用参数使用显式 `key=value`，如 `size=...`、`dtype=...`、`device=...`、`pin_memory=...`；不先组装 `options` 字典再通过 `**options` 传入。 -- 资源身份、布局范围和分配域选择等必要检查直接在使用处展开，不再包装成只有断言、布尔判断或转发循环的小函数。同一操作已校验的槽位继续复用，避免逐页、逐 State 重复校验分配器。 -- 每个公共接口使用中文 docstring 说明操作对象、索引含义和数据范围；在有差异时明确量化参数、前缀长度、MTP 游标及完整逻辑窗口的处理。私有实现使用独立分块，关键布局或状态转换另加说明。 -- bounded 运行态重置、大小页保存与恢复、MTP 提交和运行态复制集中排列。pool 只保存快照 buffer；前缀与槽位的关联、页面发布和引用生命周期由调用流程管理。 - -## 2. 整体结构与职责边界 - -```mermaid -flowchart TD - Model[模型初始化] --> Concrete[共享 AttentionStateManager] - Context[InferenceContext] --> Req[ReqManager] - Req --> Slots[请求槽位与 KV 映射] - Context -->|前缀复用| Radix[Radix Cache 功能] - Context -->|多级缓存| CPU[CPU Cache 功能] - Radix -->|调用通用状态页操作| Concrete - CPU -->|缓存数据读写| Concrete - Concrete --> Alloc[按需初始化或绑定的共享 allocator] - Concrete --> Internal[模型自行组织的内部资源] - Internal --> Full[full attention 实现] - Internal --> Bounded[bounded attention：KDA/GDN、SWA、窗口压缩器] - Concrete -->|绑定具体视图| Forward[模型与 attention backend] - Transport[DP/PD 调度] -->|调用具体资源操作| Concrete -``` - -图中的资源类别按模型需要选择，可以缺省或包含多个角色、方案不同的实例。箭头表示调用或组合关系；main/draft 推理模型共享同一个 `AttentionStateManager`，分别绑定各自的计算 State。 - -| 层次 | 组件 | 职责 | -| --- | --- | --- | -| 请求流程 | `InferenceContext`、`ReqManager` | 创建、暂停、恢复和完成请求；由 ReqManager 管理请求槽位、KV 索引映射及采样参数 | -| 资源数据操作 | `AttentionStateManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | -| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用`AttentionStateManager`操作资源数据 | -| 资源实现 | 模型所需的 State、存储布局、checkpoint pool | 持有 buffer 和布局，实现读写、复制、重置和快照；IPC 导出范围由管理器统一处理 | -| 传输调度 | DP/PD 与多级缓存模块 | 管理任务、队列、通信、事件同步及完成通知 | - -`AttentionStateManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，`AttentionStateManager`执行清零、复制、保存和恢复。请求阶段的组织和缓存节点策略由相应调用流程负责。 - -代码依赖从调用流程向资源层单向传递。radix cache 依赖管理器的通用状态页操作；管理器与 State 不导入或访问 radix 节点、树和策略。资源实现也不反向导入请求调度模块；传输调度使用数据读写与复制接口。 - -## 3. 对象结构与资源归属 - -```text -ReqManager -├── request_slots -└── req_to_token_indexs # 共享请求到 KV 的索引映射 - -AttentionStateManager # 公共接口与 allocator 入口 -├── full_att_allocator: FullAttKVAllocator | None -├── bounded_att_big_page_allocator: StatePageAllocator | None -└── bounded_att_small_page_allocator: StatePageAllocator | None - -AttentionStateManager # 模型实现，内部结构由其自行定义 -├── 按需创建或绑定所需 allocator -├── 初始化模型所需的资源组合 -├── 实现内部状态页聚合操作与 CPU / DP / PD 数据传输 -└── 绑定推理视图、组织传输所需状态 # 搬运直接调用 State 方法 -``` - -纯 full-attention 与混合模型都使用 `AttentionStateManager`，区别在于模型工厂返回的 State 组合。管理器按角色持有各 State，共享 allocator 按分配域创建；公共调用方使用同一组聚合接口。 - -以下列出可复用资源组件的结构示例。linear attention 直接使用 `KDAState`、`GDNState` 等具体状态对象，由 `BoundedAttentionState` 约定必需操作，不增加资源包装实例。管理器内部的成员遵守统一命名及唯一性约束，checkpoint pool 仅在相应功能启用时存在： - -```text -FullAttentionState -├── capacity / page_size / device # 普通 token 容量、HOLD 页大小及分配设备 -├── config # 每种具体方案的独立配置；packed 的物理 dtype/行宽由编码推导 -├── k_storage / v_storage # GQA 及其量化方案分别分配 K/V；MLA 使用单份 latent storage -└── layer_mapping # 本模型的层映射 - -KDAState # 一种具体 attention 状态实现 -├── conv_state: torch.Tensor # 卷积历史窗口及 MTP 尾部空间 -├── ssm_state: torch.Tensor # KDA 递推状态及 MTP 状态槽位 -├── mtp_state_index: Tensor | None -├── big_pages_pool # 按需创建，使用共享 allocator 分配的槽位 -│ ├── conv_state -│ └── ssm_state -└── small_pages_pool # 按需创建 - ├── conv_state - └── ssm_state - -GDNState # 独立定义 GDN 状态、布局和快照操作 -└── 按其实现组织 Conv / SSM / MTP 状态及可选大小页 pool - -WindowCompressorState # bounded 窗口候选与 gate，供 V4 压缩和 GLM K-pool 使用 -└── buffer # 同 dtype 的 candidate|gate 拼接 Tensor -``` - -同一共享分配域中的 allocator 只初始化一次，由 main/draft 共用的唯一 `AttentionStateManager` 管理生命周期。调用方解除相关引用后，向对应 allocator 单次归还索引，不因模型数增加而重复分配或释放。 - -| 资源 | 持有者 | 分配单位 / 生命周期 | -| --- | --- | --- | -| full-attention 空闲页与容量状态 | `AttentionStateManager.full_att_allocator` | 每次为请求统一分配/回收物理页索引，main/draft 共用 | -| 大页 checkpoint 空闲槽位 | `AttentionStateManager.bounded_att_big_page_allocator` | 为 main/draft 参与的大页 pool 统一分配槽位 ID | -| 小页 checkpoint 空闲槽位 | `AttentionStateManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | -| 请求 KV 索引表 | `ReqManager` | 记录逻辑 token 到共享物理索引的映射；负责绑定、更新和清理 | -| HOLD 保留物理页及索引描述 | `AttentionStateManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | -| KV、量化 scale、packed indexer buffer 与层映射 | `AttentionStateManager` 组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | -| KDA/GDN 的 Conv/SSM、MTP 游标 | `AttentionStateManager`持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | -| 大小页 checkpoint buffer | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID 保存数据；不持有独立空闲列表或长度记录 | -| indexer 原始尾部窗口 | `AttentionStateManager` 组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需窗口布局 | -| 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | -| 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复关联资源数据，解除引用后直接归还 allocator | - -### 3.1 main/draft 共享 full-attention allocator - -调用方直接通过 `AttentionStateManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，`AttentionStateManager`恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： - -```text -逻辑 token 位置 t - ↓ ReqManager.req_to_token_indexs -统一物理索引 i - ├── main 对应的 KV 存储位置 - └── draft 对应的 KV 存储位置 -``` - -二者的层数、head 维度、dtype 和量化布局可以不同；同一索引域内需要使用兼容的 page_size，并且各存储都覆盖 allocator 可返回的索引范围。同一分配域统一容量和物理页索引规则，不引入模型间索引重映射。 - -共享索引表示存储位置的对应关系，不表示 main/draft 在每个位置都已经写入有效 KV。MTP 预留区和回退后的有效长度仍由请求执行流程维护。 - -释放时，请求流程与`AttentionStateManager`确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估累计各模型 State 独立分配的存储，不能因共享 allocator 而漏算。 - -### 3.2 main/draft 共享 linear checkpoint allocator - -main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAState`，draft 使用 `GDNState` 或无需 linear 状态。各状态对象持有自身运行态及启用的大页或小页 pool。对同一前缀的一个大页或小页 checkpoint，调用方直接从管理器持有的对应 allocator 分配一次槽位 ID，再将该 ID 传给管理器的聚合数据操作，供参与的两侧 pool 使用。两侧的状态布局、dtype 和单槽位字节数可以不同，各自的存储必须覆盖共享 allocator 可返回的槽位范围。 - -例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。`AttentionStateManager`负责找到参与操作的 pool，请求流程无需读取模型内部字段。缓存调用流程维护前缀到槽位的关联，并在两侧都完成对应状态的保存后发布；恢复时通过该关联选择槽位。仅一侧具有 linear 资源时，只处理该侧。 - -共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 - -linear 运行态与两侧索引池化窗口均由 `ReqManager` 提供的 `req_idx` 定位。索引池化窗口不单独分配 checkpoint 槽位；其状态布局和恢复规则由各自实现描述。 - -### 3.3 分配状态与数据状态分离 - -allocator 是索引分配的唯一事实来源，资源对象持有存储和布局。缓存功能流程通过 radix 维护前缀关联、页/槽位的引用与发布关系；`AttentionStateManager` 组织所有参与 State 的数据操作；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 - -同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 - -对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，索引池化窗口资源管理增量计算需要的原始短窗口；二者的组合关系由`AttentionStateManager`维护。 - -### 3.4 按具体 attention 实现定义状态 - -定义只约束操作的 `BoundedAttentionState` 抽象基类，由 `KDAState`、`GDNState` 等子类负责自己的成员、布局、重置、快照与恢复语义。`AttentionStateManager` 初始化并组合所需对象，公共缓存流程通过管理器访问它们。不同模型也可以复用同一种状态实现。 - -`KDAState` 的成员定义如下： - -```python -class KDAState(BoundedAttentionState): - config: ConvSSMStateConfig # 仅包含算法布局和 TP 信息 - state_name: str | None # 构造时为 None，由管理器绑定模型角色时设置 - request_capacity: int # 包含 HOLD 请求槽位 - conv_state: torch.Tensor - ssm_state: torch.Tensor - mtp_state_index: torch.Tensor | None - mtp_step: int # 该实例的 MTP 布局参数 - big_pages_pool: KDAPagesPool | None - small_pages_pool: KDAPagesPool | None -``` - -State 直接持有按层组织的 Tensor，不额外设置通用状态管理对象。该示例中的 Conv、SSM 和 MTP 索引驻留 GPU；checkpoint pool 独立持有快照 buffer。设请求槽位容量为 `N`（包含保留槽位），该实例的 `mtp_step` 为 `S`： - -| 成员 | 内容与索引方式 | -| --- | --- | -| `conv_state` | 对应算法的卷积历史窗口，窗口宽度为 `conv_kernel_size - 1 + S`；用 `buffer[layer, req_idx, ...]` 访问 | -| `ssm_state` | KDA 的递推矩阵状态，每请求预留 `S + 1` 个状态槽位；第 `j` 个槽位位于 `buffer[layer, req_idx * (S + 1) + j, ...]` | -| `mtp_state_index` | 形状为 `[N]` 的 `int32` Tensor；记录各请求接受的 MTP 状态索引。`S == 0` 时为 `None` | -| 大小页 pool | 按需创建，保存该算法指定前缀位置的状态；槽位来自管理器的对应 `bounded_att_*_allocator` | - -`GDNState` 按 GDN 算法独立定义 Conv/SSM 和 MTP 状态，其大小页使用 `GDNPagesPool`，支持不同的 K/V head 数及维度。`KDAPagesPool` 按 KDA 的相同 Q/K/V head 数和方阵递推状态定义布局。两类 pool 各自分配快照存储，不额外记录长度或页面有效性；前缀与槽位的关联由缓存调用流程维护。两种 State 可以复用兼容的 Conv/SSM 数据操作，共同实现 `BoundedAttentionState` 接口；基类不规定 Conv/SSM 等成员，也不意味着两种算法的状态内容可以互换。 - -运行态按 `ReqManager` 的槽位容量预分配，状态对象接收 `req_idx` 执行操作。main/draft 可以持有同一种状态类的两个独立实例，也可以使用不同状态类型；各实例的层数、维度、dtype 和 MTP 配置由对应模型决定。allocator 仍由 `AttentionStateManager` 创建或绑定并统一管理。 - -后续 sliding-window attention 可以定义实现 `BoundedAttentionState` 的 `SlidingWindowState`，按窗口 KV、有效范围和游标等实际需求组织成员及自己的页池。它可以使用适用的 bounded-attention 分配域，其页布局、恢复边界和已接受状态选择须按窗口语义实现；不要求包含 Conv/SSM 字段，也不复用 KDA/GDN 的页池。统一操作接口见 6.7 节。 - -### 3.5 State 成员命名、资源身份与唯一性 - -`AttentionStateManager(request_capacity)` 只记录共享请求容量，三个 allocator 初始为 `None`，State 由各模型创建后通过 `register_state(model_role, state)` 绑定。`model_role` 为 `main` 或 `draft0` 到 `draft35`；接口根据 `state.get_state_type()` 生成规范成员名，并设置 `state.state_name`。各具体 full-attention 配置、KDA/GDN 的 `ConvSSMStateConfig` 和 `WindowCompressorStateConfig` 均不保存资源身份，绑定时保留原配置对象；这些 State 未绑定时的 `state_name` 为 `None`，资源名称只保存在 State 上。管理器持有传入的原对象，不复制 Tensor、不重建 State、不创建额外的 buffer 视图，meta State 和已分配运行态均可绑定。每个角色的同一方案只能绑定一次，同一个 State 对象不能分配给多个角色；`finish_state_registration()` 后不再添加 State。allocator 仍由后续容量与存储初始化流程设置。 - -管理器成员名与 `state.state_name` 完全一致，统一使用 `{model_role}_{算法及存储方案}_state`。模型角色为 `main` 或 `draft0` 到 `draft35`，最多 36 个 draft。资源类别由 `FullAttentionState`、`BoundedAttentionState` 类型表达，名称中不再重复 `full_att`、`bounded_att`。 - -两类 State 都必须实现 `get_state_type()`，只返回中间的算法及方案标记，使用小写 snake_case。方案标记区分量化位宽、group size、scale 方式/精度，KDA/GDN 分别标注 Conv 和 SSM 精度，K-pool 标注压缩比例与存储精度。层号、容量、设备、TP rank 和请求信息不进入名称。 - -同一角色、同一标记只能注册一个 State，包含该方案的全部逻辑层；不同存储方案可以共存，通过算法字符串、层归属及可选的 State 类别筛选。名字用于资源定位，不能替代完整布局兼容性检查。标记注册后保持稳定，meta 定义、实际存储和 CPU/DP/PD 进程使用相同名称，不提供旧名称别名。 - -| 方案 | State 成员名及 `state_name` | -| --- | --- | -| 普通 BF16 KV | `main_gqa_bf16_state` | -| INT8，group size 8，BF16 scale | `main_gqa_int8_group8_bf16_scale_state` | -| 静态 FP8，per-head scale | `main_gqa_fp8_static_per_head_state` | -| MLA BF16 | `draft0_mla_bf16_state` | -| 带 indexer 字段的 packed MLA FP8 | `main_mla_packed_fp8_group128_state` | -| KDA / GDN | `main_kda_conv_bf16_ssm_fp32_state` / `draft0_gdn_conv_bf16_ssm_fp32_state` | -| K-pool W4/S4 索引池化窗口 | `main_compressor_indexer_w4_s4_bf16_state` | - -普通 KV 使用 `gqa`，涵盖 MHA/GQA/MQA；MLA 使用 `mla`。`packed` 区分 KV 与 indexer 联合存储，FP8 packed MLA 还标注量化 group size。KDA/GDN 使用 `conv_*_ssm_*` 明确两类状态的精度；窗口压缩器使用 `compressor_{target}_wW_sS_*` 标明用途、窗口、步长与存储精度。完整 shape、标定方式等仍保存在 State 配置中。 - -以启用 draft 的 GLM BF16 配置为例，管理器的 State 成员是： - -```text -AttentionStateManager -├── full_att_allocator -├── bounded_att_big_page_allocator -├── bounded_att_small_page_allocator -├── main_mla_packed_bf16_state -├── main_kda_conv_bf16_ssm_fp32_state -├── main_compressor_indexer_w4_s4_bf16_state -├── draft0_mla_packed_bf16_state -└── draft0_compressor_indexer_w4_s4_bf16_state -``` - -当前 GLM draft 不使用 KDA，因此不创建 `draft0_kda_conv_bf16_ssm_fp32_state`。Qwen 的线性状态为 `main_gdn_conv_bf16_ssm_fp32_state`；Full Attention 按实际配置创建 `main_gqa_bf16_state` 等成员。MLA 与普通 KV 都可使用同一个 `full_att_allocator` 分配域。 - -`vanilla_no_att`、`eagle_no_att` 等仅执行 FFN 的 draft 不创建 KV State，也不初始化 attention backend；其计算引用可以为 `None`，仍共享管理器的 allocator 与请求映射。缺少实际需要的 draft attention State 则应在初始化时断言失败。 - -因此同一角色可以同时拥有 `main_gqa_bf16_state` 和 `main_gqa_int8_group8_bf16_scale_state`；这是两种不同方案，各自覆盖对应的层，聚合复制及传输必须覆盖二者。现有内置模型按配置选择自己的方案；未来混合量化模型在管理器中显式定义多个规范成员，基类通过聚合查询将它们传给相同 helper，不需要增加通用 State 注册表。 - -main/draft 必须拥有不同逻辑 State，不能将覆盖两侧层的同一个 State 同时命名为 main 和 draft。当前启动流程为各模型分别创建 State 和存储；GLM 的每个模型也分别创建 indexer ring。三种 allocator 的名称不带模型角色，同一分配域只初始化或绑定一个实例。 - -`AttentionStateManager`只持有实际启用的成员。当前配置选择的类型名可作为普通元数据保存，用于找到对应的规范字段；参与操作的集合由基类查询接口从显式成员生成，State 对象创建完成后可缓存只读 tuple 查询结果。Model 在资源层只持有 `att_state_manager`，不另存 `full_att_state` 或游标集合。main 在准备资源时创建唯一的管理器，所有 draft 从准备 State 开始就引用同一实例，并通过 `register_state` 添加各自的规范成员。每个 draft 使用独立的 `draft0`、`draft1` 等角色，最多 36 个（`draft0` 到 `draft35`）；即使方案相同，也分别持有 State，不合并层或运行态。State config 完整保存全局 head 数和 TP 分布规则，全部对象创建完成后再统一分配缓存存储。管理器的拥有成员必须使用完整规范名称。 - -`InferStateInfo` 保存 `att_state_manager`、`is_draft_model` 和资源角色 `att_state_role`，不额外绑定 `full_att_state`。各层直接访问管理器上对应方案的 State 成员，再通过局部引用取得层视图或写入 KV；FP8 backend 从具体量化方案的成员读取固定参数。每个 draft 的层由自身编号对应的 State 管理；`att_state_role` 为 `main` 或 `draft0` 到 `draft35`，不能仅根据 `is_draft_model` 选择资源。是否初始化 full-attention backend 根据管理器中该角色实际启用的资源判断；FFN-only 和 linear-only 模型不创建 full-attention backend，linear backend 由相应模型独立初始化。 - -资源身份用于 CPU/PD 资源布局与 DP 映射，所有传输按完整角色和方案匹配。部署配置确定各资源的页布局；PD 参数兼容性由 PD master 在连接阶段校验，CPU 服务绑定时核对页大小，磁盘缓存按 main/draft 模型版本隔离。 - -## 4. 目录设计 - -目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器使用 `manager.py`，两类公共 State 契约使用各自目录下的 `base_state.py`；allocator 和 pool 文件按对应类名命名，具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 - -```text -lightllm/common/ -├── basemodel/ -│ ├── basemodel.py # 模型基础流程与模型专属配置 hook -│ └── basemodel_utils.py # 配置归一化、KV 选项与 full/linear State 定义构建 -├── attention_state/ -│ ├── __init__.py -│ ├── utils.py # State 命名、传输视图尾维合并及搬运 dtype 选择 -│ ├── manager.py # AttentionStateManager：State 装配、容量与存储初始化、缓存和传输接口 -│ ├── types.py # TensorField / StateTransferLayout 等 -│ ├── allocators/ # 管理器直接持有，同一分配域共享实例 -│ │ ├── __init__.py -│ │ ├── full_att_kv_allocator.py # FullAttKVAllocator / KVCapacityReader：full_att 分配域及共享容量读取 -│ │ └── state_page_allocator.py # StatePageAllocator：bounded_att 大小页各一个实例 -│ ├── full_att/ -│ │ ├── __init__.py -│ │ ├── base_state.py # FullAttentionState：抽象行为接口 -│ │ ├── gqa/ # GQA 及其量化变体的配置、State、创建入口和传输实现 -│ │ │ ├── create_gqa_state.py -│ │ │ ├── gqa_state.py -│ │ │ ├── group_quantized_state.py -│ │ │ ├── group_quantized_transfer.py -│ │ │ ├── static_fp8_state.py -│ │ │ ├── static_fp8_transfer.py -│ │ │ └── gqa_transfer.py -│ │ ├── dp_transfer.py # 各 full-attention 方案的直接索引 DP 聚合复制 -│ │ ├── mla/ # 普通 MLA latent 配置、State 和创建入口 -│ │ ├── compressed_att/ # V4 C4/C128 压缩块配置、State 和创建入口 -│ │ └── packed_mla/ # packed MLA 及其 FP8/E8M0 变体的配置、State、创建入口和传输 -│ ├── bounded_att/ -│ │ ├── __init__.py -│ │ ├── base_state.py # BoundedAttentionState：有界状态的抽象操作契约 -│ │ ├── utils.py # 页池分配策略：DP fetch 使用共享 pinned 存储 -│ │ ├── transfer_utils.py # bounded_att_buffer_copies：生成快照字段的复制视图 -│ │ ├── copy_kernel.py # copy_bounded_att_tensors：聚合多页、多字段的 Triton 拷贝 -│ │ ├── conv_ssm_types.py # KDA/GDN 的 Conv/SSM 布局配置 -│ │ ├── gdn/ # GDN 配置、State、页池和创建入口 -│ │ ├── kda/ # KDA 配置、State、页池和创建入口 -│ │ ├── sliding_window/ # SWA 配置、State、创建入口和传输 -│ │ ├── window_compressor/ # 窗口压缩器配置、State、创建入口和传输 -│ │ └── window_pages_pool.py # SWA 与压缩器共用的时间有序窗口快照池 -``` - -`allocators/` 与两类 State 目录平级。`full_att_kv_allocator.py` 提供 full-attention 的 KV 分配器，对外使用 token 数量和索引，内部遵守整页分配约束；`state_page_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `BoundedAttentionState`。 - -`full_att/base_state.py` 只定义 `FullAttentionState` 行为接口；每种 full-attention 子 State 使用独立目录，目录内放置配置、State、创建入口及专用传输实现。`bounded_att/base_state.py` 定义有界运行态的公共接口；每种 bounded 子 State 同样使用独立目录，公共的页池工具和聚合 kernel 留在分类目录根部。布局或算法差异由各子目录自己的 `*_state.py` 实现，不设置另一套独立的 storage 公共继承体系。 - -`KDAState` 使用 `KDAPagesPool`，`GDNState` 使用 `GDNPagesPool`；每个 State 的 `big_pages_pool` 与 `small_pages_pool` 是对应 pool 类型的独立实例。两个 State 的 `initialize_page_pools()` 分别直接创建 `KDAPagesPool` 和 `GDNPagesPool`；管理器确定共享槽位容量后调用该公共方法，不直接构造某种通用 pool。State 页接口与管理器统一使用 `save_bounded_att_runtime_to_pages`、`restore_bounded_att_runtime_from_pages` 等名称，pool 不持有 allocator。K-pool 使用 W4/S4 的 `WindowCompressorState`，在整组边界无重叠历史，不创建大小页 pool。 - -`ConvSSMStateConfig` 描述 KDA/GDN 共享的 Conv/SSM 布局，不作为所有 bounded attention 的配置契约。KDA/GDN 构造函数的 `config` 明确标注为 `ConvSSMStateConfig`，入口通过 `assert isinstance(config, ConvSSMStateConfig)` 约束类型；配置只描述层、head、维度、dtype 和 TP，不包含 `state_name`。其余构造参数为 `request_capacity`、`mtp_step` 和 `device`；State 只保存一份请求容量和 MTP 步数，设备从 Tensor 获取。大小页池初始为 `None`，容量确定后通过 `initialize_page_pools(big_page_capacity, small_page_capacity, *, device="cpu", pin_memory=True)` 创建，容量保存在 pool 中。大小页恢复直接访问 Conv/SSM 基准存储切片，不为恢复操作构造字段字典或 `TensorField`。`get_layer_view(logical_layer_id)` 直接返回 `(conv_state, ssm_state)` 两个 Tensor 引用，MTP 游标与步数从 State 读取，不另外包装视图对象。`BoundedAttentionState` 不声明抽象构造函数，仅约定操作；构造参数由具体 State 和算法 create 入口明确标注，KDA/GDN 分别实现运行态、大小页、MTP 与传输方法，不再设置中间实现基类；新的算法可以定义自己的配置、层视图和页池。管理器统一使用 `get_bounded_attention_states(role=None)` 枚举参与对象；模型通过 `create_bounded_att_states(..., device="meta")` 返回 CPU 布局定义。KDA/GDN 在各自 State 中构建 CPU/PD 传输区域；其他 bounded 布局分别实现 `get_cpu_transfer_regions()` 和 `get_pd_transfer_regions()`，管理器使用同一组接口。 - -KDA/GDN 的大小页保存与恢复、CPU cache 读写、DP 大页与运行态复制、PD 打包与读回，统一调用 `copy_bounded_att_tensors()`。调用方将 Conv、SSM 分别收集为同形状批次；CPU/PD 按传输区域收集页面视图，Q/K/V 切片与 SSM 分别提交。入口按三个阶段顺序执行：公共 `utils.merge_contiguous_tail_dims` 合并连续尾维，公共 `utils.view_max_copy_dtype` 为整批选择最大搬运宽度视图，`_launch_copy_kernel` 上传指针并发射算子。公共函数只读取首对 Tensor 的 shape/stride，从末轴向前扫描一次，判断两端共同连续的尾部区间，再调用 `view` 创建整批视图。普通轴须满足外层 stride 等于已合并尾部元素数乘尾部元素 stride；长度为 1 的轴不约束 stride，内侧仅有一个元素时改用当前非单元素轴的 stride。这样可正确保留 padding，也支持末轴步长不为 1 的视图。不复制数据，不依赖 PyTorch 内部 API，不以异常作为正常回退路径，也不使用可能复制数据的 `reshape`。`start_dim` 限制允许参与合并的最前轴，例如 `[layer, window, head, dim]` 使用 `start_dim=2` 保留环形 window，只合并 head/dim；`start_dim=-2` 按两侧末尾两维对齐，允许前缀维数、token 容量不同。默认直接减少维数；`keepdim=True` 将被合并的前面各轴保留为 1，使元素集中到末轴，供固定维数的 token/page kernel 使用。实现仅依赖公开的 Tensor shape、stride 和 view 接口，不再需要针对 PyTorch 内部接口的版本断言。每批只包含一种 shape 和 dtype，只从最后一维向前合并两端共同满足 stride 连续关系的尾部维度，遇到任一侧无法合并就停止，保留前面的维度；尾部单元素轴按上述规则处理，再按行宽选择候选类型，通过首对布局过滤及整批 `view` 校验尝试创建最大 8 字节的整数视图，宽类型转换后保留维数，即使末轴长度变为 1 也不再降维，最后按降维后的视图选择算子，不做形状或维数分组，每批默认 `grid_num=24`、`BLOCK_ELEMENTS=1024`、`num_warps=8`、`num_stages=1`，通过 `tl.range(..., num_stages=NUM_STAGES)` 跳步循环覆盖全部数据；源、目标指针分别上传为两个独立 Tensor，1～6 维分别使用独立算子，每个算子按“计算索引和偏移、读取源/目标指针、load 数据、store 数据”的顺序直接展开，整批共享一份 `SHAPE` 和 `DTYPE`，`total_element_count` 指定整批传输单元总数（转换后视图的 numel × 复制项数），算子按全局元素索引计算复制项编号及项内位置，源、目标分别共用一份 stride 编译期常量；入口仅检查首尾两对 Tensor 的 shape、dtype 及同侧 stride 一致，不分配状态数据中间 Tensor。按转换后视图的元素宽度选择 `uint8/uint16/uint32/uint64` 搬运，所有转换均为不分配数据的 view，保留原始位模式及 padding。运行时参数统一位于 `tl.constexpr` 参数之前，支持 MTP 窗口裁剪、TP 切片和外部轴顺序变化。ready 页在构建复制列表前跳过。游标重置仍在同一流中独立提交。 - -默认参数在 H200 上固定 grid=24 联合调优，覆盖 Conv BF16 / SSM FP32、三种规模和 CUDA↔CUDA / CUDA↔pinned CPU / pinned CPU↔pinned CPU，共 24 个场景。第一轮比较 BLOCK_ELEMENTS=512/1024/2048/4096 与 warps=1/2/4/8，并保留 256/8/1 对照；第二轮为入选组合比较 num_stages=1/2/3/4，两轮共测试 29 个不同组合。每个候选均校验数据，预先创建指针表，以随机顺序进行 CUDA Graph 计时并取中位数。按场景等权的速度几何平均选择 1024/8/1；相比旧默认 grid=16、BLOCK_ELEMENTS=256、warps=8、stages=1，第二轮综合速度约为 1.58 倍。该组合并非每项最快，小块场景最多增加约 1.35 μs（Conv CPU→GPU 约 2.42→3.77 μs）。`grid_num`、`num_warps`、`num_stages` 可由调用方覆盖；流水深度用于 `tl.range` 的循环属性。 - -另外使用不同地址的独立块连续发射复测：Conv 为 64×3 MiB，SSM 为 16×64 MiB，块间不做同步，九轮取平均。下表为有效载荷带宽（GB/s），不含 Python 准备、指针上传和网络传输；旧默认与新默认在同一轮实验对照。未测量与推理计算并发时的性能影响。 - -| 数据类型 | 方向 | 旧默认 16/256/8/1 | 新默认 24/1024/8/1 | -| --- | --- | ---: | ---: | -| Conv | GPU→GPU | 16.54 | 59.78 | -| Conv | GPU→CPU | 16.74 | 40.73 | -| Conv | CPU→GPU | 6.22 | 26.70 | -| Conv | CPU→CPU | 19.51 | 36.32 | -| SSM | GPU→GPU | 76.83 | 303.01 | -| SSM | GPU→CPU | 50.61 | 51.99 | -| SSM | CPU→GPU | 22.51 | 50.11 | -| SSM | CPU→CPU | 19.91 | 36.88 | - -CPU 页池和 CPU cache 页必须是当前进程已 pinned/registered 的存储；两端均为 CPU 时也由 GPU kernel 访问 mapped host memory。所有视图直接使用 `tensor.data_ptr()`，调用方保证该地址对当前 GPU 可访问；跨 GPU 访问依赖调用方已有的 peer/IPC 映射，拷贝入口不建立设备访问关系。所有操作在当前 CUDA stream 排队，调用方负责生产流依赖、存储及页面引用生命周期，以及 CPU 读取/复用前的完成事件；拷贝入口不查询主机映射、不调用 `record_stream()`，也不再依赖 CPU→CPU `copy_` 的同步语义。 - - -`FullAttentionState`、`BoundedAttentionState` 的具体子类直接实现 CPU、DP、PD 读写；无重叠的窗口 State 在 CPU 对齐边界不需要快照；IPC 导出前由管理器统一创建浅拷贝并排除进程私有资源，State 不设置额外的序列化钩子。各具体 State 通过 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` 选择传输字段、形状及 TP 切片；字段字节对齐、区间计算和 `TransferRegion` 构造直接在 State 内展开，不再通过 `region()` 包装。各具体 full-attention State 对应的 `*_transfer.py` 由 Triton 直接按 token 和页索引搬运 KV 及配套量化字段，每个字段一次 launch 覆盖全部页，不创建 gather/scatter 中间 KV。CPU 页必须是 pinned 或已通过 `cudaHostRegister` 注册的内存；调用方在 stream/event 完成后才能读取或复用页面。每种 full-attention 方案的 CPU/PD 路径共用本方案的 `transfer_*_pages` 入口和 kernel，直接按 token 和页索引读写外部页。两条路径均兼容不足整页的有效 token；`bounded_att/transfer_utils.py` 复制完整状态快照，不接收 token 数量；`bounded_att/window_compressor/window_compressor_transfer.py` 接收本地 ring 和 prefill 前缀长度，单次 Triton launch 搬运外部快照所需的逻辑窗口；两个方向均只处理逻辑窗口，并将短前缀对应的窗口内位置补零。bounded-attention 的传输函数使用 `page_buffers` 和必传的 `page_index`，不带 CPU/PD 前缀；页存储必须至少为二维，第一维表示页，统一通过 `page_buffers[page_index]` 取页。PD State 与 manager 读写入口统一接收 `pd_page_buffers` 和必填的 `pd_page_index`，将页集合与真实页索引直接传给共用搬运函数；不接受单页 Tensor，也不在入口添加页维度。CPU 专属读写接口仍使用 `cpu_page_buffers` 和 `cpu_page_indices`。Full attention 的各方案 CPU/PD 共用入口直接提交 Triton kernel；bounded attention 的复制视图按字段或传输区域收集为同形状批次，提交专用 kernel；窗口压缩器直接提交环形窗口 kernel，全部层的 key/gate 窗口搬运在一次 launch 中完成。取页和外部字段视图在各方案中直接展开：先选择页，再按字节起止位置切片并解释为字段 dtype/shape，不额外逐区域断言边界、对齐或轴定义。`utils.view_max_copy_dtype` 先按末轴 `size * element_size` 是否能被 8、4、2 整除，直接选择对应的候选 tuple：8/4/2、4/2 或仅 2 字节；例如行宽 12 字节只尝试 INT32、INT16，不尝试 INT64。使用各 State 正常分配的存储，不额外遍历或检查数据指针；任一侧末轴 stride 不为 1，或元素已占至少 8 字节时直接返回原视图，先用首对 Tensor 的外层 stride 和 Storage 偏移排除必然失败的候选，其余切片的偏移继续交给 `view(dtype)` 校验，仅捕获创建视图的 `RuntimeError` 后整批退回较窄 dtype，不复制数据。该函数由 bounded 聚合拷贝、full-attention CPU/PD、DP token 拷贝及 K-pool 共用；每个同布局批次统一选择 dtype，不同字段仍按各自编码处理。五种 full-attention 的 CPU/PD 页传输和 DP token 复制，在 TP 切片后先调用 `merge_contiguous_tail_dims(start_dim=-2, keepdim=True)`，只合并两侧均连续的 head/width；合并后的 head 轴保留为 1，继续使用原有四维/五维 kernel，再选择最大搬运 dtype。任一侧存在 head 间隙时沿用原布局。layer/token/page 等独立寻址轴及两侧各自容量保持不变；K-pool 的尾部只有 channel 一维，window 仍独立，不再跨 window 合并。full-attention 在各方案的 `*_transfer.py` 中直接应用 TP 切片，kernel 按索引长度裁剪有效 token；bounded-attention 在本地 `transfer_utils.py` 直接展开字段定位、轴对齐和本地/外部 TP 切片，不再通过 `_copy_pair` 跳转。压缩器使用固定 `window_size` 的逻辑窗口，按传入的前缀长度把逻辑位置映射到本地 ring 槽位；短前缀在外部窗口前部补零，读回时保留逻辑窗口之外的本地 MTP 槽位。full-attention 和 bounded 路径检查实际复制视图的 shape 一致；窗口搬运检查 layer/channel 相同，window 按两侧各自容量作环形映射。读回时反向执行同一组轴映射，保留本地和外部的固定角色。full-attention DP 直接复制同一 TP 分片的显式存储，原根目录 `transfer_utils.py` 的字段映射复制已无调用方并删除。State 的 KV 写入与量化方法直接调用相应计算模块中的 kernel。具体方案标记统一由 `get_state_type()` 返回，不再保留无读取方的 `layout_name`。这些聚合操作由 `AttentionStateManager` 提供默认实现，统一组织多资源读写、CPU 联合页、调用方传入的长度和位置映射；模型子类仅覆盖存在差异的部分。网络通信、任务队列与完成通知由 server 层管理。 - -CPU 服务通过模型配置与 meta device 上的 State 描述生成页格式，无需分配 GPU 存储。服务端将 main/draft 的 meta State 装配到管理器并标记创建完成，复用 `AttentionStateManager.get_cpu_page_layouts()`，与推理端使用同一条布局聚合路径。该接口直接遍历 full-attention、bounded State，调用各自的 `get_cpu_transfer_regions()`，同时平移字段区间的起点和终点；每个 State 的最大字段终点按 8 字节对齐后作为下一 State 的起点。CPU 页容量由全部布局的最大 `byte_end` 按 16 字节对齐得到，TP 副本不重复占用外部页；启动绑定时核对页字节数。 - -无需模型特有初始化的场景直接使用 `AttentionStateManager`，main 和各 draft 通过自身的准备流程向同一个管理器添加 State。它可以同时管理 full attention 与 bounded attention（包括压缩窗口），不按 main 的 attention 类型再分一层管理器。构造函数只接收 `request_capacity`；各模型通过 `register_state` 绑定已创建的 State，allocator 在后续容量与存储初始化流程中设置。 - -模型特有的 State 组合由 `model.py` 中两类创建接口确定，模型专有配置及派生参数放在各自的 `config.py`。Qwen3Next 创建 GDN，GLM 创建 KDA 和 K-pool 窗口；main/draft 分别返回各自的 State。所有模型共用 `AttentionStateManager` 的聚合数据操作，不再保留 Qwen3Next/GLM 专用管理器及旧 `bind_storage()` 流程。 - -token 复制、bounded 运行态重置、大小页保存/恢复、MTP 提交直接在 `AttentionStateManager` 中遍历对应 State，不再增加仅转发调用的 helper。CPU/DP/PD 传输组织也直接在 `AttentionStateManager` 的现有接口中实现，通过聚合查询选择参与 State；不再通过独立 helper 转发 State 集合及 allocator，模型子类仍可按特有需求覆盖操作。显式成员决定 State 归属，不另外维护 `_full_states`、`_bounded_states` 或 `_states` 注册表，也不按 `id(state)` 静默去重。基类允许在对象创建完成后缓存从这些成员派生的查询结果，缓存不参与资源装配。传输按 state_name 匹配时可构建临时映射，它不决定模型的资源组织方式。 - -#### 按逻辑算法创建 State - -每种逻辑算法提供一个公开的 `create` 入口。量化属于该算法的存储选择，不按量化类型拆分创建入口；调用方提供逻辑配置及编码参数，由入口返回具体 State。 - -| 逻辑类型 | 创建入口 | 参数决定的实现 | -| --- | --- | --- | -| GQA（含 MHA/MQA） | `full_att.create_gqa_state(config, *, kv_type, page_size, device, ...)` | 普通浮点、INT4/INT8 分组量化、per-tensor/per-head 静态 FP8 | -| MLA | `full_att.create_mla_state(config, page_size, device)` / `create_packed_mla_state(config, page_size, device)` | 普通 latent MLA 与带 indexer 的 packed MLA 分开创建 | -| 序列压缩 attention | `full_att.create_compressed_att_state(config, page_size, device)` | FP8/E8M0 packed 压缩 KV 与可选 indexer；V4 CSA/C4、HCA/C128 由模型分别配置 | -| KDA | `bounded_att.create_kda_state(config, request_capacity, mtp_step=0, *, device)` | KDA Conv/SSM 运行态 | -| GDN | `bounded_att.create_gdn_state(config, request_capacity, mtp_step=0, *, device)` | GDN Conv/SSM 运行态 | -| SWA | `bounded_att.create_sliding_window_state(config, request_capacity, mtp_step, device)` | 近期原始 KV 窗口，当前采用 FP8/E8M0 packed 编码 | -| 窗口压缩器 | `bounded_att.create_window_compressor_state(config, request_capacity, mtp_step, device)` | 按窗口长度和产出步长管理候选/gate，KV 与索引 key 池化分别创建实例 | - -每个算法的创建入口独立放在对应的 `create_*.py` 中，通过各分类包导出，包括 GQA、MLA、压缩历史、KDA、GDN、SWA、窗口压缩器。K-pool 复用窗口压缩器入口。一个算法的普通精度与量化实现共用同一 create 入口，由参数选择具体 State；State 实现文件不负责选择其他实现。不增加工厂类或注册表。入口不读取请求信息，不绑定角色，不创建 allocator；返回对象后由模型调用 `register_state`。 - -GQA/MLA 的创建入口分别接收 `GQAAttentionConfig`、`MLAConfig` 或 `PackedMLALogicalConfig`，每种具体 State 使用独立配置，不继承通用配置。GQA 的 `kv_type` 使用现有编码名 `"None"`、`"int4kv"`、`"int8kv"`、`"fp8kv_spt"`、`"fp8kv_sph"`；`group_size` 仅对 INT4/INT8 生效。静态 FP8 构造函数和创建入口不接收标定参数，`scales`、`q_scales` 初始均为 `None`；全部 State 就绪后，统一加载流程只能通过 `bind_calibration()` 初始化标定。各 DP 使用相同的标定配置，不额外生成或校验标定内容指纹。普通 MLA 只描述 latent；packed MLA 只描述带 indexer 的逻辑 format，工厂内部再转换为具体物理布局。 - -MLA 的 packed FP8 编码由 `fp8_format` 显式指定：默认 `"group128_fp32"` 使用原有 -`PackedMLAFullAttentionState`；`"group64_e8m0"` 使用独立的 `FP8E8M0PackedMLAState`， -其配置为 `FP8E8M0PackedMLAConfig`。创建入口不通过 `latent_dim == 512` 等维度条件推断格式。 -GLM main/draft 在各自的模型工厂中传入 `fp8_format="group64_e8m0"`；未量化时仍返回普通 packed State。 - -`FP8E8M0PackedMLAState` 固定使用 448 维 FP8、64 维 BF16 和七个 E8M0 scale（补齐为八字节）， -KV 共 584 字节，随后保存 indexer key 及一个 FP32 scale；最后 64 维保留 BF16 不代表使用 RoPE。 -配置独立定义，不继承其他方案的配置；只记录逻辑层分组、indexer 维度和 TP 坐标。 -固定 latent 维度、单 head、副本规则和 uint8 编码由实现确定，不作为重复输入。 -`token_bytes` 是含 16 字节对齐 padding 的物理行宽，`payload_bytes` 是不含行尾 padding 的有效字节数。 -例如 indexer 维度为 128 时，GPU 行宽为 720 字节，CPU/PD 每个 token 只传 716 字节。 -State 将连续的 KV/scale/indexer 字节作为一个传输区域，复用 `transfer_packed_pages` 的索引拷贝 kernel； -同卡 token 复制与 DP 迁移仍复制完整物理行。原 packed State 及其传输实现无需识别 E8M0 编码。 -定制量化写入放在 `fp8_e8m0_packed_mla_write.py`，State 不依赖 GLM 模型目录中的实现。 - -各模型的两类 State 工厂负责构造配置并调用对应算法入口；`lightllm/utils/cpu_cache_layout.py` 复用模型工厂创建 meta State，生成 CPU cache 联合页布局，不加载标定。full-attention 的实际存储由管理器后续统一分配；GPU 启动时 bounded 按请求容量直接创建运行态，大小页池另行初始化。GPU 启动与 CPU cache 元数据使用同一套算法创建逻辑。 - -```python -state = create_gqa_state( - config=full_config, kv_type="int4kv", group_size=8, - page_size=16, device="meta", -) -manager.register_state("main", state) -# 全部 State 准备完成后，由 manager 统一规划容量并分配存储。 -``` - -根目录的公共帮助代码收敛到 `utils.py`，该文件负责 State 角色与精度命名及传输视图的尾维合并、整批最大搬运 dtype 选择，测试专用角色视图和成员筛选帮助函数放入测试辅助文件；full-attention 的 CPU/PD 页复制放在各方案的 `*_transfer.py` 中，bounded 的页复制放在各自目录的 `transfer_utils.py` 中；CPU/DP/PD 传输及 State 聚合循环直接放在 manager 中。共享资源初始化由 `AttentionStateManager.initialize_storage()` 统一组织。各文件用中文大分块注释区分职责，公共函数在前、私有辅助在后；索引直接使用调用方提供的有效值，仅在需要列表的操作处展开 Tensor 转换，不再集中重复校验索引类型、范围和唯一性。共用的 Conv/SSM 布局配置放在 `conv_ssm_types.py`，由 `bounded_att` 包统一导出；该文件不包含 State 操作实现。具体 State、具体页池与公共数据类型保持各自的定义边界。 - -管理器聚合调用全部参与 State：CPU 页包含全部参与 linear State，全部页面读回后恢复目标压缩窗口；IPC 导出不携带 allocator 所有权。调用方在全部数据操作完成后发布页面，失败时不发布。radix 和 CPU cache 通过`AttentionStateManager`操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。 - -`types.py` 仅放跨资源使用的纯数据类型，包括 `TensorField`、`AxisMapping`、`TransferRegion`、`StateTransferLayout` 和 `PageKind`。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求两类 State 继承同一个大接口。 - -## 5. allocator 与状态页 pool - -### 5.1 allocator 的公共入口与按需初始化 - -| 公共入口 | 类型 | 初始化条件 | -| --- | --- | --- | -| `full_att_allocator` | `FullAttKVAllocator \| None` | 组合中存在需要页式 KV 分配的 full-attention 资源 | -| `bounded_att_big_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 大页 | -| `bounded_att_small_page_allocator` | `StatePageAllocator \| None` | 启用 bounded 小页 | - -`AttentionStateManager`初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 - -allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由`AttentionStateManager`明确描述。 - -大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 - -`FullAttKVAllocator` 的申请量和返回索引均以 token 表示，分配和回收须满足整页约束；`StatePageAllocator` 返回状态页槽位 ID。GPU 运行态（包括压缩窗口）使用请求槽位，不使用上述状态页分配器。 - -HOLD 页位于普通 KV 分配域之后，由 full-attention 存储预留。初始化流程将 allocator 的 `hold_token_indices` 显式交给 `ReqManager.bind_hold_page` 建立 padding 请求映射；ReqManager 不持有 allocator，也不负责释放 KV。`FullAttKVAllocator` 构造时直接创建 `shared_can_use_token_num: SharedInt` 并发布初始容量，在 `alloc`、`free` 和 `free_all` 中直接更新。对应的 `KVCapacityReader` 定义在同一文件中，router 按本地 DP 分组读取共享计数；Reader 不重置计数。读写双方使用同一服务前缀与 `rank_in_node` 命名，因此构造 allocator 前须完成服务名和 rank 初始化，无需外部绑定或容量回调。 - -`FullAttKVAllocator` 沿用 token 索引栈 `mem_state`、`mark_start/mark_end` 游标和循环返回缓冲区；每次分配复制索引到预分配的返回缓冲区，避免逐次申请 pinned memory。分配和归还均以整页为单位，归还时保持每页内部 token 索引连续且有序，页之间可以任意排列。调用方负责资源所有权、单次归还及索引的类型、范围与页内顺序；allocator 只保留容量与页对齐、int32 索引上限及分配/归还数量边界的断言，不逐元素校验、不另建逐页分配位图，也不对归还索引排序。返回缓冲区按原有三倍容量循环复用，调用方须在缓冲区复用前完成异步索引读取；需要长期保留索引时自行复制。 - -### 5.2 管理器、allocator 与 pool 分工 - -CPU 大小页池通过 `bounded_att/utils.py` 的 `create_page_pool_tensor()` 统一选择分配方式。未启用 `enable_dp_prompt_cache_fetch` 时沿用 `torch.empty(..., pin_memory=True)`；启用后直接创建共享 Storage 并原地调用 `cudaHostRegister(Portable | Mapped)`，避免先申请普通 pinned 页池再转共享存储造成整块复制。显式关闭 pin_memory 时只创建共享 CPU 存储，GPU 存储保持原分配方式。注册随本地 Storage 生命周期管理，最后一个视图释放时等待 GPU 使用结束再注销；接收进程需要自行注册其映射地址。KDA/GDN 页池均提供 `register_pin_memory()`，复用 `register_page_pool_tensor()` 注册 Conv/SSM 的整个 Storage；已 pinned 或空存储重复调用为空操作。`create_dp_ipc_copy()` 保留共享 CPU 页池，`create_pd_ipc_copy()` 始终排除全部大小页池。DP 接收端设置本地设备并读出资源后，调用 manager 的 `register_page_pools_pin_memory()`，统一注册所有 bounded State 的大小页池。`Portable` 使本进程其他 CUDA context 也识别其 pinned 属性，`Mapped` 建立设备可访问映射；它们不会替另一个进程注册地址。 - -`StatePageAllocator` 只维护可用槽位。`KDAPagesPool`、`GDNPagesPool` 等具体页池只负责分配和持有各自的状态 buffer；State 使用外部传入的槽位完成保存和恢复。`AttentionStateManager` 持有 allocator 并提供跨 State 的聚合数据操作，调用方直接使用 allocator。操作顺序如下： - -1. 调用方直接从 `bounded_att_big_page_allocator` 或 `bounded_att_small_page_allocator` 分配槽位。新分配的槽位不代表已有可复用内容，必须先完成保存或载入才能发布。 -2. `save_bounded_att_runtime_to_pages` 聚合调用参与 State 保存；异步操作通过 stream/event 保证后续读取顺序。保存失败直接向上传递异常，由调用方停止发布并处理槽位生命周期。 -3. `restore_bounded_att_runtime_from_pages` 按调用方选定的槽位恢复全部参与 State，调用方保证前缀关联正确并已满足同步条件。 -4. 调用方确认引用已解除且异步使用完成，直接向同一个 allocator 归还共享槽位一次；pool 中的数据可保留至下一次写入。 - -CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排除在大页 allocator 的普通分配范围外；参与操作的 main/draft 大页 pool 分别为这些槽位提供存储。成员名为 `bounded_att_cpu_load_big_page_slot`（从 CPU 加载末尾不足一页的状态时使用）和 `bounded_att_cpu_offload_big_page_slot`（为卸载到 CPU 预留，当前尚无调用点），均属于 bounded-attention 大页池。 - -`AttentionStateManager`声明数据可捕获及可恢复的条件，并实现保存和恢复操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 直接操作 pool 中的快照数据。 - -### 5.3 前缀长度与状态页的关联 - -前缀长度与大小页槽位的关联由请求及缓存流程维护。`big_pages_pool` 和 `small_pages_pool` 仅保存算法快照，不重复持有长度数组或页有效性标记。 - -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 - -调用方只能发布已完成全部参与 State 保存的页面，并通过 stream/event 保证异步操作顺序。尚未捕获的中间边界不能发布为可复用快照；失败时不发布对应页面，等待已提交操作结束后再回收槽位。前缀身份、引用与就绪状态均由调用流程管理。 - -## 6. 对外接口与 attention 资源暴露方式 - -`AttentionStateManager` 负责共享资源初始化及缓存和传输操作，模型工厂决定 State 组合。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，`AttentionStateManager` 组织内部 State 完成操作。 - -### 6.1 FullAttentionState：公共接口与不同 storage 的完整处理 - -本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttentionStateManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 - -#### 逻辑层与物理存储层 - -具体 State 配置中的 `logical_layer_groups: tuple[tuple[int, ...], ...]` 是 full-attention 层映射的唯一配置来源。GQA/MLA 分别使用各自的独立配置，序列压缩历史使用独立的 `CompressedAttentionConfig`，不要求所有配置继承同一基类。外层每项代表一份物理层存储，顺序对应 K/V、scale 和 packed indexer 等 Tensor 的 layer 轴；内层列出共享这份存储的逻辑层号。每组必须非空，同一个逻辑层不能重复出现，层号不要求连续。 - -```python -logical_layer_groups = ( - (0, 4), # 逻辑层 0、4 共用物理层 0。 - (8,), # 逻辑层 8 使用物理层 1。 -) -# 不共享时使用 ((0,), (4,), (8,))，对应三份物理层存储。 -``` - -层命名统一使用 `logical_` 和 `storage_` 区分含义：`logical_layer_groups` 描述共享分组，`logical_layer_id` 是推理层传入的逻辑层号，`storage_layer_index` 是 Tensor layer 轴上的物理下标。State 的 `_logical_to_storage_layer` 保存两者的对应关系。KDA/GDN 和 K-pool 当前仍是一对一存储，配置使用 `logical_layer_ids`；两类 State 的 `get_layer_view(logical_layer_id=...)` 使用相同参数语义。FP8 标定的 `storage_layer_count`、`storage_layer_start` 表示物理行数和偏移，不能传逻辑层号。Tensor 轴标签 `"layer"` 保留，表示存储层轴。 - -静态 FP8 的 FA3/FlashInfer 后端通过 `AttControl.logical_layer_id` 接收当前逻辑层号，再由 State 的 `get_storage_layer_index()` 查询标定行，不通过 KV 指针反查层号。该接口属于 `FullAttentionState` 的公共契约，所有具体实现按自己的层映射返回物理行；层视图和 KV 写入也使用此接口。 - -State 根据分组展开逻辑层到物理下标的映射。上述配置的 `get_layer_view(logical_layer_id=0)` 与 `get_layer_view(logical_layer_id=4)` 返回同一份存储视图；`logical_layer_id=8` 返回另一份。分配、单 token 字节数、CPU/PD 布局和静态量化标定行数均按 `len(logical_layer_groups)` 计算，不按组内逻辑层总数计算，不重复保存物理层数或另一份配置映射。 - -同组逻辑层共享该 State 的全部配套字段，必须使用兼容的布局和参数。分组只表达数据共享关系，不指定写入者或执行顺序；模型负责跳过共享消费者的 KV 写入，避免覆盖生产层数据。普通模型每个逻辑层单独一组；Gemma4 将后部共享层加入同类型的最后一个非共享层所在组。推理层按自己的逻辑层号访问，CPU 服务与模型使用同一套配置构建物理分组。DP/CPU/PD 传输按相同的物理层顺序搬运，main/draft 仍各自拥有独立 State。 - -#### State 层归属与模型装配 - -两类 State 均提供 `contains_logical_layer(logical_layer_id) -> bool`，通过自身的逻辑层映射判断归属，不要求存储已分配。full-attention 的共享消费者和生产层都属于同一个 State；调用方不展开 `logical_layer_groups`，也不选择物理层。 - -`AttentionStateManager.get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` 通过三部分选择资源：`model_role` 按完整角色前缀匹配，避免 `draft1` 匹配 `draft10`；`state_type_filter` 在 `get_state_type()` 返回值中做字符串包含匹配，如 `gqa`、`gqa_int8`、`mla_packed`、`compressor_indexer_w4_s4`；`contains_logical_layer()` 判断层归属。可用 `state_class` 进一步限定资源类别；通用推理模板传 `state_type_filter=""`，仅通过角色、逻辑层和类别选择。最终必须恰好匹配一个 State，不存在或多个匹配均断言失败。 - -该单对象查询使用 `lru_cache(maxsize=None)`，要求 `_state_registration_finished=True`，与聚合查询一样只能在全部 main/draft State 创建完成后调用。实际存储可以尚未分配，后续原地初始化不改变缓存中的 State 对象。两类模型工厂只返回 State 列表，`prepare_attention_states()` 统一注册，不查询或绑定推理层资源。完成全部 State 创建和统一分配后，推理层直接调用 `get_state()`；首次按角色、逻辑层号和资源类别匹配，后续相同参数命中缓存。无需单独的 State 绑定阶段。 - -推理层与 indexer 都不保存 `full_att_state_name`，也不提供 `bind_full_attention_state()`。两者通过推理上下文的模型角色、自身逻辑层号及 `FullAttentionState` 类别取得同一资源；同层窗口 State 不会被误选。共享消费者是否写入 KV 仍由具体模型的推理层控制。 - -#### 继承与内部 storage 的关系 - -建议以 `FullAttentionState` 作为公共接口，由具体布局子类实现差异。子类内部的 storage 持有 buffer 和布局信息，operator 可以复用兼容的写入、量化及复制 kernel。布局选择集中在具体子类及其配套实现，公共调用方使用稳定接口。 - -```text -FullAttentionState -├── GQAFullAttentionState # MHA / GQA / MQA 共用的 KV 存储 -├── GroupQuantizedFullAttentionState # INT4 / INT8 等 -├── StaticFP8FullAttentionState -├── MLAFullAttentionState -├── PackedMLAFullAttentionState # 普通精度与 group128/FP32-scale 的 packed 布局 -├── FP8E8M0PackedMLAState # 独立 group64/E8M0 packed MLA，供 GLM main/draft 使用 -└── CompressedAttentionState # V4 的 C4/C128 压缩块 -``` - -所有具体 full-attention State 必须直接继承 `FullAttentionState`，具体实现之间禁止继承，也不通过调用其他具体 State 的方法复用行为。每种方案在自己的文件中完整实现配置初始化、容量计算、层视图、写入、token 复制以及 CPU/DP/PD 操作，防止某个方案的修改隐式改变其他方案。公共复用限于明确的工具函数和底层 kernel，不引入带具体存储行为的中间基类。配置也不建立继承关系：普通 GQA 使用 `GQAAttentionConfig`；分组量化使用 `GroupQuantizedAttentionConfig`，以 `scale_dtype` 明确 scale 精度；静态 FP8 使用 `StaticFP8AttentionConfig`，不要求调用方传入固定的物理 dtype。普通 MLA 使用 `MLAConfig`，packed MLA 使用带 indexer 的 `PackedMLALogicalConfig`；工厂再转换为具体 State 的物理配置。具体 packed 配置使用 `kv_encoding` 显式选择 `fp16`、`bf16` 或 `fp8_group128`，`dtype` 由编码提供只读 property。FP8 行的 uint8 承载类型只在存储分配函数中确定，不提供 `storage_dtype` 配置或属性；物理行宽及 indexer 偏移由编码推导。MLA 固定单份 latent 并在 TP rank 间复制，不提供 GQA 的 head 数或 `replicated` 开关。 - -`full_att/gqa/gqa_state.py` 中的 `GQAFullAttentionState` 使用广义 GQA 命名。设 TP 切分前的 Query 头数为 `H_q`、KV 头数为 `H_kv`，每个 KV 头对应一组 `H_q / H_kv` 个 Query 头，要求 `H_q` 能被 `H_kv` 整除。MHA 是 `H_kv == H_q`、每组一个 Query 头的特例；MQA 是 `H_kv == 1`、全部 Query 头共享同一组 KV 的特例；中间情况为多组 Query 分别共享 KV。 - -该 State 只保存 KV，`config.head_num` 表示当前 TP rank 切分或复制后持有的 KV 头数。K 和 V 各自独立分配一个 `[layer, token_slot, local_kv_heads, head_dim]` buffer，不再保存 `heads_per_kv`。写入接口的输入仍按 K heads、V heads 前后排列，State 将两部分分别写入各自存储；计算 backend 直接取得 K/V 视图。Query 头到 KV 头的分组关系由 attention backend 处理，因此三者共用同一套 State 接口与存储实现。 - -#### 必须提供的公共接口 - -```python -# ViewT 为具体布局定义的类型，下面是公共签名草案；构造参数由各具体 State 定义。 -class FullAttentionState: - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - @abstractmethod - def get_state_type(self) -> str: ... - - def allocate_storage(self, token_capacity: int) -> None: - ... - - def get_token_bytes(self) -> int: - ... - - # ======================================================================== - # 二、Full attention：计算视图、KV 写入与完整 token 复制 - # ======================================================================== - def contains_logical_layer(self, logical_layer_id: int) -> bool: ... - - def get_storage_layer_index(self, logical_layer_id: int) -> int: ... - - def get_layer_view(self, logical_layer_id: int) -> ViewT: - ... - - @abstractmethod - def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: - ... - - def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: - ... - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - raise NotImplementedError - - # layout 描述该 State 在外部页中的位置及本地轴到页内轴的映射。 - def write_to_cpu( - self, token_indices, cpu_page_buffers, cpu_page_indices, - *, layout: StateTransferLayout, cpu_page_ready_flags: List[bool] | None = None, - ) -> None: - raise NotImplementedError - - def read_from_cpu( - self, cpu_page_buffers, cpu_page_indices, token_indices, - *, layout: StateTransferLayout, - ) -> None: - raise NotImplementedError - - # ======================================================================== - # 四、DP 迁移 - # ======================================================================== - # 其他 DP 的同一 TP 分片到当前 State；调用方准备各自设备上的 CUDA 索引，目标为 int64。 - def copy_tokens_from_dp( - self, src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, - remote_dp_full_att_state: FullAttentionState, - ) -> None: - raise NotImplementedError - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - raise NotImplementedError - - # pd_page_buffers 为页集合，pd_page_index 选择物理页；layout 区域偏移相对该页起点。 - def write_to_pd_page( - self, token_indices, pd_page_buffers, - *, pd_page_index: int, layout: StateTransferLayout, - ) -> None: - raise NotImplementedError - - def read_from_pd_page( - self, pd_page_buffers, token_indices, - *, pd_page_index: int, layout: StateTransferLayout, - ) -> None: - raise NotImplementedError -``` - -“必须提供”表示每个具体布局直接实现相同的抽象操作契约，不能从其他具体 State 继承实现。buffer 分配属于初始化；请求使用哪些物理 token 槽位由管理器的 allocator 决定。 - -| 接口 | 完整语义 | -| --- | --- | -| 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | -| `allocate_storage(token_capacity)` | 仅接收 token 容量，原地创建全部 buffer，保留 State 对象身份和 config；页大小、设备及配套参数使用已设置的存储配置。GQA/MLA 分配 KV，分组量化同时分配动态 scale，静态 FP8 同时绑定固定标定参数；计算视图在实际分配后获取 | -| `get_token_bytes()` | 返回当前 TP rank 单个 token 的完整存储字节数，包含动态量化参数和 packed 行 padding；不包含静态标定参数。管理器按模型组合汇总后直接乘容量或 page size | -| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | -| `get_layer_view(logical_layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | -| `write_full_att_kv(logical_layer_id, token_indices, kv)` | 将 KV 写入已分配位置，完成该方案要求的编码及关联参数更新；融合压缩路径不调用本接口，具体实现显式拒绝误调用 | -| `copy_full_att_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | - -容量、device、层映射及布局标识应从同一份稳定配置和描述中取得，避免多处维护互相矛盾的副本。配置相同且只改变容量时，容量计算和真实 buffer 创建应覆盖相同的数据组成。共享底层内存按实际分配去重记账。 - -#### 存储字段与 State 元数据 - -两类 State 不提供完整存储枚举接口，各实现显式确定参与传输的固定字段。Full attention 的布局直接根据配置计算形状、dtype 和字节区间；PD/DP 调用处显式列出实际字段：GQA 为 K/V，MLA 为 latent，分组量化增加动态 scale，packed MLA 包含对应的 KV、scale、RoPE 与 indexer 字段。静态标定参数独立保留，不加入 token 传输字段。KDA/GDN 直接从配置描述 Conv/SSM，K-pool 直接描述完整压缩器 ring。 - -`TensorField` 仅作为内部搬运数据描述，仅保留 `tensor`、`axes`，字段名称由所在字典的 key 提供；Tensor 自身提供 dtype、shape、stride、device 和底层存储视图。运行态、MTP 游标及大小页池由具体 State 直接持有，不额外导出一份完整字段列表。 - -CPU/PD 的字段选择由 State 显式定义，公共区域接口如下： - -| State 类型 | CPU 区域接口 | PD 区域接口 | -| --- | --- | --- | -| `FullAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | -| `BoundedAttentionState` | `get_cpu_transfer_regions()` | `get_pd_transfer_regions()` | - -两类 State 的 CPU/PD 区域接口统一使用内层 `@lru_cache(maxsize=None)`、外层 -`@deepcopy_result`：缓存首次构建的字段描述，每次查询返回独立副本，包含区域内的轴映射。 -首次查询后配置与全局页规格保持不变；实际存储分配完成后，管理器清除各 State 的两类区域缓存， -再清除管理器的派生布局和页大小缓存。子类覆盖区域接口时也须保留这两个装饰器。 - -所有接口返回 `tuple[TransferRegion, ...]`，可以基于 meta State 调用，不分配传输 buffer。State 决定哪些字段参与传输、字段区域的 dtype/shape、TP 切片和共享页的副本写入者；每份描述以 `byte_start` / `byte_end` 记录左闭右开的字节区间，两个位置均从该 State 的区域起点计算。两类 State 的 CPU/PD 区域接口统一无参数。full-attention 从全局启动参数读取 `cpu_cache_token_page_size` 或 `pd_kv_page_size`，按固定页容量展开形状与字段偏移；全局 head 数与本地 TP 切片均由自身 config 推导。bounded（包括压缩窗口） 每次描述一份运行态。 - -本次有效 token 数属于读写流程，不参与布局生成。CPU cache 使用固定的全局 CPU 页布局；State 的传输接口在不足整页时仍使用相同的字段形状和偏移，只搬运有效 token。缓存页的保存边界和发布条件由外层流程控制。 - -GQA、MLA 及各量化方案只导出 KV、动态 scale 和随 token 保存的 indexer 字段，不传静态标定参数或 packed padding。KDA/GDN 只导出有效 Conv 窗口和一份已接受 SSM，不传 MTP 候选区、游标或整个大小页池。K-pool 导出完整原生 key/gate ring。CPU 与 PD 可以在具体子类中分别覆盖；当前字段布局相同的实现共用自己的私有构建方法。 - -manager 将 State 内区间平移为页内区间：CPU/PD 均按对应传输字段的实际最大结束位置安排后续资源，并通过布局终点计算整页大小；manager 不解析 Conv/SSM/K-pool 的内部结构。CPU 服务复用相同 State 接口；当前 CPU 联合页包含 full/bounded 中需要保存的字段，无重叠窗口不占 CPU 页。DP 映射只包含 full/bounded attention，调用它们无参数的 PD 区域接口，只取字段名和本地轴切片，不使用区域的页形状与字节偏移；本次复制的 token/request 数量由调用参数决定。 - -资源身份使用 `state.state_name`；层映射、容量和 TP 分片事实继续由 State 的 config、capacity、page_size 及计算接口维护。CPU/DP/PD 由抽象接口约定并由具体 State 实现，管理器直接调用对应方法，不再维护额外的传输能力集合。信息保留在其所属对象，传输布局仅描述当前路径所需字段。 - -字段描述只读，引用的 Tensor 可以由合法计算操作更新。字段列表不表示数据已就绪，stream/event 同步和前缀关联仍由调用流程负责。每次增加辅助 buffer 或 packed 字段，都需要检查固定传输字段、容量计算、写入和复制/传输是否一起覆盖它。公共层不根据固定属性名推测资源组成。 - -#### 计算视图与完整存储范围 - -| 布局 | 计算视图的主要内容 | 完整复制需要覆盖的内容 | -| --- | --- | --- | -| 普通 K/V | K、V | 对应位置的全部 K/V | -| 分组 INT4/INT8 | K/V 与动态 scale | packed/量化 K/V 加上对应动态 scale | -| 静态 FP8 | 以正确数值类型解释的 KV 及静态 scale | 本地 token 复制只搬 KV；跨资源使用须验证或传递固定量化参数 | -| MLA | 压缩 KV 或其所需分段 | 本 State 负责的完整压缩 KV 表示 | -| 带 indexer 的 packed 布局 | attention KV 子视图、indexer key/scale 子视图 | KV 及全部相关 packed 辅助数据；按实际存储区域处理别名 | - -`get_layer_view()` 可以返回 `DenseKVView`、`QuantizedKVView`、`MLAView` 等具体小对象。视图类型由模型/backend 在初始化绑定时确认。各布局不必返回相同 Tensor 个数，也不必填充包含所有可选字段的大对象。 - -带 indexer 的 packed 布局中，attention 与 indexer 可以分别读取 buffer 的不同区域。因此不能用 attention 的读取视图代替完整复制范围。类似地，静态 FP8 的物理存储 dtype 和计算时解释的 dtype 可以不同，视图必须明确相应参数。 - -`write_full_att_kv` 保留为 `FullAttentionState` 的抽象接口，各具体 State 实现该方法。 -压缩历史等复杂布局通过自己的计算视图或定制接口完成融合写入，不调用普通 KV 写入接口； -其实现显式抛出 `NotImplementedError`，避免误调用时静默丢失写入。 -indexer 池化等专用 kernel 可以通过专用视图写入另一区域。 -管理器与执行流程在相关写入都满足同步条件后才允许复用或发布。 - -#### 基类与子类的分工 - -| 部分 | 基类约定 | 具体 State 负责的内容 | -| --- | --- | --- | -| 初始化 | `allocate_storage`、`get_token_bytes` | 强类型构造配置、层映射和实际 Tensor 分配 | -| 层访问 | `contains_logical_layer`、`get_storage_layer_index`、`get_layer_view` | 自身布局的层映射和计算视图 | -| 写入 | `write_full_att_kv` | 实现普通写入；融合压缩路径使用定制方法或视图，普通入口显式拒绝误调用 | -| 完整复制 | `copy_full_att_tokens` 的索引与完整性契约 | 直接实现布局正确的复制，或调用兼容的公共 kernel | -| 传输 | CPU/DP/PD 直接读写契约 | 直接使用自身 storage 和具体 kernel，不从其他具体 State 继承行为 | - -具体 State 可以复用满足其布局要求的 kernel；基类不规定构造参数或内部受保护方法。 -通用复制工具只有在理解 token 维、stride、数据区域和别名关系时才适用。复杂 packed 布局由具体 State 保证复制粒度与字段完整性。 - -#### 索引、异步执行与生命周期 - -构造函数显式接收 `page_size` 和 `device`，不接收外部 KV buffer 或 token 容量，不再使用 `FullAttentionStorageSpec` 或其量化子类。层、head、dtype、TP 与量化方案继续放在 config；实际容量只保存为 `capacity`，页大小只保存为 `page_size`。构造阶段 `capacity=0`、各数据 buffer 为 `None`，不调用 `allocate_storage()`，也不分配 HOLD 页。单 token 字节数直接由配置计算；传输布局查询直接按配置计算字段形状和 dtype，不创建 meta Tensor，也不依赖实际存储。所有模型准备完成后，管理器统一规划容量并调用 `allocate_storage(token_capacity)` 分配实际存储。`device` 指定分配目标，管理器从 meta 准备阶段切换到实际设备时更新它，再调用只接收容量的 `allocate_storage(token_capacity)`。HOLD 索引由管理器决定，State 在普通容量之后为其预留一页。整页容量约束由 `FullAttKVAllocator` 检查，State 分配时不重复执行通用配置校验；量化格式由配置检查；静态标定参数在 `bind_calibration()` 中检查。 - -State 在后续 `allocate_storage()` 中自行分配数据 buffer：GQA 与静态 FP8 分别分配 `k_storage`、`v_storage`，分组量化另外分别分配 `k_scales`、`v_scales`；MLA/packed MLA 保留单份 latent `storage`。动态 scale 初始为 `None`，与 KV 同阶段分配；构造函数不接受 `buffer` 或 `scale_buffer`。每个 State 独立分配 KV 和动态 scale，容量规划累计各 State 的存储开销；main/draft 共享 allocator，不共享 KV buffer。静态 FP8 构造时不接收 `scales` 或 `q_scales`；只能通过 `bind_calibration()` 绑定真实的 CPU 标定 Tensor，实际分配时再迁移到目标设备，不把标定值转换成 meta Tensor。CPU/PD 的 K/V 字段顺序和全局形状不随物理存储拆分而改变，单 token 容量仍累计两份 KV 及其配套动态 scale。静态 FP8 写入 kernel 支持独立的 K/V 输出，保留原先的 per-tensor/per-head 标定规则；backend 通过 K/V 视图共同定位物理层，不依赖两个地址的大小关系。正常推理期间保持 buffer 地址稳定。`copy_full_att_tokens()` 复用 `kv_move.copy_kv_buffer_to_kv_buffer`，逐个 token buffer 按字节直接搬运，包含动态 scale 与 packed 行 padding，不创建中间 KV Tensor，也不在设备端检查索引集合是否重叠。调用方保证源、目标集合不相交且目标索引唯一；索引迁移到存储所在 GPU 后直接提交 Triton 算子。 - -`logical_layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 - -混合模型将实际 full-attention 逻辑层号及追加的 draft 层号写入 State 配置；例如逻辑层 3、7、8 可以分别对应物理行 0、1、2。计算侧直接以逻辑层号调用 `write_full_att_kv` 和 `get_layer_view`，由 State 完成映射，CPU/PD 描述复用同一份层标识。KV/Q 标定按物理存储行组织，backend 必须使用对应行的参数。 - -写入与复制接受与资源同设备的一维整数索引 Tensor。源和目标索引长度相同，目标唯一；基础复制先约定源目标不重叠，重叠移动另行定义。索引还须满足具体布局的对齐和粒度条件。 - -数据操作在调用方当前 CUDA stream 上提交，方法返回不代表设备操作完成；跨 stream 消费由调用方建立依赖。销毁存储前，管理器等待相关计算和传输结束并解除推理视图及跨进程存储引用。热路径不支持隐式 resize；需要重建存储时必须重新绑定视图及相关执行结构。 - -#### CPU / DP / PD 直接操作 - -三种路径直接作为 `FullAttentionState` 的方法暴露，由子类覆盖实现。CPU/PD 接口通过 `StateTransferLayout` 描述联合页与 TP 切片；DP State 接口直接复制同一 TP 分片的完整 token，不需要字段映射。布局是只读数据参数，不包含算子、buffer 所有权或执行方法。 - -| 接口 | 数据方向与位置参数 | -| --- | --- | -| `write_to_cpu(..., layout)` | 当前 State → CPU 缓存；layout 描述外部页区域、层/轴映射和写入副本选择 | -| `read_from_cpu(..., layout)` | CPU 缓存 → 当前 State；layout 描述该目标 State 从页中读取的对应范围 | -| `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` | 其他 DP 同一 TP 分片 → 当前 State，复制完整 token 的 KV、动态量化参数及 indexer 字段，不接收 mapping | -| `write_to_pd_page(..., layout)` | 当前 State → `pd_page_index` 指定的 PD 页；按 layout 写入指定区域 | -| `read_from_pd_page(..., layout)` | `pd_page_index` 指定的 PD 页 → 当前 State；按 layout 恢复对应区域 | - -`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，`AttentionStateManager`按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 - -基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttentionStateManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 - -每次调用完整处理其映射涉及的数据及关联的动态 scale、packed 字段；一次组合传输可以由若干次调用完成，上层检查所有必需数据区域均已完成。固定参数按约定传递或验证兼容性。位置映射的组成、单位及生成责任见 6.5 节。 - -方法统一在调用方当前 CUDA stream 上提交操作，返回 `None`。若内部使用其他 stream，子类须将完成依赖汇合到调用方 stream，使调用方随后记录的 event 覆盖本次操作。网络收发、任务排队和完成通知由调度层负责。 - -形状相同不能证明不同模型的 KV 语义可互换。模型/资源身份和 token 前缀匹配由管理器与缓存流程校验；State 检查布局及固定参数兼容性。跨布局重编码需要明确实现，本地 `copy_full_att_tokens` 不承担这一职责。 - -### 6.2 AttentionStateManager：内部资源的聚合操作与数据传输 - -`AttentionStateManager` 对内部资源组合提供统一的数据操作。一个模型可以包含不同角色或不同具体方案的 FullAttentionState、KDAState、GDNState 等状态实现，同一角色和具体方案最多一个；token 复制覆盖 full-attention 参与集合，清零操作覆盖需要初始运行态的参与集合，状态页操作覆盖指定大小页域的参与集合。外部调用方无需知道具体数量、类型或成员名称。 - -数据操作的依赖方向是 `radix cache / CPU cache / 其他调用方 → AttentionStateManager → 各 State`；槽位分配与归还由调用方直接访问管理器持有的共享 allocator。状态页接口不接收 radix 节点、树或前缀匹配对象，也不读取 radix 内部状态；radix 只保存并使用管理器提供的页类别、槽位 ID 和状态长度等数据。状态页在未启用 radix 时同样可以保存、恢复或用于 CPU 缓存。 - -#### 资源初始化、State 查询与传输布局查询 - -| 接口 | 语义 | -| --- | --- | -| `AttentionStateManager(request_capacity)` | 仅接收共享请求容量，初始化空资源管理器；State 与 allocator 在后续流程中设置 | -| `register_state(model_role, state)` | 将已创建的 full/bounded State 原对象绑定为规范成员，同步 state_name 与管理器成员名；只允许在创建完成前调用 | -| `finish_state_registration()` | 全部 main/draft State 对象装配完成后，先断言未重复标记创建完成，再设置 `_state_registration_finished = True`；允许 full-attention 存储仍为 `None`，只能调用一次 | -| `load_full_att_calibration(args)` | 在 `finish_state_registration()` 后、存储分配前加载静态 FP8 标定；文件只读一次，逐 State 绑定物理层区间及 TP 参数；没有静态 FP8 State 时直接返回 | -| `initialize_storage(args, *, token_capacity, mem_fraction, device="cuda")` | 共享资源初始化入口：在完成创建标记和标定绑定后，确定容量并分配存储；PD 模式在容量规划时按需构建 meta 布局。不接收或修改模型对象，返回 `None` | -| `get_state(model_role, state_type_filter, logical_layer_id, *, state_class=None)` | 全部 State 创建完成后按模型角色、方案子串、逻辑层及可选资源类别选择唯一 State；断言 `_state_registration_finished` 并使用 `lru_cache` 缓存 | -| `get_full_attention_states(role=None)` | 返回指定角色或全部 `FullAttentionState` 的 tuple | -| `get_bounded_attention_states(role=None)` | 返回指定角色或全部 `BoundedAttentionState` 的 tuple | -| `get_full_att_token_bytes()` | 返回当前 TP rank 全部 full-attention State 的单 token 存储字节数，包含动态 scale 和行 padding | -| `get_cpu_page_bytes()` | 依赖 `get_cpu_page_layouts()`，取全部 State 的最大 `byte_end` 并按 16 字节对齐；包含字段间及 State 间的必要 padding，页内不保存格式指纹、长度或版本 | -| `get_cpu_page_layouts()` | 按全局 `cpu_cache_token_page_size` 生成 CPU 联合页中各 State 的字段偏移和 TP 映射，使用 `lru_cache` 缓存 | -| `get_pd_token_page_layouts()` | 从全局启动参数读取 `pd_kv_page_size`，按 State 字段的实际偏移和结束位置生成 `kv` 布局及最终页跨度；不依赖 `get_pd_page_bytes()`，按 `state_name` 索引并缓存 | -| `get_pd_runtime_page_layouts()` | 按实际字段结束位置独立生成 `att_state` 布局及其自身所需大小，不查询 KV 布局或执行组合容量校验；每个 State 保存一份运行态，`tokens_per_page=None` | -| `get_pd_page_bytes()` | 查询两类布局大小并集中决定实际 PD 块容量：返回 KV 页与完整运行态页占用的最大值；有 KV 且 KV 页较小时记录 warning 并扩大物理页；无资源时返回 0。包含 field 与 State 间的 padding，可基于 meta State 查询，不分配传输 buffer | -| `allocate_storage(token_capacity, page_size, *, big_page_tokens, small_page_capacity, device)` | 为已创建 State 分配实际 full-attention 存储、共享 allocator 和 bounded 大小页 | -| `create_pd_ipc_copy()` | PD 导出：保留运行态 Tensor 和容量信息，排除 allocator 及所有大小页池 | -| `create_dp_ipc_copy()` | DP shared 导出：保留运行态和共享 CPU 大小页池，仅排除 allocator | -| `register_page_pools_pin_memory()` | 在当前 GPU 上统一注册全部 bounded State 的大小页池，已注册存储可重复调用 | - -聚合查询、页大小计算、传输布局 getter 和数据操作均由基类提供默认实现。CPU 布局 getter、两个 PD 布局 getter 及 CPU/PD 页大小查询分别使用 `lru_cache` 缓存，存储初始化后同步调用 `cache_clear()`；子类若覆盖布局 getter，需保留缓存清理能力。各模型通过 `register_state` 建立规范成员，组织各自资源。`role` 支持 `None`、`"main"`、`"draft0"` 到 `"draft35"`；默认返回全部，统一按完整 `state_name` 字符串排序，CPU 布局构建采用相同规则。该顺序只用于保持布局和聚合遍历的一致性，不表示模型执行顺序。不存在对应资源时返回空 tuple。集合接口覆盖同一角色多个 attention 方案，用于初始化校验、组合传输等聚合流程，不另设 main/draft 两套函数。`get_state` 提供创建完成后的单对象缓存查询，推理层通过该接口直接取得当前层使用的资源，不另存成员名。 - -State 装配期间由管理器内部遍历 `__dict__`，公共聚合查询及 `get_state` 均会断言失败，避免缓存不完整的 main/draft 集合；模型准备阶段只创建和注册 State，不查询推理资源。装配完成后先调用 `finish_state_registration()`，再通过聚合查询接口进行容量计算和存储初始化；此时查询可以返回 meta State。创建完成标记表示对象及组合关系稳定，不表示实际 buffer 已经分配。 - -容量确定后，管理器先统一设置 State 的页大小和实际设备，再调用各 FullAttentionState 的 `allocate_storage(token_capacity)` 原地分配其 KV 和配套参数，保持 State 对象身份不变；bounded（包括压缩窗口） 的固定运行态保留，通过 `initialize_page_pools()` 只补齐 bounded 大小页。因此创建完成后建立的 tuple 缓存始终引用同一组对象，不需要清除缓存或替换成员。实际存储分配前取得的 Tensor/层视图不用于后续计算；计算 backend 在实际存储就绪后建立视图。 - -实际存储分配完成后执行 warmup。CPU 布局由第一次容量查询或读写时按需生成；PD 容量规划会通过 `get_pd_page_bytes()` 在 meta State 上提前生成 PD 布局，不分配实际传输 buffer。PD 页的 token 容量由全局 `get_env_start_args().pd_kv_page_size` 唯一确定，两个 PD 布局 getter 均不接收页 token 数，按 manager 实例缓存并直接返回 `Mapping[str, StateTransferLayout]`，没有对应资源时返回空字典。 - -PD 的依赖方向为：`get_pd_page_bytes()` 分别查询两个独立的布局 getter。full-attention 和 bounded（包括压缩窗口） 各自生成偏移及本类布局所需大小，互不查询，也不调用 `get_pd_page_bytes()`。实际块大小取两类布局占用的最大值，统一在 `get_pd_page_bytes()` 完成。State 提供包含内部字段 padding 的相对偏移，manager 按字段的实际最大结束位置确定下一个 State 的起点，不再用裸存储字节数预留 PD 区域。PD 收发流程按 `kv` / `att_state` 任务选择布局。 - -CPU getter 从全局启动参数读取 `cpu_cache_token_page_size`，同样按 manager 实例缓存。不保存独立的 `cpu_page_bytes`、`pd_page_bytes`、`cpu_transfer_layouts`、`pd_transfer_layouts` 或 `cpu_page_tokens` 成员。存储分配可能改变 State 的物理页规格，因此分配完成后同时清除布局及派生 CPU/PD 页大小缓存。初始化时的异常直接向上传播，初始化失败即终止该流程；不增加属性拦截或回滚逻辑。 - -查询检查与遍历集中在 `_get_states`，通过 `@lru_cache(maxsize=None)` 缓存结果，键包含 manager 实例、资源类别和角色。返回不可变 tuple 中的对象引用，不复制存储，也不修改资源归属，不再维护实例缓存字典。manager 按进程生命周期使用；若在进程内卸载或重建管理器，需对 State 查询、页大小和布局查询方法调用 `cache_clear()`，释放缓存持有的 manager 和 State 引用，该操作会清除对应方法所有实例的查询缓存。 - -IPC 发布按用途调用 `create_pd_ipc_copy()` 或 `create_dp_ipc_copy()`，两者均浅拷贝管理器并将三个 allocator 字段置为 `None`。PD 导出会分别浅拷贝 bounded State，再清空大小页池；即使源页池是共享内存，PD 副本也不携带。DP shared 导出保留运行态及大小页池引用，要求页池已使用共享 CPU 存储，避免序列化私有页池时隐式分配、复制。两种导出都不修改源管理器，也不转移槽位所有权。 - -`shared_att_resources_utils.py` 提供 `publish_att_resources` / `load_att_resources`，用于发布和读取 attention 资源及请求索引映射的 IPC 副本。PD/DP 的发布与读取显式传入 `kind`，分别使用 `rank_storage_pd_` 和 `rank_storage_dp_`，锁及剩余读取计数也各自独立，允许同时启用且不依赖读者到达顺序。DP 仅导入本节点其他 DP 中对应 TP rank 的副本，每个发布者的读者数为 `dp_size_in_node - 1`；资源管理器和请求映射使用按节点内 DP rank 索引的字典，每项均对应本进程的 TP 分片；节点内进程 rank 仅用于加载 IPC 资源。DP 接收端设置当前设备后读取 DP 副本，再调用 `register_page_pools_pin_memory()`；manager 遍历各 State 的大小页池，复用 pool 的 `register_pin_memory()`。注册只作用于接收进程的映射，不随对象序列化；PD 接收端无需注册页池。 - -`lru_cache` 属于方法，不进入实例的序列化数据；副本及接收方首次查询时以各自实例为键建立缓存。是否允许查询由创建完成标记判断，不依赖 allocator 是否存在。 - -#### 数据操作入口 - -| 操作 | 接口 | 聚合语义 | -| --- | --- | --- | -| full-attention token 复制 | `copy_full_att_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | -| bounded 空前缀清零 | `reset_bounded_att_state(req_indices)` | 将指定槽位的全部参与 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页 | -| bounded 状态页保存 | `save_bounded_att_runtime_to_pages(...)` | 按 state_token_lens 保存全部参与 bounded State 的同类 pool、同一组槽位 | -| bounded 状态页恢复 | `restore_bounded_att_runtime_from_pages(...)` | 按 state_token_lens 将选定快照恢复到对应请求运行态及 ring 位置 | -| bounded 验证后提交 | `commit_bounded_att_state(req_indices, accepted_steps, *, role)` | 将指定角色的接受步数分发给参与 bounded State；由 State 选择已接受运行态，不暴露游标 Tensor | -| CPU 写出 / 读回 | `write_req_to_cpu(...)` / `read_req_from_cpu(...)` | 组合 KV 和状态页，按调用参数提供的长度读写各 State 在外部页中的区域 | -| DP 复制 | `copy_from_dp(remote_att_manager, src_token_indices, dst_token_indices, src_big_page_ids, dst_big_page_ids, src_req_idx, dst_req_idx)` | 直接遍历同名 State，复制 token、大页快照及指定请求的 prefill 基准运行态，复制目标窗口运行态；不分配资源 | -| PD 写出 / 读回 | `write_to_pd_page(...)` / `read_from_pd_page(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | - -#### 必须实现的聚合接口 - -基类源码划分七个大块：基础管理、Full attention、Bounded attention、CPU Cache 数据交换、DP 迁移、PD 分离传输、私有辅助实现。Bounded attention 分块依次放置运行态重置、大小页保存与恢复、MTP 提交接口。各块内保留接口的具体说明，接口顺序保持 public 在前、private 在后。 - -管理器的分类操作使用 `full_att`、`bounded_att` 标明作用范围；CPU/DP/PD 接口继续表示跨资源传输。具体 State 的分类操作与管理器使用相同命名；State 只操作自身数据，管理器负责聚合参与对象。 - -以下仅展示接口签名，实际方法在基类中提供聚合实现。 - -```python -from collections.abc import Mapping -from typing import Any, List, Literal - -PageKind = Literal["big", "small"] - - -class AttentionStateManager: - full_att_allocator: FullAttKVAllocator | None - bounded_att_big_page_allocator: StatePageAllocator | None - bounded_att_small_page_allocator: StatePageAllocator | None - - def __init__(self, request_capacity: int): ... - - def register_state( - self, model_role: str, state: FullAttentionState | BoundedAttentionState, - ) -> None: ... - - # ======================================================================== - # Full attention：token 数据复制 - # ======================================================================== - def copy_full_att_tokens( - self, src_indices: torch.Tensor, dst_indices: torch.Tensor, - ) -> None: ... - - # ======================================================================== - # Bounded attention：运行态重置、大小页保存与恢复、MTP 提交 - # ======================================================================== - - # 空前缀初始化，只处理 bounded 运行态；req_indices 定位已分配的槽位。 - def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... - - # Bounded Attention 状态页：供 radix/CPU cache 复用，调用方直接分配和归还槽位。 - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: ... - - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: ... - - # MTP 推测解码：只提交指定角色；接受步数包含本轮已经执行的基准输入步。 - def commit_bounded_att_state( - self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, *, role: str, - ) -> None: ... - - # ======================================================================== - # ======================================================================== - - # 基类使用 lru_cache 缓存布局，存储初始化后清除缓存。 - def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... - - # CPU 联合页：KV + 全部参与 linear State 的有效页。 - # 纯 full-attention 模型没有 linear 页；混合模型不能省略它们。 - def write_req_to_cpu( - self, token_indices, cpu_page_buffers, cpu_page_indices, - *, layouts: Mapping[str, StateTransferLayout], - cpu_page_ready_flags: List[bool] | None = None, - big_page_slot_ids: List[int] | None = None, - small_page_slot_ids: List[int] | None = None, - ) -> None: ... - - # 一次读取同一前缀的有序页面；全部页面读回后从末页恢复请求运行态。 - def read_req_from_cpu( - self, cpu_page_buffers, cpu_page_indices, token_indices, - *, layouts: Mapping[str, StateTransferLayout], - big_page_slot_ids: List[int] | None = None, - req_idx: int | None = None, - state_token_len: int | None = None, - ) -> None: ... - - # DP 缓存迁移：索引已由调用方分配，空 Tensor/列表表示不搬运该类数据。 - def copy_from_dp( - self, remote_att_manager: AttentionStateManager, - src_token_indices: torch.Tensor, dst_token_indices: torch.Tensor, - src_big_page_ids: List[int], dst_big_page_ids: List[int], - src_req_idx: int, dst_req_idx: int, - ) -> None: ... - - # 基类直接构建两类 PD 布局，各自使用 lru_cache 缓存；无对应资源时返回空字典。 - def get_pd_token_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... - - def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: ... - - # PD 分离：pd_page_buffers 为页集合，pd_page_index 必填；kind 沿用 kv / att_state 两类任务。 - def write_to_pd_page( - self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], - layouts: Mapping[str, StateTransferLayout], - token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, - ) -> None: ... - - def read_from_pd_page( - self, pd_page_buffers, *, pd_page_index: int, kind: Literal["kv", "att_state"], - layouts: Mapping[str, StateTransferLayout], - token_indices=None, req_idx: int | None = None, state_token_len: int | None = None, - ) -> None: ... -``` - -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 - -`page_kind="big"` 对应 `bounded_att_big_page_allocator` 和参与 State 的 `big_pages_pool`；`"small"` 对应小页 allocator 与 `small_pages_pool`。两者是独立分配域，即使槽位数字相同，也不是同一份资源。 - -调用方根据 `page_kind` 直接选择相应 allocator；该成员为 None 时不能申请该域。分配与归还使用 allocator 自身的接口，管理器不再增加包装方法。allocator 只管理空闲索引，不修改 State 的页数据或有效长度。 - -`req_indices` 只定位运行态存储，不传递请求对象。保存/恢复时，请求状态槽位与页槽位逐项对应；目标槽位唯一且已分配。调用方维护快照与实际前缀的关联，并通过 `state_token_lens` 提供各快照的绝对长度，供 ring 等状态定位和恢复；页池不重复保存长度元数据。 - -#### bounded 空前缀运行态初始化 - -`reset_bounded_att_state(req_indices: List[int])` 用于从前缀长度 0 开始计算时初始化 bounded 运行态。`req_indices` 仅接受 Python 整数列表，不支持 `torch.Tensor`；实现直接使用列表，不做类型转换。调用方已通过 ReqManager 取得请求槽位，并确认这些槽位此前的计算和传输使用结束；管理器只按传入索引初始化 bounded 数据，不查询缓存命中情况或请求对象。 - -- 对全部参与的 KDA/GDN 等状态调用 `reset_bounded_att_state(req_indices)`，将 Conv/SSM 恢复到零初值，并重置 MTP 游标及影响后续读取的候选信息。未来算法按自身契约恢复空前缀初态。 -- full-attention 没有需要清零的空前缀递推状态；该操作不遍历或清空历史 KV，不更改请求映射，也不失效大小页快照或归还索引。纯 full-attention 管理器合法地不做数据操作。 -- 目标请求槽位须唯一且在有效范围内，main/draft 所需运行态全部初始化后才可使用。重复执行在没有并发访问时具有相同结果；方法返回仍遵守 stream 完成约定。 - -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 - -#### 验证后运行态提交 - -`commit_bounded_att_state(req_indices, accepted_steps, *, role)` 只执行 bounded 运行态选择，不承担 token 验证、采样、KV 回收或请求长度更新。调用方在验证完成后传入逐请求的槽位与接受步数，两个参数必须是当前 State 设备上的一维 CUDA Tensor；请求索引为 int32/int64，接受步数为 int32，长度一致，接口内部不转换设备或类型。语义与 State 的 `commit_bounded_att_state` 接口一致；接受步数包含已经执行的基准输入步。 - -main/draft 共享管理器，因此必须显式指定 `role="main"` 或具体 draft 编号（例如 `role="draft1"`）。管理器只调用该角色参与 bounded State 的 `commit_bounded_att_state`，不将主模型验证结果直接提交到 draft 的候选状态；该角色没有需要提交的 State 时为空操作。 - -`mtp_state_index` 及其他提交控制信息只由具体 State 持有。BaseModel 不维护 `mtp_state_cursors` 或其他游标引用集合，验证流程也不直接读写游标。State 自行决定通过游标选择、数据搬移或其他方式提交；计算 backend 可通过 State 或层视图读取所需信息。 - -#### full-attention 的聚合 token 复制 - -`copy_full_att_tokens(src_indices, dst_indices)` 的两个索引均为存储设备上的一维 CUDA 整数 Tensor；直接使用输入 dtype 和 stride，不转换设备、类型或连续布局。源和目标位置不重叠。该接口在本管理器持有的 full-attention 存储内部执行位置复制。参与集合在初始化时按共享 token 索引域确定，可以包含 main/draft 的不同布局，也可以包含同一模型的多个 FullAttentionState。管理器将同一组物理索引交给各参与 State 的 `copy_full_att_tokens`；每个 State 在自己的存储中从源位置复制到目标位置，不在不同模型之间互拷数据。 - -- 复制覆盖全部参与存储层的 KV、动态 scale 和 packed indexer key/scale 等随 token 变化的数据，保持各自存储表示；共享物理区域的别名按所有权去重。 -- 源和目标为逐项对应的一维整数索引 Tensor，长度相等、目标唯一，且满足各布局的复制粒度与对齐要求；基础契约要求源目标不重叠。源数据须已就绪并保持有效，目标物理位置须已分配且允许写入。 -- 调用前校验全部参与 State 的参数和能力；无 full-attention 参与者时明确报错。方法在当前 stream 提交并返回 None，全部参与复制完成前不能消费目标数据；中途失败时目标视为未就绪。 - -例如 main 和 draft 各有一个 FullAttentionState，复制物理位置 `[10, 11]` 到 `[20, 21]` 时，一次管理器调用会分别搬运 main 和 draft 各自的对应数据。共享的是位置变换，不要求两者具有相同存储布局。 - -#### 多个 State 的参与规则 - -`AttentionStateManager`在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员按第 3.5 节命名；同一角色、同一具体方案只有一个 State，不增加公共 State 注册层。 - -例如，一个管理器同时管理两个 KDAState 和一个 GDNState，大页槽位 17 对应： - -```text -bounded_att_big_page_allocator 分配一次 slot_id = 17 - ├── KDAState A.big_pages_pool[17] - ├── KDAState B.big_pages_pool[17] - └── GDNState C.big_pages_pool[17] -``` - -调用 `save_bounded_att_runtime_to_pages(..., state_token_lens=[256], page_kind="big", page_slot_ids=[17])` 保存三个 State；`restore_bounded_att_runtime_from_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 - -调用方通过已发布的槽位引用选择可复用页面；发布引用前须保证全部参与 State 已保存对应前缀的状态；回收前解除相关引用，重新分配后先写入再发布。`restore_bounded_att_runtime_from_pages` 和 CPU 写出直接使用指定槽位的数据，不另查页池长度记录。应参与的 State 缺少 pool 属于初始化配置错误，不能跳过该 State。 - -full-attention 历史 KV 和没有独立 pool 的 WindowCompressorState 不属于这组状态页操作。需要同时恢复这些数据的调用方，使用相应资源操作，并保证历史 KV、bounded 快照和窗口历史 对应同一前缀且均已就绪。 - -#### 保存与直接归还槽位的完整性 - -- 调用方只从共享 allocator 分配一次槽位，全部参与 pool 覆盖同一索引范围。分配成功仅表示槽位可用，不表示其中已有有效状态。后续数据操作失败时，本次数据不得发布；调用方在已提交操作结束后决定重写或归还槽位。 -- `save_bounded_att_runtime_to_pages` 将同一组槽位传给全部参与 State 保存。只有所有参与者都完成对应位置的数据保存，调用方才能发布页面；一项失败时直接透传异常，不发布本次结果。 -- `restore_bounded_att_runtime_from_pages` 使用调用方选定的快照恢复各参与 State。复制阶段失败时直接透传异常，源页保持不变，调用方不得消费半恢复的数据。 -- 归还槽位前，调用方必须解除相关使用引用并满足完成条件，再直接向对应 allocator 归还一次。只归还指定大小页域的槽位，不操作相同数字的 KV 索引或请求槽位；槽位中的旧内容不能作为重新分配后的有效快照使用。 - -保存时全部参与 State 必须实际具备所要求位置的状态；不能从较新的状态补造较早快照。main/draft 可以分别计算，但聚合操作的完整性条件不变；若具体实现分阶段捕获参与者，未全部完成前必须保持聚合页不可用。 - -GPU 数据操作遵守当前 stream 的依赖约定。调用方通过 stream/event 在读取、发布或释放前确认异步复制完成。数据操作失败后，等待已提交操作结束再复用或释放相关存储。 - -共享域跨多个管理器实例时，调用方使用同一个 allocator 实例分配和归还。只有所有使用方均已结束且全部相关引用已解除后，才归还该槽位一次；不能仅因某个管理器已经结束使用就提前回收。allocator 的创建和销毁责任仍由装配时约定的唯一所有者承担。 - -#### CPU / DP / PD 聚合接口的参数与数据路径 - -`layouts` 是按完整 `state_name`（包含 main/draft 角色）索引的 `StateTransferLayout` 映射，用于 CPU/PD 外部页。DP 不构造字段映射：manager 直接遍历本地 State，通过同名成员取得远端 State，将对应 token、页或请求索引交给具体实现。DP 不再维护映射缓存，也不通过 helper 转发。 - -CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。CPU 读写由管理器按 full-attention、bounded 状态页及末页运行态恢复分阶段组织，高频读写路径不记录流程日志；buffer 与字段搬运约束交给 State，管理器检查恢复所需的目标请求、非空前缀及完整页槽位数量。异步搬运的完成仍由调用方的 stream/event 保证。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 +- 同一管理器可以同时持有 main、draft 及其他模型角色的 State。 +- full attention、bounded attention 和不同量化/编码方案通过独立 State 实现组合。 +- State 不拥有请求到物理 token 的映射，也不负责分配或回收外部页索引。 +- 共享 allocator 只管理索引和容量；State 只管理自己的 Tensor、布局和算法状态。 +- 新增算法优先新增独立 State 和创建入口，避免通过大量算法分支修改公共基类。 +- 所有跨模块数据操作都通过显式索引、长度、布局和资源引用完成，不从全局推理上下文隐式读取状态。 -**CPU 联合页** +## 2. 架构和职责 -包含 linear attention 的模型，每个已发布 CPU 缓存页都必须包含全部参与 linear State 对应页终点的有效快照，且长度一致。调用方在写出前通过前缀与槽位的关联选取完整快照，读回时按布局写入目标 pool；不完整状态不能发布。纯 full-attention 模型仍使用自身的 KV 页格式。接口不增加缺失 linear 快照的 CPU 页语义，也无需为选择有效快照终点新增元数据查询方法。 - -- `token_indices: torch.Tensor` 是由调用方准备在 CUDA 上的 full-attention 物理 token 索引，`cpu_page_indices: List[int]` 是 CPU 页池的页 ID，仅接受 Python 列表，不接受 Tensor 或其他序列，也不在接口内部转换。索引按列表顺序排列，从各页槽位 0 开始；长度与页数遵守 6.5 节的整页/尾页约定。页列表按逻辑前缀顺序排列，不要求物理页 ID 连续或递增。只传状态的 CPU 格式使用空整数 Tensor 作为 token 索引。 -- 格式包含 linear 页时，管理器写出接口分别接收 `big_page_slot_ids` 与 `small_page_slot_ids`；`cpu_page_indices` 前段对应大页、后段对应小页，每段与槽位列表一一对应。未参与的页池传 None 或空列表。full-attention KV 按完整页面列表写出一次，各 bounded State 再分别从大小页池写出快照；读回接口固定写入大页池，不接收 `page_kind`。管理器的 `big_page_slot_ids: List[int]` 只包含完整 CPU 页对应的普通大页槽位，不包含末尾碎页的临时槽位；没有 linear 页时为 None。槽位参数仅接受 Python 列表，不接受 Tensor。快照搬运不需要逐页长度数组。 -- `cpu_page_ready_flags: List[bool] | None` 是本批 CPU 页的跳过写入标记，顺序与 `cpu_page_indices` 一致；True 保留整张联合页，None 表示全部写出。缓存层保留完整 token/page 序列并一次调用资源管理器，不逐页发起传输，也不通过过滤页 ID 改变 token 映射。资源层不查询或发布 ready 状态，全部 ready 时直接返回。 -- FullAttentionState 将 ready 标记传给各自 CPU/PD 共用的传输入口；CPU 写出按 State 构造一次 GPU bool Tensor，全部字段复用。在 kernel 内按逻辑页屏蔽源 KV/量化字段读取和目标写入；若当前 block 没有待写元素，直接跳过后续索引读取、地址计算和数据搬运。跨页 block 混有 ready 与未就绪元素时保留逐元素 mask。PD 与读回不传标记、不分配此 Tensor，也不执行 ready 归约判断。bounded State 接收对应大小页段的 ready 列表，在取快照和构造复制视图前跳过 ready 页。 -- CPU 写出只读取大小页快照，不接收请求索引；无重叠的窗口 State 返回空区域，不占联合页。 -- `read_req_from_cpu` 一次处理同一前缀的有序页面，将全部 KV 写到给定物理位置，将各页 linear 快照写入对应 pool 槽位。管理器根据 `token_indices.numel()` 和布局的页 token 数判断是否存在末尾碎页，将其完整 linear 快照写入预留 LOAD 大页；完整页写入 `big_page_slot_ids` 指定的大页槽位。读回后直接使用最后一页恢复 `req_idx` 的 linear 运行态，不修改调用方的槽位列表。仅含碎页时传 `big_page_slot_ids=[]`；只传状态的格式使用空 token 索引，按传入的普通槽位恢复。 -- 恢复终点须满足各 State 的快照边界；窗口压缩器要求完成一个 compress_stride，最后一页可以是大页的一部分。 -- 任一资源读入失败时异常直接透传，已写入的数据不得作为就绪数据使用。写出失败时不得发布不完整联合页。 - -例如恢复三个 CPU 页，终点长度依次为 `[P, 2P, 3P]`：先读回各页 KV 和状态快照，再从最后一个大页槽位恢复请求在 `3P` 的运行态。无重叠窗口在该边界无需历史快照，后续计算使用新输入；调用方不再额外恢复一次。 - -**DP 直接复制** - -DP fetch 的来源是 radix cache 命中的前缀。缓存页大小须对齐压缩窗口的产出步长；GLM 启动时检查 `linear_att_hash_page_size % index_kpool == 0`。manager 接收远端管理器、CUDA token 索引、源/目标大页 ID 及源/目标请求索引，按同名 State 依次复制 token、大页快照与请求运行态。窗口 State 和其他 bounded State 使用同一流程，不再另设窗口重置阶段。 - -KDA/GDN 的源游标须为 0；窗口 State 的两端配置和 MTP 容量须相同。复制操作不创建历史快照，调用方负责前缀对应关系、合法缓存边界、源数据就绪及引用有效期。不需要复制 token 或大页时传入空 Tensor 或空列表；指定请求的运行态仍会复制。 - -**PD 页读写** - -| `kind` | 必需参数 | 处理对象 | -| --- | --- | --- | -| `"kv"` | `token_indices`；`req_idx=None` | 按单页容量读写全部参与 full-attention 的 KV、scale 和 packed 字段 | -| `"att_state"` | `req_idx`；`state_token_len`；`token_indices=None` | 读写 prefill 基准 bounded 状态（包括固定压缩窗口），不经过大小页 pool | - -PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 是有效槽位，不能当成缺失参数。`pd_page_buffers` 第一维是页维度，调用方必须传入 `pd_page_index`；资源层先选 `pd_page_buffers[pd_page_index]`，再应用 layout 中相对该页的区域偏移，各偏移只应用一次。连接参数兼容性由 PD master 校验，资源层不再逐页读写布局指纹、长度或版本字段。传输任务通过 `end_kv_index` 管理请求进度；`att_state` 任务的该值作为 `state_token_len` 传给 manager，再转交各 bounded State，用于各自的窗口定位及恢复。full-attention 的 token 传输接口保持不变；页内只保存状态数据。无运行态资源的管理器不接受 `att_state` 任务。 - -CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 - -NCCL Xfer 用 mode 区分 send/recv 句柄，二者共用 finish_cuda_event。接收线程在 recv 后创建并记录 event，record 返回后才发布；在此之前 recv 句柄的 check_status() 保持 PROC，避免未记录的 CUDA event 被误判为完成。Decode 与 NIXL 统一轮询 check_task_status()，直到状态为 DONE 后才读取 PD 页。接收提交失败时返回 ERR，由读页线程转入已有失败队列，不读取无效页面。 - -#### 调用方与资源实现的边界 - -radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存和恢复状态页时调用上述通用接口，前缀与槽位的对应关系由 radix 调用流程保证；槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 - -CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，`AttentionStateManager`负责组合布局、字段覆盖和分片对应。请求阶段、MTP 验证决策及 ReqManager 映射维护由外部流程组织；验证得到接受步数后，通过管理器提交，由各 State 更新自己的运行态选择信息。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 - -### 6.3 模型推理使用具体资源视图 - -Model 和 InferStateInfo 通过 `att_state_manager` 访问资源，不额外保存 full-attention State 别名。执行 KV 写入或读取层视图时，直接访问管理器上的具体成员，例如 `main_gqa_bf16_state`、`draft0_mla_packed_bf16_state`。InferStateInfo 的 `att_state_role` 指定当前模型的资源前缀；`is_draft_model` 仍用于推理流程判断。 - -视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与池化窗口的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 - -普通 KV 层可支持多种存储精度。推理层使用 `get_state(model_role=infer_state.att_state_role, state_type_filter="gqa", logical_layer_id=self.layer_num_, state_class=FullAttentionState)` 取得资源。MLA 推理层和 GLM indexer 使用 `state_type_filter="mla"`，indexer 使用自己的 `layer_idx`。通用写入模板使用空字符串和 `FullAttentionState` 类型，不推断算法。匹配结果由管理器缓存，推理层和 indexer 不保存重复成员名，也不需要模型启动时逐层绑定。不同精度的 State 由逻辑层归属区分。 - -Gemma4 共享 KV 层按自己的逻辑层号查询 State 并取得视图，由 State 将共享消费者映射到源物理层。FP8 FA3 和 FlashInfer 分别直接访问对应角色的 per-head 和 per-tensor State 成员获取标定参数；GLM indexer 按窗口和步长访问对应角色的 WindowCompressorState。物理层索引在具体 State 内确定。同一角色的不同量化方案可以分别覆盖不同层；复杂模型仍需选择支持各层方案的计算 backend。 - -成员访问只读取 Python 属性，不复制 buffer，也不执行 GPU→CPU 同步。KV 写入直接调用该 State 的方法，packed KV/indexer 通过其层视图访问。DFlash 的特殊 KV 写入上下文同样设置管理器和 draft 标记。需要主机标量的 backend 在 State 初始化时准备标定镜像，避免 forward 中逐层执行 GPU→CPU 同步。 - -### 6.4 调用边界与独立传输进程 - -单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttentionStateManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 +```mermaid +flowchart TD + Model[模型装配] --> Manager[AttentionStateManager] + Manager --> Full[FullAttentionState] + Manager --> Bounded[BoundedAttentionState] + Manager --> Alloc[共享 allocator] + Request[ReqManager] --> Mapping[请求槽位与 token 映射] + Cache[Radix / CPU cache] --> Manager + Transport[DP / PD 调度] --> Manager + Forward[模型推理与 attention backend] --> Manager + Manager -.不拥有.-> Mapping +``` -具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttentionStateManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该`AttentionStateManager`的具体实现提供。 +### 2.1 AttentionStateManager -跨进程共享通过管理器的 `create_pd_ipc_copy()` / `create_dp_ipc_copy()` 保留 State Tensor 引用并由 IPC 传递句柄，不调用 `allocate_storage()` 重新分配 KV；两者均排除 allocator，PD 副本排除全部大小页池，DP 副本保留共享 CPU 大小页池，不导出整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 +管理器负责： -### 6.5 复杂 attention 组合与 TP 的位置映射 +- 注册、查询和组织模型 State。 +- 按模型需要创建或绑定 full-attention、bounded-attention allocator。 +- 在容量确定后统一初始化 State 存储和大小页池。 +- 聚合多个 State 的重置、复制、保存、恢复及 CPU/DP/PD 操作。 +- 生成并缓存各 State 的传输布局。 +- 创建 PD/DP IPC 视图并重新绑定需要本地设备注册的页池。 -#### 位置映射的范围 +管理器不负责： -传输需要分别描述逻辑层到物理层的映射、组合页中的资源区域、各字段的 TP 分片与副本关系，以及本次操作的 token/page 索引。KV、scale、Conv、SSM 和窗口 可以使用不同的 dtype、stride、对齐和轴划分，不能统一简化为一个 `offset` 或一个全局 `tp_rank`。 +- 请求 KV 映射表、radix 树、缓存命中策略和页面引用计数。 +- 选择请求要使用的 token 或状态页索引。 +- 解释具体算法的 Tensor 字段、ring 位置或压缩规则。 +- 在 State 未提供能力时通过通用分支猜测其行为。 -#### 信息由谁维护 +### 2.2 State -| 信息 | 维护位置 | 示例 | -| --- | --- | --- | -| 本地存储布局与分片事实 | 具体 State 的配置、存储参数及内部 Tensor | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | -| 传输字段及本地 TP 切片 | 具体 State 的 `get_cpu_transfer_regions()` / `get_pd_transfer_regions()` | 字段选择、全局形状、State 内字段的起止区间和各 rank 的读写切片 | -| 组合页布局及资源间分区 | 模型专属 `AttentionStateManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main_gqa_bf16_state、draft0_mla_fp8_state、Conv、SSM、窗口各区域的起止位置、对齐和整页容量 | -| 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `StateTransferLayout`；按资源与 CPU/PD 路径分别保存 | -| 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | +每个 State 代表一种逻辑 attention 算法和存储编码方案。State 必须独立实现自己的配置、存储、计算视图、量化/编码、快照和恢复逻辑。 -资源的全局身份须包括模型角色及资源标识，不能只用层号识别 main/draft。同一模型角色可以有多个存储或量化方案不同的 State，但每个 `(role, state_type)` 最多一个；同方案的所有层归入该 State。管理器按需组织成员，不增加固定的六个成员变量。共享 allocator 只保证索引域，不能推导层、head、字段或传输字节位置。 +同一种逻辑算法的不同存储编码属于同一算法目录。例如： -同一 State 可能服务 CPU cache 和多个 PD 对端，不应把组合页的绝对位置固化成唯一的 `self.pd_offset`。本地存储事实固定在 State，外部位置由该次传输选择的 layout/mapping 提供。算子不得用进程全局的 TP rank 替代被操作 State 的分片描述，尤其是在传输进程导入多个 rank 的 State 时。 +- `full_att/gqa/` 包含普通 GQA、INT4/INT8 分组量化和静态 FP8。 +- `full_att/mla/` 包含普通 MLA。 +- `full_att/packed_mla/` 包含 Packed MLA 及其 FP8/E8M0 变体。 +- `bounded_att/gdn/`、`bounded_att/kda/`、`bounded_att/sliding_window/` 和 `bounded_att/window_compressor/` 分别包含对应算法及其传输实现。 -#### StateTransferLayout：单个 State 的传输布局 +量化、压缩或 FP8 变体不单独建立只按精度命名的顶层目录；它们应放在所属逻辑算法目录中。 -`StateTransferLayout` 通过 `state_name` 绑定一个具体 State，描述它在 CPU/PD 外部页中的传输布局。它通过 `regions` 组织字段级映射，每个 `TransferRegion` 描述 State 内一个字段与外部页区域之间的对应关系，`AxisMapping` 进一步描述轴上的拷贝范围： +### 2.3 allocator、页池和请求映射 ```text -StateTransferLayout # 一个 State 的传输布局 -├── state_name # 对应的 State -├── tokens_per_page # 外部页的 token 容量 -├── byte_start / byte_end # 本 State 的页内预留区间 [start, end) -└── regions: tuple[TransferRegion, ...] - └── TransferRegion # 一个字段与外部区域的映射 - ├── field / byte_start / byte_end / dtype / axes / shape - ├── axis_mappings: tuple[AxisMapping, ...] - └── write_enabled -``` - -同一类布局页可包含多个 State 的 layout，共用该类布局的 token 容量。每个 layout 的 `[byte_start, byte_end)` 只描述本 State 的预留区间，包含该 State 的字段、对齐及预留空间；其中每个 region 的区间只包含对应字段的实际数据。后一个 State 从前一个 State 的终点开始排列，所有 layout 的最大终点表示该类布局的总占用。 - -CPU buffer 使用 `(page_num, page_bytes)` 的二维页池，通过 `cpu_page_buffers[page_id]` 选择物理页；页间距由实际 buffer 决定，不保存在 State 布局中。CPU 整页大小按全部 State 的占用向上对齐到 16 字节。PD 的 `kv` 与 `att_state` 分别生成自身区间，`get_pd_page_bytes()` 统一选择实际传输块大小，可能大于运行态布局的占用。同一字段也可通过多个 region 描述不同范围，例如 Conv 的 Q/K/V channel 切片。读回时反向使用字段与外部区域之间的映射。 - -这些类型均为只读数据参数，不持有 buffer 或执行搬运。具体 State 的方法根据 layout 完成数据操作。 - -| 字段 | 语义与单位 | -| --- | --- | -| `state_name` | 与管理器成员名和 State 的资源名称一致，用于选择当前资源的页区域 | -| `tokens_per_page` | 含 KV 的页格式为正整数 token 容量；纯递推状态页为 `None`，按请求状态或 checkpoint 槽位处理 | -| `byte_start` | 本 State 相对整页的起始字节，包含该位置；region 已使用页内坐标，搬运时不再重复加上此起点 | -| `byte_end` | 本 State 的结束字节，不包含该位置，涵盖自身字段、对齐及预留空间；与起点之差是本 State 占用，不是物理页间距 | -| `regions` | 本 State 参与传输的物理字段或片段列表，同一字段可分为多个映射项 | -| 每项的 `field`、`dtype`、`axes`、`shape` | 标识 State 中的字段及外部页区域；axes 命名各维度，shape 用元素数量；外部区域采用连续布局，stride 由 shape 和 dtype 推导 | -| 每项的 `byte_start`、`byte_end` | 字段的左闭右开字节区间，组合后相对整页；长度等于 `prod(shape) * dtype.itemsize`，不包含字段间的对齐空隙，必须位于所属 State 区间内 | -| 每项的 `axis_mappings` | `tuple[AxisMapping, ...]`，通过 `axis`、`local_start`、`external_start`、`length` 描述轴名、本地起点、外部起点和元素数量；一个 region 内同一轴只出现一次，local_circular/external_circular 描述两侧环形寻址；普通非连续范围拆成多个 region | -| 每项的 `write_enabled` | 指定 TP 副本是否负责写入该区域，避免多个副本重复写入；恢复时各目标副本均可读取 | +ReqManager +├── request_slots +└── req_to_token_indexs -以下冻结数据类只描述传输位置，不包含搬运方法： +AttentionStateManager +├── full_att_allocator +├── bounded_att_big_page_allocator +└── bounded_att_small_page_allocator -```python -from dataclasses import dataclass, replace -import torch - - -@dataclass(frozen=True) -class AxisMapping: - axis: str - local_start: int - external_start: int - length: int - local_circular: bool = False - external_circular: bool = False - - -@dataclass(frozen=True) -class TransferRegion: - field: str # 对应 State 内部搬运字段，如 k/v/scale/packed - byte_start: int # 字段起点；State 返回局部位置，manager 组合后相对整页 - byte_end: int # 字段终点，不包含；区间长度等于 shape 对应的数据字节数 - dtype: torch.dtype # 外部区域的数据类型 - axes: tuple[str, ...] # 如 (token, layer, head, dim) - shape: tuple[int, ...] # 外部区域在各轴上的完整尺寸，不含页池的 page 轴 - axis_mappings: tuple[AxisMapping, ...] - write_enabled: bool = True - - -@dataclass(frozen=True) -class StateTransferLayout: - state_name: str # 含模型角色，如 main_gqa_fp16_state - tokens_per_page: int | None # KV 页为 token 容量，纯状态页为 None - byte_start: int # 本 State 的页内起点，包含 - byte_end: int # 本 State 的页内终点，不包含，含自身对齐与预留空间 - regions: tuple[TransferRegion, ...] +State +├── runtime tensors +├── optional page pools +└── transfer layouts ``` -`field` 在初始化时与 State 字段匹配；本地 Tensor 继续保留自身的 dtype、stride 和字段起点，可以是 packed buffer 的非连续视图。外部页中每个 region 的完整区域采用连续布局，`axes` 与 `shape` 长度相同且轴名唯一，逐维 stride 由后续维度大小乘 dtype 字节数推导，不再作为输入字段。多个 region 可以通过切片引用同一外部区域，例如 Conv 的 Q/K/V 分片。部署配置确定布局，CPU cache 读写直接使用字段区域与 TP 切片。 - -`AxisMapping` 的起点和长度均以元素为单位。在 `TransferRegion` 中，映射固定描述本地 State 与外部页的对应关系；写出和读回共用两侧起点，仅反转数据流向。构造时使用字段名，搬运时通过属性访问，不再依赖四元组的位置。 +- `ReqManager` 是请求槽位和 token 映射的唯一所有者。 +- `FullAttKVAllocator` 管理 full-attention token 索引。 +- `StatePageAllocator` 管理 bounded 状态页槽位。 +- 各页池只保存快照数据，不维护独立空闲列表和前缀引用关系。 +- 缓存调用方保存“前缀长度—页槽位”的关联，并在最后一个使用方结束后归还索引。 +- main/draft 可以共享 allocator；共享只表示索引域相同，不表示两个 State 的数据自动一致。 -层维度统一完整、同序复制，搬运前检查源目标层数一致。逻辑层到物理存储位置的映射由 State 自身维护，例如逻辑层 `[3, 7]` 紧凑存为 `[0, 1]`，外部页沿用相同层顺序。CPU/PD 区域和 DP State 复制均不提供层子集选择或层重排参数。 +## 3. State 命名、注册和查询 -token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_mappings`。其余搬运轴通过 `AxisMapping` 指定范围，同一组切片应用于全部层；不同 TP rank 的 head/channel 差异由这些切片表达，不改变层顺序。普通非连续范围使用多个连续片段表达；环形轴使用 `local_circular` / `external_circular`，按两侧 Tensor 的实际轴容量对 `start + i` 取模。标记始终按本地 State → 外部页解释，读回时反向执行；搬运长度不超过各环形轴容量。 +### 3.1 State 类型名称 -环形标记只描述寻址约定，必须由具体 State 的专用 kernel 实现；普通 full/bounded 的 TP 切片路径不支持将这两个标记设为 True。当前窗口压缩器通过 `state_token_len` 和两侧窗口容量直接完成环形搬运，不读取 `AxisMapping` 的环形标记。 - -`shape` 表达该外部字段区域的完整范围，而 `axis_mappings` 表达当前 State 实际访问的子范围；两者不能混淆。State 区域接口返回的起点、终点均相对本 State 起点；管理器组合时同时平移两个位置： +`get_state_type()` 返回稳定的小写 snake_case 算法和存储方案标记，不包含 `full_att`、`bounded_att`、模型角色或 `_state` 后缀。管理器生成： ```python -replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) +state_name = f"{model_role}_{state.get_state_type()}_state" ``` -最终 `StateTransferLayout.regions` 中的区间相对整张外部页，搬运时直接使用。共用搬运函数要求 buffer 至少为二维，并沿第一维选择对应页；PD State 入口直接接收页集合与显式页索引。取页后直接按字段字节区间、dtype 和 shape 构造视图，不重复断言布局边界。 +名称必须能够区分影响布局或数据解释的方案参数，例如量化位宽、分组方式、scale 方案和编码精度。层号、容量、设备和请求信息不进入名称。 -下面将前述结构用于 P 端 TP=2、rank 1：本地 K/V 各有 2 个 heads，外部页使用 4 个全局 heads。假设组合页共 64 KiB，其中前 32 KiB 存 main K/V，后 32 KiB 用于其他资源： +### 3.2 注册和查询 ```python -k_region = TransferRegion( - field="k", - byte_start=0, - byte_end=16384, - dtype=torch.float16, - axes=("token", "layer", "head", "dim"), - shape=(8, 2, 4, 128), - axis_mappings=( - AxisMapping(axis="head", local_start=0, external_start=2, length=2), - AxisMapping(axis="dim", local_start=0, external_start=0, length=128), - ), -) -v_region = replace(k_region, field="v", byte_start=16384, byte_end=32768) - -layout = StateTransferLayout( - state_name="main_gqa_fp16_state", - tokens_per_page=8, - byte_start=0, - byte_end=32768, - regions=(k_region, v_region), +manager.register_state(model_role, state) +manager.finish_state_registration() +manager.get_state( + model_role=model_role, + state_type_filter=state_type_filter, + logical_layer_id=logical_layer_id, + state_class=state_class, ) ``` -示例中的配置标识字符串只为说明关联关系，不规定标识编码方式。K/V 每个完整区域为 `8 * 2 * 4 * 128 * 2 = 16384` 字节；当前 rank 通过 head 映射只写各区域的一半，另一半由另一个 TP rank 写入。逻辑层 `[3, 7]` 到本地层 `[0, 1]` 的关系保留在 State，外部区域按相同顺序完整存放这两层。 - -例如当前 rank 的本地 head 0、token 在页内位置 0、物理层 0、dim 0，对应 K 区域地址偏移为 `0 + 2 * 256 = 512` 字节，V 为 `16384 + 512 = 16896` 字节。CPU buffer 每行 `65536` 字节，选择下一行即进入下一物理页；本 State 的 `byte_end - byte_start = 32768` 只表示自身占用。 - -D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸，把本地 head 映射改成 `AxisMapping(axis="head", local_start=0, external_start=3, length=1)`。P/D 根据各自的 TP 配置生成切片，完整 layout 对象可以不同。 - -字段列表按实际物理表示组织。例如 K、V、scale 可分别描述；packed 数据也可整体描述为一个字节区域。别名区域只搬运一次。Conv 的 Q/K 和 V、SSM 可能采用不同 head 映射，应分别描述，不能把一个资源的 `head_start` 套用于所有字段。数据类型、量化组粒度及索引单位由对应格式约束。 - -`AttentionStateManager`从 State 的 CPU/PD 区域接口取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 - -FullAttentionState 的 token 页操作要求 `tokens_per_page` 为正整数。CPU 参数 `cpu_page_buffers` 为完整二维页池；PD 参数 `pd_page_buffers` 同样为页集合，`pd_page_index` 指定物理页。传入完整页集合，不提前切到单页或 State 区域。对于常规 K/V，一个外部地址可以写成： - -```text -page_base = cpu_page_buffers[page_id].data_ptr() # CPU 选择对应行 -page_base = pd_page_buffers[pd_page_index].data_ptr() # PD 选择对应行 - -page_base -+ region.byte_start -+ token_in_page * token_stride_bytes -+ buffer_layer * layer_stride_bytes -+ buffer_head * head_stride_bytes -+ dim * dim_stride_bytes -``` - -`data_ptr()` 已指向传入 Tensor 的起点，不重复加 Tensor 的 storage offset。State 内部源/目标地址使用该 State 自身 Tensor 的基址及 stride；PyTorch 元素 stride 进入字节描述前乘以对应 dtype 大小。非连续区域必须保留真实 stride，不能通过会产生拷贝的 reshape/contiguous 假装仍指向原写入目标。 - -CPU/PD 的 token 页操作统一从外部页内槽位 0 开始，公共接口不设置页内 token 起点参数。令 `i` 为本次索引向量位置： - -```text -# CPU：按给定页 ID 顺序，从每页槽位 0 开始 -page_id = cpu_page_indices[i // tokens_per_page] -token_in_page = i % tokens_per_page - -# PD：通过显式索引选择本次传输页 -page_id = pd_page_index -token_in_page = i -``` - -有效 token 数由索引长度给出。普通 CPU 路径要求整页操作；支持尾页的路径允许最后一页只有前部有效，CPU 页 ID 数量须与覆盖页数一致，PD 单页要求 `len(token_indices) <= tokens_per_page`。checkpoint 对应的前缀长度由调用方关联到槽位，不作为页拷贝参数，也不保存在 pool 中。 - -上述固定的是 token 维度起点。组合资源的 `byte_start` / `byte_end`、TP head/channel 范围以及 State 的实际物理 token 索引仍可非零，应继续保留在各自的位置映射中。 - -#### DP 同名 State 的直接复制 - -DP 请求共享信息分为两个数组，均按请求在当前批次中的位置和节点内 DP 编号索引： - -- `shared_req_infos` 保持 `[max_req_num, dp_size_in_node, 2]`，两项分别为 KV 长度和请求槽位。 -- `shared_req_big_page_ids` 独立使用 `[max_req_num, dp_size_in_node, max_big_pages]`，其中 `max_big_pages = max_req_total_len // big_page_tokens`；未启用大页时不创建。 - -大页列表第 `i` 项对应前缀长度 `(i + 1) * big_page_tokens`。主 TP rank 将 radix 已命中的大页与请求持有的大页按对应边界写入，每次先把当前行清为 `-1`，避免较短请求复用共享行时读到旧页 ID。发布和读取沿用节点内 barrier，源请求引用保持到所有复制完成。 - -`TransTask` 使用 `src_token_indices` / `dst_token_indices`、`src_big_page_ids` / `dst_big_page_ids` 明确配对关系；`dst_req` 持有目标请求，`src_req_idx` 指定来源运行态，`src_dp_rank` 以节点内 DP 编号定位对应 TP 分片的远端 manager 和请求索引表；节点内进程 rank 只在初始化加载 IPC 资源时使用。目标 token 索引来自 CPU allocator，在调用资源接口前一次性转成 CUDA int64。 - -目标 DP 只复制从本地已有 KV 长度到来源 KV 长度之间新增的完整大页，不重复复制已有前缀页。先申请目标大页再申请 KV；目标大页分配必须成功，通过断言检查。复制完成后，将目标页按前缀长度登记到 `req.hybrid_len_to_big_page_id`，后续 radix 插入和释放沿用请求的现有归属管理。未跨越新大页边界时，大页列表为空。 - -DP 副本之间使用相同的 State 成员名。`copy_from_dp` 接收远端 manager、源/目标 token 索引、源/目标大页 ID，以及单个源/目标请求索引 `src_req_idx` / `dst_req_idx`，直接遍历本地 full/bounded State，通过 `getattr(remote_att_manager, state.state_name)` 定位远端同名 State。不需要复制 token 或大页时传空 Tensor 或空列表；指定请求的运行态仍会复制。manager 将单个请求索引包装为列表，调用各 bounded State 的 `copy_runtime_state_from_dp`。调用方保证源 MTP 游标为 0、前缀满足各 State 的缓存边界，以及源页就绪、引用有效。当前 DP fetch 从独立的大页共享数组读取新增前缀对应的页 ID，在同一次调用中完成 token、大页和运行态复制。 - -full-attention 使用 `copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state)` 指定源对象和配对索引；bounded State 分别通过 `copy_big_pages_from_dp` 与 `copy_runtime_state_from_dp` 指定源对象和配对索引，不接收字段映射。full-attention DP 的源、目标 token 索引必须是位于各自存储设备上的 CUDA Tensor；源、目标索引均支持 int32/int64。manager 和 State 直接使用传入索引，不进行设备迁移、dtype 转换或 Python list 转换；bounded 请求索引仍为 `list[int]`。布局及编码兼容检查使用断言。full-attention 所有 token 索引入口（KV 写入、本地 token 复制、CPU/PD 页读写及 DP 复制）都只接受 CUDA Tensor，由调用方准备好索引；State 不迁移设备、不转换 dtype，也不调用 `contiguous()` 创建副本。CPU cache 页面仍在主机内存，token 索引仍须位于 CUDA。PD 网络任务中的 `mem_indexes` 仍以列表序列化，由传输进程在调用 `copy_page` 前显式创建设备上的整数 Tensor。State 的四个公开页读写接口将显式存储成员交给本方案的 `transfer_*_pages`；CPU 传入页索引列表，PD 传入包含指定页索引的列表，不再保留 `_buffer_io()`。CPU/PD/DP 显式访问各实现的存储成员：GQA/静态 FP8 使用 `k_storage`、`v_storage`，分组量化同时操作 `k_scales`、`v_scales`，MLA 使用 `storage`；本地变量和页切片逐项列出，具体 State 不提供统一枚举存储 buffer 的接口。完整 token 复制同样直接对各存储成员调用搬运算子。所有 full-attention CPU/PD 共用 kernel 和本地 token 复制 kernel 均按输入索引的 stride 寻址，直接支持非连续 CUDA 索引，保留 int32/int64 类型；DP fetch 的源索引直接取自远端 CUDA 请求表，目标 allocator 返回 CPU int32 索引，因此由 DP 调用方在进入 manager 前一次性准备当前设备的 CUDA int64 索引，供全部 State 共用；State 仅断言索引为 CUDA Tensor，不重复转换。页布局由对应 State 生成，调用方保证读回目标索引唯一；搬运时仅保留页大小和页数等必要检查，不逐次重建字段集合或回读索引去重。 - -full-attention DP 的 `dp_transfer.py` 提供显式存储入口：`copy_gqa_tokens` 接收 K/V，供 GQA 和静态 FP8 使用；`copy_group_quantized_tokens` 另接收 K/V 动态 scale；`copy_packed_tokens` 接收完整 latent/packed 存储。三种入口共用四维 token 布局的 Triton kernel，一次 launch 搬运当前 State 的全部字段；编译时展开 Tensor tuple 为指针参数，不创建或上传指针表。按两端 token 索引及各自真实 stride 直接寻址，无 `index_select`、跨设备 payload `.to()` 或 `index_copy_` 中间存储。每个字段独立选择对齐的整数搬运宽度，支持两端不同 token 容量、非连续索引和存储视图；默认 grid=24、BLOCK=4096、warps=8、stages=1，通过跳步循环覆盖全部元素，地址计算使用 int64。静态 FP8 标定参数不搬运；packed MLA 按原 DP 语义复制完整行，包括 padding，区别于 CPU/PD 的有效字段页格式。调用方保证目标索引唯一、源目标区域不重叠，并负责远端 peer/IPC 映射、源数据就绪及存储生命周期；算子在目标设备当前流执行，不在内部建立 peer 映射或同步。 - -DP 拷贝在两张 H200 上对照旧 gather → `.to()` → scatter 路径测量：16 层、1025 个乱序 token，源/目标容量为 2048/3072；GQA/INT4/INT8 使用 8 个 KV head、head_dim=128。每种方案和方向预热后连续调用 20 次，三轮取中位数，计时包含 Python 入口及末尾设备同步，不含分配、初始化和编译；重复使用同一组存储，未测试与推理并发。 +注册完成后才允许查询。查询必须同时按模型角色、方案过滤条件、逻辑层归属和可选 State 类型筛选，最终只能匹配一个 State。`get_state` 可缓存；缓存只依赖注册期稳定信息，不依赖后续原地分配的 Tensor。 -| DP 布局 | 同卡加速比 | 跨卡加速比 | 跨卡旧→新有效载荷带宽（GB/s） | -| --- | ---: | ---: | ---: | -| GQA BF16 | 2.46× | 1.54× | 202.83 → 311.41 | -| INT4 + scale | 3.33× | 2.82× | 107.44 → 302.90 | -| INT8 + scale | 4.14× | 2.65× | 117.73 → 312.54 | -| MLA BF16 | 1.70× | 1.29× | 178.72 → 230.66 | -| Packed MLA FP8 | 2.97× | 1.68× | 129.99 → 217.73 | +State 必须提供逻辑层归属查询和访问视图： -CPU 直接索引拷贝的实现约束：每个传输字段用真实 Tensor stride 计算本地槽位和外部页地址，支持独立 K/V、动态 scale、packed MLA 的 RoPE/indexer 字段及 TP 切片。按同宽整数搬运原始编码，保留浮点 NaN 等位模式；未参与字段、其他 TP 分片、GPU packed padding 和最后一页的无效 token 均不写入。写出时遵守 `write_enabled`，读回不受它限制。写入的 CPU 页 ID、读回的目标 token 索引由调用方保证无冲突；一次调用只分配少量索引 Tensor，临时显存不随 KV 数据量增长。CPU cache 搬运仅支持 CUDA State，不提供 CPU 参考路径或回退。GQA、静态 FP8、分组量化、MLA、packed MLA 各自提供专用入口和 kernel 文件；GQA/静态 FP8 显式接收 K/V，分组量化再接收 K/V scale，MLA 接收 latent，packed MLA 接收 packed 存储及 KV/indexer 区间参数。每个 full-attention 方案的 CPU/PD 共用本方案的 `transfer_*_pages` 入口及一个双向 kernel，显式接收 K/V、token 索引、`page_buffers` 和 `page_indices: list[int]`。CPU 传入页索引列表，PD 入口将单个 `pd_page_index` 转为 `[pd_page_index]`，公共 State 接口保持原有语义。读写均直接寻址，不使用 `index_select`、`index_copy_` 或中间 KV Tensor；单页和多页统一构造 GPU 页索引 Tensor，由 kernel 按 `PAGE_IDS[token // PAGE_TOKENS]` 定位物理页，不再设置单页特判。各 `transfer_*_pages(..., grid_num=24, num_warps=8, num_stages=1)` 控制每个字段的发射配置，默认 `BLOCK=1024`；kernel 从 `program_id * BLOCK` 开始，以 `tl.range(..., num_stages=NUM_STAGES)` 按 `num_programs * BLOCK` 跳步循环，因此 grid 数量固定时仍可覆盖任意搬运范围。先构造全部参与字段的视图，再逐字段提交 kernel，保留 TP 切片、副本写出规则及尾页有效范围。网络收发和同步继续由调用方负责。静态 FP8 标定参数不随页传输，动态量化 scale 与 token 一起传输；packed MLA 只搬运有效字段，保留本地行 padding。所有 full-attention 页搬运和 KV 写入均要求 CUDA State；CPU 实例仅用于布局、容量等描述性操作与测试。DP 直接索引复制使用下述专用入口。 - -full-attention 的宽类型搬运只调整字段最内轴，不合并具有独立寻址语义的 layer/token/head/page 轴。两端末轴连续，且行宽、起点、storage_offset、外层 stride 都满足对齐时，依次尝试 8/4/2 字节的整数 view；不满足时保留原元素宽度。view 不复制数据、不转换数值，WIDTH、ELEMENT_BYTES、各轴 stride 和搬运单元总数均由转换后的视图得到，浮点 NaN、量化 payload 和 scale 的原始位模式保持不变。页起点、页跨度不满足宽对齐或最内维不连续时须正确回退；TP 切片、乱序 token/page 索引、ready 页跳过、尾页及 packed padding 的语义不变。五类 CPU/PD 方案继续使用各自的专用入口和 kernel，不新增通用 State 或合并字段 launch。 - -full-attention 参数在 H200 上实测：先对 GQA 的小/中/大规模、CPU/PD 双向搬运比较 12 个组合，再在 INT4、INT8、静态 FP8、MLA 和三种 packed 布局上比较入选组合，共 68 个场景。候选包含 grid=24/36、BLOCK=256/512/1024/2048、warps=4/8、循环 stages=1/2/3/4；固定配置 24/1024/8/1 的速度几何平均约为旧实现的 2.12 倍。小块不保证更快，本轮最多增加约 1.74 μs。旧实现为未扩宽 view 的 grid=36、BLOCK=256、默认 4 warps、普通 range 循环。计时预先准备页索引并预热编译，随机顺序执行 CUDA Graph，三轮取中位数；每个候选均与旧实现逐字节对照。 - -进一步使用 16 个独立地址数据块连续发射，覆盖五类专用实现的 CPU/PD 双向操作，共 20 个场景。每块 16 层、1025 个有效 token、64 token/页，使用乱序物理 token 与页索引；GQA/量化 GQA 为 8 个 KV head、head_dim=128，MLA 使用单 latent head。核函数七轮取平均，完整 API 五轮取平均；块间不做同步。下面是代表性结果（GB/s，按有效载荷计），完整 API 包含视图准备、页索引上传、发射及末尾同步，不含 buffer 分配或 PD 网络收发。新配置未测试与推理计算并发的影响。 - -| 布局及方向 | 旧核带宽 | 新核带宽 | 旧完整 API | 新完整 API | -| --- | ---: | ---: | ---: | ---: | -| GQA CPU→GPU | 11.72 | 43.83 | 11.63 | 43.26 | -| GQA GPU→CPU | 29.59 | 50.07 | 29.22 | 49.56 | -| GQA PD 页→KV | 30.44 | 205.40 | 30.08 | 198.41 | -| INT4 CPU→GPU | 7.34 | 43.91 | 7.03 | 42.47 | -| 静态 FP8 CPU→GPU | 6.10 | 43.73 | 6.08 | 42.13 | -| MLA CPU→GPU | 11.87 | 45.58 | 11.50 | 43.81 | -| packed FP8 MLA CPU→GPU | 7.02 | 43.55 | 6.90 | 39.12 | -| packed FP8 MLA PD 页→KV | 16.93 | 169.67 | 16.54 | 48.24 | - -packed MLA 等多字段布局的 Python 准备及逐字段发射开销仍会限制完整调用吞吐；因此核带宽不能直接当作完整 CPU cache 或 PD 吞吐。字段融合仍保留为后续优化，不在本轮引入。 - - - - -DP 直接复制两端 State 的字段，buffer 基址、容量和 stride 从各自 State 获取。层维度完整、同序复制；字段类型、形状、编码及固定参数的检查仍由具体实现完成。不保存额外的复制映射。 - -DP 只复制对应 TP rank 的完整分片，不支持通过 State 接口拼接多个 TP 分片。不同 token 来自不同 DP 时，先按来源分组索引，再逐项调用 `copy_tokens_from_dp`；同一批内配对源与目标 token 索引。 - -源与目标布局完全一致、分片一一对应时，可以采用同构布局的复制 kernel，但须先确认源目标 stride 等条件成立。重分片或编码变化由具体子类明确声明支持并提供实现。 - -#### main/draft 与不同 TP 的例子 - -假设 main 的 full-attention 逻辑层为 `[3, 7]`，本地压缩存储层为 `[0, 1]`；一页 16 个 token，全局 4 个 KV heads，每 head 64 个 FP16 元素。以 K、V 分区的示例格式组织 `kv` 页： - -```text -byte 0 : main_gqa_fp16_state K [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB -byte 16384 : main_gqa_fp16_state V [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB -byte 32768 : draft0_mla_fp8_state 的数据区域，按 draft 自身格式定义 +```python +state.contains_logical_layer(logical_layer_id) +state.get_layer_view(logical_layer_id) ``` -在该示例格式下，main 的 P 端 TP=2 时，rank 1 的本地 head `[0, 2)` 映射到页内全局 head `[2, 4)`；D 端 TP=4 时，rank 3 的本地 head `[0, 1)` 从页内全局 head `[3, 4)` 恢复。两端共享页格式，使用不同的局部分片映射。 +一个 State 可以服务多个逻辑层；共享层映射由具体配置描述，模型负责保证只有生产层写入共享数据。 -如果 draft 使用 TP 副本形式的 MLA，draft 区域由选定的 P 端副本写一次，D 端各副本读取该区域。它不沿用 main 的 head 偏移公式。main 与 draft 即使共享 token allocator，也分别取得自己的区域映射。 +## 4. 公共接口契约 -CPU 组合页包含 full-attention 和状态快照，实际 KV token 数由索引数量确定、格式由外层描述；PD 使用 `kv` 与 `att_state` 两类任务，各自拥有布局，不强制把 KV、Conv、SSM 和窗口 放进同一页。采用按 TP rank 分块的格式时，必须明确其 TP 配置约束；跨 TP 配置恢复需要相应格式和转换能力。 +以下接口是行为契约，具体配置参数和返回视图由实现定义。基类不承载具体算法字段。 -#### 初始化与提交时的校验 - -布局由部署配置和 State 生成；搬运时校验资源名称、字段覆盖、dtype、字节对齐、页内地址边界及层/轴映射范围，TP 切片指定分片范围与副本写入者。各资源的实际写入区域不得相互覆盖；packed 别名按实际物理范围去重。涉及多个布局的描述以实际访问地址为准，不能仅比较 shape。 - -数据操作保留必要的长度和对应页数检查，直接使用调用方提供的有效索引；源数据须已就绪且保持有效。`AttentionStateManager`校验所有参与 State、TP 分片及调用方传入的长度；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局或参与资源改变时，重新生成映射，使用对应的区域偏移。 - -### 6.6 子类实现与直接调用示例 - -以下为伪代码，layout/mapping 在初始化时由管理器组织并经具体 State 校验，调用过程保持直接。 +### 4.1 FullAttentionState ```python -class GQAFullAttentionState(FullAttentionState): +class FullAttentionState: + def get_state_type(self) -> str: ... + def allocate_storage(self, token_capacity: int) -> None: ... + def get_token_bytes(self) -> int: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_storage_layer_index(self, logical_layer_id: int) -> int: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: ... + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: ... + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: ... + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: ... + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: "FullAttentionState", + ) -> None: ... + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... def write_to_pd_page( - self, token_indices, pd_page_buffers, *, pd_page_index: int, layout, - ) -> None: - gqa_kv_to_pd_buffer( - kv_buffer=self.kv_buffer, - token_indices=token_indices, - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - layout=layout, - ) - + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: ... def read_from_pd_page( - self, pd_page_buffers, token_indices, *, pd_page_index: int, layout, - ) -> None: - gqa_pd_buffer_to_kv( - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - kv_buffer=self.kv_buffer, - token_indices=token_indices, - layout=layout, - ) -``` - -CPU 操作按路径选取该 State 的位置映射： - -```python -state.write_to_cpu( - src_token_indices, cpu_page_buffers, cpu_page_indices, - layout=cpu_write_layout, -) -# 当前 stream 上记录 event；所有资源写入完成后发布 CPU 页 ready。 - -state.read_from_cpu( - cpu_page_buffers, cpu_page_indices, dst_token_indices, - layout=cpu_read_layout, -) -# 后续计算与写入建立 stream/event 依赖。 -``` - -DP 操作按同一 TP 分片的来源执行；多来源时重复调用，批次由`AttentionStateManager` 组织： - -```python -dst_state.copy_tokens_from_dp( - src_token_indices=src_token_indices, - dst_token_indices=dst_token_indices, - remote_dp_full_att_state=src_state, -) -``` - -PD 两端分别传入完整发送/接收页集合与本地页索引；两端页索引可以不同，区域与 TP 位置来自各自 layout： - -```python -p_state.write_to_pd_page( - src_token_indices, send_pages, - pd_page_index=send_page_index, - layout=p_layout, -) -# 全部参与资源/分片的打包完成后，调度层发送。 - -# D 端接收完成并建立 stream 依赖后执行。 -d_state.read_from_pd_page( - recv_pages, dst_token_indices, - pd_page_index=recv_page_index, - layout=d_layout, -) -# 所有目标资源/分片恢复完成后允许 decode。 -``` - -调度层负责网络、同步和生命周期。方法返回不等于数据已经就绪，发送页在网络发送完成前不能复用，源槽位在读取完成前不能释放。KDA/GDN 状态的统一操作见 6.7 节，窗口压缩器的操作见 6.8 节；它们共用适用的位置描述，不继承 `FullAttentionState`。 - -### 6.7 BoundedAttentionState：有界状态的统一操作接口 - -`BoundedAttentionState(ABC)` 约束有界 attention 状态的操作，当前由 `KDAState`、`GDNState` 实现，后续 sliding-window 等实现也按相同契约组织资源操作。基类不声明 Conv、SSM、窗口 KV 或 MTP buffer，不包装另一层状态对象。`AttentionStateManager`直接持有子类实例，计算 backend 通过具体层视图访问所需运行态。以下存储与 MTP 示例以 KDA/GDN 为例，不规定其他算法必须使用相同布局。 - -```text -BoundedAttentionState # 抽象操作契约 -├── KDAState # Conv + KDA 递推状态 + 原生 MTP 布局 -├── GDNState # Conv + GDN 递推状态 + 原生 MTP 布局 -└── 后续 SlidingWindowState 等有界状态实现（各自定义布局与页池） + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: ... ``` -`bounded_att_*_allocator` 的范围仍可覆盖 sliding-window 等有界状态，但使用该 allocator 不要求继承 `BoundedAttentionState`。 - -#### 状态操作约定 - -- forward 通过层视图更新 Conv/SSM；MTP 提交选择已接受状态。提交可以只更新游标。当前大小页保存和 DP/PD 导出仅处理 prefill 基准状态，KDA/GDN 的源游标必须为 0，不支持 MTP 提交后保存或导出。 -- 页保存只接受实际捕获的状态位置。跨过中间边界而未保留其状态时，只能预留槽位，不能发布为有效页。 -- 大小页 pool 保存独立快照，默认使用 CPU pinned memory；DP/PD 操作请求运行态。两者使用不同索引和生命周期。 - -#### KDAState 内部布局示例 - -设线性层数为 `L`，请求槽位容量为 `N`（已经包含 HOLD），本地 head 数为 `H`，head 维度为 `D`，Conv 历史宽度 `W = kernel_size - 1`，MTP 额外候选数为 `S`。该 KDA 布局中，Q/K/V 均为 `H * D` 通道： - -| 存储 | 示例形状 | 内容 | -| --- | --- | --- | -| GPU `conv_state` | `[L, N, 3*H*D, W+S]` | 历史窗口及验证时暂存的候选尾部 | -| GPU `ssm_state` | `[L, N*(S+1), H, D, D]` | 每请求 `S+1` 个递推状态槽位 | -| GPU `mtp_state_index` | `[N]`，int32；S=0 时可为 None | 当前已接受状态的原生选择索引 | -| 某个 pool 的 Conv | `[C, L, 3*H*D, W]` | 一个确定前缀的规范化历史窗口 | -| 某个 pool 的 SSM | `[C, L, H, D, D]` | 同一前缀对应的递推矩阵 | - -大、小 pool 的 C 分别取管理器给出的容量；两者保存同种状态快照，区别主要在分配域、保存边界和容量，不是按 token 页大小重复展开运行态。pool 不持有空闲列表。默认使用 CPU pinned 快照存储，设备和 pin memory 选项由 `initialize_page_pools()` 的调用方指定，不在 State 中重复保存。 +所有表示 token 的索引参数由调用方准备为 CUDA Tensor。State 不在接口内部做设备、dtype 或索引格式转换，不分配请求映射。量化和 packed State 的复制必须保持其完整物理表示，包括必要的 scale、indexer 和 padding。 -GDN 的 Q/K 与 V head 数和维度可以不同，其 Conv/SSM 形状由自身实现确定，不把上述 `3*H*D` 或方阵要求放进基类。KDA/GDN 即使 shape 相同，也不能互换状态内容。 - -#### 必须实现的公共接口 - -所有下列方法均为抽象方法。子类可继承已验证兼容的具体实现；对于未支持的 checkpoint、MTP 或传输功能，也必须明确实现为报错并在对应接口中说明限制，不能依靠基类空方法静默成功。初始化只要求实际启用的功能可用。 +### 4.2 BoundedAttentionState ```python -from abc import ABC, abstractmethod -from typing import Any, List, Literal - -PageKind = Literal["big", "small"] - - -class BoundedAttentionState(ABC): - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - @abstractmethod +class BoundedAttentionState: def get_state_type(self) -> str: ... - - @abstractmethod def initialize_page_pools( - self, big_page_capacity: int, small_page_capacity: int = 0, *, device="cpu", pin_memory: bool = True + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, ) -> None: ... - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - @abstractmethod def contains_logical_layer(self, logical_layer_id: int) -> bool: ... - def get_layer_view(self, logical_layer_id: int) -> Any: ... - - @abstractmethod - def reset_bounded_att_state(self, req_indices: List[int]) -> None: ... - - @abstractmethod + def reset_bounded_att_state(self, req_indices: list[int]) -> None: ... def save_bounded_att_runtime_to_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], + page_slot_ids: list[int], ) -> None: ... - - @abstractmethod def restore_bounded_att_runtime_from_pages( self, - req_indices: List[int], - state_token_lens: List[int], + req_indices: list[int], + state_token_lens: list[int], page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: ... - - @abstractmethod - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: ... - - # ======================================================================== - # 三、CPU Cache 数据交换：读写大小页快照 - # ======================================================================== - @abstractmethod - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... - - @abstractmethod - def write_state_pages_to_cpu( - self, page_kind: PageKind, page_slot_ids, cpu_page_buffers, cpu_page_indices, - *, layout: StateTransferLayout, cpu_page_ready_flags: List[bool] | None = None, - ) -> None: ... - - @abstractmethod - def read_state_pages_from_cpu( - self, cpu_page_buffers, cpu_page_indices, - *, big_page_slot_ids: List[int], layout: StateTransferLayout, - ) -> None: ... - - # ======================================================================== - # 四、DP 迁移 - # ======================================================================== - @abstractmethod - def copy_big_pages_from_dp( - self, src_page_ids: List[int], dst_page_ids: List[int], - remote_dp_bounded_state: BoundedAttentionState, + page_slot_ids: list[int], ) -> None: ... - - @abstractmethod - def copy_runtime_state_from_dp( - self, src_req_indices: List[int], dst_req_indices: List[int], - remote_dp_bounded_state: BoundedAttentionState, + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, ) -> None: ... - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - @abstractmethod + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_state_pages_to_cpu(...) -> None: ... + def read_state_pages_from_cpu(...) -> None: ... + def copy_big_pages_from_dp(...) -> None: ... + def copy_runtime_state_from_dp(...) -> None: ... def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... - - @abstractmethod def write_to_pd_page( self, req_idx: int, @@ -1507,8 +227,6 @@ class BoundedAttentionState(ABC): state_token_len: int, layout: StateTransferLayout, ) -> None: ... - - @abstractmethod def read_from_pd_page( self, pd_page_buffers: torch.Tensor, @@ -1519,392 +237,153 @@ class BoundedAttentionState(ABC): ) -> None: ... ``` -bounded State 使用 `write_state_pages_to_cpu` / `read_state_pages_from_cpu` 表达页快照的写出与读回，区别于请求运行态操作；资源管理器通过 `write_req_to_cpu` / `read_req_from_cpu` 操作显式传入的 token 索引和状态页槽位，不接收 `InferReq`。`MultiLevelKvCacheModule` 提供接收 `InferReq` 的 `write_req_to_cpu` / `read_req_from_cpu`：负责查找 radix/请求持有的页、分配目标资源、跳过 ready 页和登记请求页引用；资源层负责跨 State 搬运、末页运行态恢复。CPU 读回含 bounded State 时，调用方还须传入 `state_token_len`，取最后一个 CPU 页的逻辑终点；即使本批仅加载前缀末段，也不能用 `token_indices.numel()` 替代。无需独立的 hybrid CPU cache helper 文件。State 写出接口保留 `page_kind` 和 `page_slot_ids`，由管理器按来源页池分别调用；管理器对外通过两类槽位列表表达来源。 - -CPU 页读写的 `cpu_page_indices`、写出的 `page_slot_ids` 与读回的 `big_page_slot_ids`，以及大小页保存和恢复的 `req_indices`、`state_token_lens`、`page_slot_ids` 均仅接受 `List[int]`，单请求也传长度为 1 的列表；实现只断言 list 类型，不进行 Tensor 或其他序列转换。`state_token_lens` 表示各快照已经执行到的绝对 token 长度，不包含新采样但尚未执行的 token，也不是本批搬运的 token 数。调用方保证长度与选定快照一致；页池不新增长度记录。KDA/GDN 的 Conv/SSM 布局不依赖长度定位，保留此参数供 ring 等实现使用。 - -PD 状态页按单请求处理，全部 bounded State 的读写均接收 `state_token_len: int`，取自同一 PD 任务的 `end_kv_index`。管理器透传该值，页内不新增长度元数据。`req_indices` 是 ReqManager 已分配的请求槽位，`page_slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。DP 运行态复制的 `src_req_indices`、`dst_req_indices` 和大页复制的 `src_page_ids`、`dst_page_ids` 同样仅接受 `List[int]`，实现不进行类型转换。manager 的 `copy_from_dp` 统一聚合 token、大页和请求运行态复制；manager 接收单个 `src_req_idx` / `dst_req_idx`，向 State 传入长度为 1 的列表，不隐式创建历史快照。 - -#### 运行态与推理访问 - -| 接口 | 必需语义 | -| --- | --- | -| 初始化 | 构造函数按 `request_capacity` 和 `mtp_step` 创建运行态及层映射，大小页池先为 `None`；不创建 allocator，不分配请求号或页槽位 | -| `initialize_page_pools(...)` | 容量规划完成后创建大小页池，覆盖共享 allocator 的槽位域；容量为零的页池保持 `None`，不重建运行态或重置 MTP 游标 | -| `contains_logical_layer(logical_layer_id)` | 判断指定逻辑层是否属于当前 State 的层映射，不依赖实际存储分配，共享层同样返回 True | -| `get_layer_view(logical_layer_id)` | 返回该层计算需要的直接引用；KDA/GDN 返回该层的 Conv、SSM 两个 Tensor；MTP 游标和步数直接从 State 读取。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | -| `reset_bounded_att_state(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | -| `commit_bounded_att_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | - -`accepted_steps` 为逐请求接受的、已经执行过 forward 的输入步数，包含本轮基准输入步，不是“接受的 draft token 数”或新采样输出数。取值范围为 `[1, 本轮实际验证宽度]`；KDA/GDN 的原生索引为 `accepted_steps - 1`。实现不应把不存在的零步回滚当作有效操作。未启用 MTP 时，合法的一步提交可为无数据移动的操作。 - -GDN/KDA Backend 不保存 State 引用。prefill/decode 按角色、算法和逻辑层号调用管理器的 `get_state()`；MTP decode 的 `init_state()` 在需要游标时局部获取当前角色唯一的对应算法 State，将 `mtp_state_index` 显式传给 fixed/dynamic 初始化函数。此初始化仍要求每个角色只有一个对应算法 State，不假定多个 State 共用游标。 - -KDA 的提交可以只更新 `mtp_state_index`，下一轮 kernel 根据该游标选择输入状态。当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0，直接读取 Conv 的 `[:W]` 窗口和 SSM 的 `req_idx*(S+1)` 基准槽位。导入时写到目标的基准位置，并将游标设为 0。实现中以 TODO 标明：若以后支持 MTP 提交后传输，需要按已接受索引 j 选择 Conv 的 `[j:j+W]` 窗口和 SSM 的 `req_idx*(S+1)+j` 槽位。 - -prefill 的最终状态、普通 decode 的结果和已提交验证状态都可以成为“当前已接受状态”；执行流程必须保证所有相关层已完成且选择信息与数据一致。验证结果尚未提交时，不允许将临时候选作为当前状态发布。大小页保存仅发生在 prefill 阶段，此时 `mtp_state_index` 为 0，直接复制 Conv 的 `[:W]` 窗口及 SSM 的 `req_idx*(S+1)` 基准槽位，不需要先捕获候选状态。实现中以 TODO 标明：若以后支持 decode/MTP 后保存，需要改为按已接受游标选取数据。 - -#### 状态页：只保存真实存在的前缀状态 - -| 接口 | 必需语义 | -| --- | --- | -| `save_bounded_att_runtime_to_pages(...)` | 将每个请求 prefill 完成后的基准状态直接写到指定 pool 槽位；req、state_token_lens 与 slot 逐项对齐，ring 实现按长度确定窗口；前缀关联由调用方维护 | -| `restore_bounded_att_runtime_from_pages(...)` | 按 state_token_lens 从选定槽位恢复运行态及 ring 位置等控制信息；源快照保持不变，目标 MTP 从规范化状态继续 | - -调用方负责将快照关联到实际执行后的位置；KDA/GDN 的页保存契约只覆盖传入请求的 prefill 完成状态；`state_token_lens` 必须是该实际状态的绝对长度，传入其他长度不能生成历史快照。chunk 跨过中间边界但没有保留对应运行态时，预留槽位不能作为已保存的页面发布。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 - -保存或 CPU 导入时直接写入指定槽位；异步完成前由上层保持 pending，完成后才发布或读取。发生失败直接透传异常，调用方不发布该批页面，并等已提交操作结束后才可回收。恢复前由调用方保证前缀关联和来源就绪，State 按兼容布局搬运数据。 - -同一槽位数字在 big/small pool 中属于不同分配域。main/draft 的相关 State 共享分配的槽位 ID，各自保存快照数据；管理器组织所有参与 State 的保存和恢复，缓存调用流程在数据就绪后发布；最终释放时，调用方解除引用并确认使用结束，再直接向 allocator 归还一次。 - -#### CPU、DP、PD 的状态单位 - -CPU 页与本地大小页池之间使用 `get_cpu_transfer_regions()` 返回的映射：`region.field` 选择页池中的 Conv/SSM 字段，`byte_start/byte_end` 和 `shape` 定位 CPU 全局字段，`axis_mappings` 将本地 TP 分片对应到全局 head/channel 范围。`local_start` 固定指本地字段，`external_start` 固定指外部页，实际读回时只反转搬运方向，不交换两者含义。`read_state_pages_from_cpu` 显式准备目标 `pool_page_fields`，共用搬运流程直接展开这些映射，不通过 `_copy_pair` 隐藏切片过程。 - -窗口顺序的约定为：请求 ring → 按长度保存为逻辑时间顺序的大小页快照 → CPU 页；恢复时为 CPU 页 → 大页快照（完整页用普通槽位，碎页用 LOAD 槽位）→ 按绝对长度恢复请求 ring。CPU 页与页池之间只处理 TP 映射，不执行环形展开或恢复。 - -| 路径 | 传输单位与接口职责 | -| --- | --- | -| CPU 写出 | `write_state_pages_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，长度由缓存任务携带，不写入页内；每个待发布页的参与状态必须完整；CPU pool 到 CPU 页同样是独立的格式搬运 | -| CPU 读回 | State 的 `read_state_pages_from_cpu` 按 `big_page_slot_ids` 将快照写入大页池；管理器补上碎页 LOAD 槽位，并调用 `restore_bounded_att_runtime_from_pages` 恢复最后一页运行态 | -| DP 大页复制 | `copy_big_pages_from_dp(src_page_ids, dst_page_ids, remote_dp_bounded_state)` 在同一 TP 分片的 DP 副本之间复制大页 Conv/SSM 快照；不分配页、不修改请求运行态或小页，前缀关联与源页生命周期由调用方维护 | -| DP 运行态复制 | `copy_runtime_state_from_dp(src_req_indices, dst_req_indices, remote_dp_bounded_state)` 复制同一 TP 分片的 prefill 基准 Conv/SSM；KDA/GDN 源游标必须为 0，目标游标归零，不修改大小页或其他候选区 | -| PD 写出 | `write_to_pd_page` 将一个请求的 prefill 基准状态写入 `pd_page_index` 指定的状态页区域，按协议关联其前缀长度 | -| PD 读回 | `read_from_pd_page` 按 state_token_len 将完整快照恢复到指定请求槽位及对应 ring 位置，不读取页内长度、版本字段；不自动生成历史 checkpoint | - -CPU 尾页若需要保存 prefill 完成后的运行态，管理器先使用保留槽位调用 `save_bounded_att_runtime_to_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 - -CPU/PD 位置描述使用 `StateTransferLayout`；bounded DP 接口不需要字段映射，直接复制同一算法、同一配置及 TP 分片的对应存储。Conv 区域可用 `(layer, channel, window)` 轴，SSM 用 `(layer, head, key_dim, value_dim)` 轴。它们没有历史 KV 的 token 轴。含 KV 的 CPU 页仍有 token 容量，checkpoint 组件不按该容量重复展开；纯 PD 状态页的 `tokens_per_page` 为 None。各资源偏移、TP 范围和副本写入者由管理器组合，State 直接执行方法。 - -当前 DP/PD 导出 prefill 基准状态，调用方保证 KDA/GDN 源游标为 0，不传 MTP 候选区，也不支持 MTP 提交后导出。目标可按自己的 MTP 步数初始化工作区；DP 两端配置及 TP 分片一致，按各自 MTP 步数计算基准槽位。PD 则通过全局布局与 TP 切片处理两端分片差异。不能因为 KDA/GDN 的 Conv/SSM 形状相同就跨算法恢复。 +`state_token_len` 表示状态已经执行到的绝对前缀长度，不表示当前搬运的 token 数。ring、窗口或游标位置必须由具体 State 根据该长度恢复。请求列表、长度列表和页槽位列表按位置一一对应，关联关系由调用方维护。 -多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；`AttentionStateManager`完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 +### 4.3 AttentionStateManager 聚合接口 -#### 依赖、同步与使用示例 - -资源方法接收请求槽位、状态页槽位、绝对前缀长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 WindowCompressorState；管理器统一遍历 bounded State，各实现分别处理自身快照和运行态迁移。 - -GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 +管理器提供按资源组合调用 State 的入口，至少包括： ```python -# ReqManager 已分配请求槽位；调用方已从管理器持有的 allocator 分配状态页槽位。 -state.reset_bounded_att_state(req_indices) -# backend 经层视图完成 prefill；调用方维护实际 chunk 终点。 -state.save_bounded_att_runtime_to_pages( - req_indices=req_indices, - state_token_lens=state_token_lens, - page_kind="big", - page_slot_ids=page_slot_ids, -) -# 确认所有参与状态已按指定长度保存并完成传输后，缓存调用方发布联合页。 - -state.restore_bounded_att_runtime_from_pages( - req_indices=dst_req_indices, - state_token_lens=state_token_lens, - page_kind="big", - page_slot_ids=page_slot_ids, -) -# target 请求在相关 KV、linear、窗口 均就绪后继续执行。 - -# 验证 forward 完成后，将接受步数按请求提交。 -state.commit_bounded_att_state(req_indices, accepted_steps) -# 后续计算 kernel 使用已提交游标；当前不支持在 MTP 提交后保存大小页或进行 DP/PD 导出。 +manager.initialize_storage(...) +manager.allocate_storage(...) +manager.get_full_attention_states(role=None) +manager.get_bounded_attention_states(role=None) +manager.get_cpu_page_layouts() +manager.get_pd_token_page_layouts() +manager.get_pd_runtime_page_layouts() +manager.copy_full_att_tokens(...) +manager.reset_bounded_att_state(...) +manager.save_bounded_att_runtime_to_pages(...) +manager.restore_bounded_att_runtime_from_pages(...) +manager.commit_bounded_att_state(...) +manager.write_req_to_cpu(...) +manager.read_req_from_cpu(...) +manager.copy_from_dp(...) +manager.write_to_pd_page(...) +manager.read_from_pd_page(...) ``` -各具体实现必须满足数据完整性、接受位置选择和跨布局复制约束。 - -### 6.8 WindowCompressorState：窗口候选与 gate - -K-pool 与 V4 的压缩器共用 `bounded_att/WindowCompressorState`，具体池化、归一化、 -RoPE 和量化仍由各模型自己的计算算子完成。资源层只负责窗口存储、快照及传输。 -模型通过 `create_bounded_att_states()` 返回窗口 State，与 KDA/GDN、SWA 使用同一管理流程。 - -| 场景 | window_size | compress_stride | candidate_dim | 缓存边界历史 | -| --- | ---: | ---: | --- | --- | -| GLM K-pool | 4 | 4 | indexer key 的 head_dim | 无重叠，无需快照 | -| V4 CSA KV / indexer | 8 | 4 | 输出向量维度的两倍 | 保存最近 4 个位置 | -| V4 HCA KV | 128 | 128 | 输出向量维度 | 无重叠，无需快照 | - -`candidate_dim` 不包含 gate；候选和 gate 具有相同形状和存储精度，由同一个 `dtype` -描述。GLM K-pool 随模型精度配置（通常 BF16），V4 使用 FP32。运行态保持一个 Tensor，形状为 -`[layer, request, window_size + mtp_step, 2 * candidate_dim]`，末轴前半为候选、后半为 gate。 -`get_layer_view(logical_layer_id)` 返回该层的拼接视图，计算算子的输入方式保持不变。 -`target="kv"` 表示 attention KV 候选，`target="indexer"` 表示索引 key 候选， -仅区分 State 成员名，不选择模型算法。例如 K-pool 使用 `compressor_indexer_w4_s4_bf16`。 - -大小页保存和恢复必须对齐 `compress_stride`。有重叠时，通过 `WindowPagesPool` 保存 -时间有序的重叠历史，CPU 页搬运同一份拼接字段；无重叠时不分配页池、不占 CPU 页， -整组缓存边界后的计算只依赖新输入。空前缀仍通过 `reset_bounded_att_state()` 初始化， -候选清零，gate 设为负无穷。DP 运行态复制使用相同配置和 MTP 容量的物理 ring。 - -管理器只为声明了 CPU 快照区域的 bounded State 调用 `initialize_page_pools()`。 -窗口 State 在返回布局前断言小页 token 数可被 `compress_stride` 整除；大页由整数个小页组成， -因此也满足此边界。配置了 CPU 页 token 数时，同样检查该页边界，即使窗口没有快照区域也不能跳过。 -共享槽位 allocator 继续供缓存引用和其他 State 使用；不需要快照的 State 不分配对应的页池 Tensor。 - -PD 在任意前缀长度导出完整 `window_size` 逻辑窗口，例如 K-pool 的长度 6 前缀仍需 -携带未完成池化组的原始 key/gate。P/D 的 MTP 容量不必相同：搬运按绝对前缀长度映射 -到各自的 ring，不传输或清理逻辑窗口外的候选槽位。环形搬运 kernel 位于 -`bounded_att/window_compressor/window_compressor_transfer.py`，一次 launch 搬运全部层的候选与 gate。 -MTP 回退依赖保留的历史和绝对位置，不调用缓存恢复入口。 - -## 7. 资源初始化与调用约定 - -### 7.1 初始化与绑定 - -BaseModel 提供只操作当前模型的公共阶段接口；main 调用自己的接口，各 draft 调用各自的接口，模型方法不接收其他模型或组织跨模型循环。`ModeBackend.init_model()` 直接组织 main/draft 构造、共享请求管理器创建与绑定、资源准备、共享分配和各模型运行时初始化，并在全部初始化成功后发布就绪标志。benchmark 和资源验证脚本在各自调用点展开相同流程，不另设跨模型运行时启动封装。模型构造保留配置、权重和计算层的准备，不创建或读取请求管理器。 +管理器只做聚合和公共前置检查；字段编码、ring 恢复、压缩边界和特殊传输由具体 State 实现。某种 State 不参与某条路径时返回空布局或由管理器跳过，不通过无意义的通用适配层伪造支持。 -| 单模型公共接口 | 封装的内部操作 | -| --- | --- | -| `get_req_manager_max_seq_len()` | 返回请求映射所需长度容量，取本模型 `batch_max_tokens`、`max_seq_length` 与 0 的最大值；未配置的值按 0 处理，不表示模型最大上下文能力 | -| `bind_req_manager(req_manager)` | 绑定调用方创建的共享实例，不判断 main/draft，也不通过 `main_model` 获取资源 | -| `prepare_attention_states()` | main 直接创建空的 `AttentionStateManager`，draft 引用 main 的实例；在 KV 内存分区内调用两类 State 工厂并统一注册返回的列表 | -| `init_runtime()` | 绑定本模型 HOLD 页、检查容量、初始化 attention backend 和 hidden collector，完成 kernel 调优、padded-request 和 decode/prefill CUDA Graph 初始化；main 在自身图初始化后执行最大长度检查 | +## 5. 布局和传输契约 -1. 构造 main model，再构造所有 draft model。各模型加载自己的权重，准备配置和计算层，此时不创建 attention State。随后 backend 检查各模型的 `max_req_num` 一致，取所有模型 `get_req_manager_max_seq_len()` 的最大值，在主模型的 `MemoryTag.KV_CACHE` 分区中创建唯一 `ReqManager`，逐模型调用 `bind_req_manager()`。main/draft 共用请求槽位与逻辑 token 到物理索引的映射，模型不再提供 `_init_req_manager()` 私有初始化或子类共享覆盖。 -3. backend 检查 main 和所有 draft 引用同一个管理器，显式调用 `finish_state_registration()`，随后在 KV 内存分区中调用 `load_full_att_calibration(args)`。文件只读一次，管理器按 main、draft0..draft35 顺序累计全部 full State 的物理层数；同角色多 State 按最小逻辑层号排列，同起点时按规范成员名排序。每个 State 内保留 `logical_layer_groups` 的物理层顺序，同组共享逻辑层只占一行，bounded 层和无 attention 的 draft 不占行。现有标定文件仍按物理行组织，没有将逻辑层号直接当文件行号。 -4. 静态 FP8 State 的 `bind_calibration(calibration, *, storage_layer_start)` 负责解析自身标定区间、校验编码和数值、切分本地 TP head，并绑定 K/V、Q 参数。仅此具体实现提供该接口，其他 State 无需空实现。无文件时 per-tensor 使用单位 scale；per-head 必须有 K/V 标定，Q 标定可选：有 Q 标定时使用静态量化，未提供时按本次请求的 Q 分组动态量化，prefill/decode 使用同一规则。随后 backend 调用 `initialize_storage(args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction)` 规划容量并分配存储；原地分配不改变 State 对象身份，所有 full State 共享 allocator 索引域。静态 FP8 在尚未绑定标定时禁止分配实际存储。 -5. 各 Model 继续使用同一个共享管理器。backend 按 main、各 draft 的顺序逐个调用 `init_runtime()`，其中 `_check_mem_size()` 从 `full_att_allocator.size` 更新当前模型的 `max_total_token_num` 并检查容量，随后完成 backend 绑定、warmup、padded-request 和图初始化，再初始化下一个模型。main 在自己的 `init_runtime()` 末尾执行最大长度检查；draft 不执行该检查。每次调用只操作当前模型，计算上下文和 attention backend 按本模型角色访问 State；MTP 游标仍由具体 State 持有。 -6. 全部模型的 `init_runtime()` 成功返回后，backend 清理临时显存并发布就绪标志，再建立 radix cache、共享存储导出及缓存/传输服务。初始化失败直接向调用方传播，不发布就绪标志；重复启动在修改现有资源前断言失败。 +### 5.1 传输类型 -BaseModel 提供两个默认返回空列表的 `classmethod` 扩展点，具体模型只覆盖自己支持的类别。创建接口不读取模型实例、管理器或其他创建接口的临时属性，也不注册 State、不设置 `state_name`。 +- `TransferRegion` 描述一个 State 字段在外部页中的范围、dtype、轴顺序、形状和 TP 映射。 +- `StateTransferLayout` 描述一次完整页面传输的统一页规格和区域集合。 +- `AxisMapping` 描述本地和外部轴的起点、长度及是否循环；循环窗口必须通过显式映射表达，不能依赖调用方猜测。 +- State 返回的布局定义必须稳定；使用 `lru_cache` 缓存，并通过独立副本返回，避免调用方修改缓存对象。 +- 存储分配或页面规格变化后，管理器负责清理相关布局缓存。 -| 创建接口 | 返回值与职责 | -| --- | --- | -| `create_full_att_states(...)` | 返回 `list[FullAttentionState]`；只准备布局，不加载静态标定或分配 KV 容量；可返回多种量化或存储方案 | -| `create_bounded_att_states(...)` | 返回 `list[BoundedAttentionState]`；按请求容量创建运行态，大小页延后分配；CPU 查询使用 `device="meta"` | +CPU、DP、PD 使用同一套布局语义，但入口不同： -`model_logical_layer_ids` 表示当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；main/draft 保留现有全局编号机制。具体模型自行选择各类 State 的参与层和共享分组，该参数不代表物理存储层数，也不要求每个 State 使用全部层。 +- CPU 传输可以处理联合页和 CPU 页就绪标记。 +- DP 传输使用源、目标 token 或页槽位的显式映射。 +- PD 传输必须接收明确的页索引；State 不分配网络页、不修改请求映射。 -`Qwen3NextResourceConfig` 及其 GLM 配置只描述当前模型，保留 `from_model_config()`、`load_from_args()` 和 `get_full_att_layer_indices()` 三个入口。层查询返回模型内部序号，由模型工厂转换为实际逻辑层 ID；不追加 draft 层，不计算物理层映射、Conv/SSM 存储形状或字节数。配置只保存全局 head 数，本地 TP 分片由 State 或后端校验调用方计算。full-attention 维度保存原始特征宽度，GLM 的 indexer 字段和 packed padding 由具体 State 的创建入口处理。 +### 5.2 异步和设备约束 -两个接口均为公共类方法，供模型启动和外部 CPU cache 布局构建调用；子类只覆盖自己支持的资源类型，基类默认返回空列表。接口显式接收模型配置、启动参数、逻辑层号和 TP 信息；bounded 接口另收请求容量和设备，full 工厂不接收标定起点。`prepare_attention_states()` 为 main 直接创建空的 `AttentionStateManager`，draft 引用同一实例；随后用真实请求容量和 CUDA 设备调用模型工厂，并统一按模型角色注册所有返回对象。无 attention 的 draft 仅共享管理器，不创建自身资源。所有模型统一使用 `AttentionStateManager`，模型差异由两类 State 工厂及具体 State 实现表达。ReqManager 的创建与共享关系仍由 backend 组织。 +- State 的 GPU 操作提交到调用方当前流；State 不隐式同步设备或主机。 +- 传输调用方负责建立生产流、通信流、接收流和消费流之间的依赖。 +- NCCL/NIXL 的完成语义由传输层提供；读 PD 页前必须确认接收端数据已完成，不以发送端完成代替接收端完成。 +- 跨设备复制要求调用方保证 peer/IPC 可访问、源数据生命周期有效、索引和目标区域满足实现约束。 +- CPU 页的 pinned/registered 方式由缓存和进程启动策略决定，State 只使用传入的 buffer。 +- 需要 kernel 加速的聚合复制应保持异步语义，避免在公共 State 接口中加入主机同步。 +### 5.3 数据所有权 -需要保留的跨模型阶段边界是：全部模型完成权重加载及 State 准备后才规划和分配共享容量。State 资源在创建完成后通过缓存查询取得；分配完成后的 attention backend 初始化、warmup、padded-request 和图初始化按单模型顺序执行，不要求全部模型先 warmup 再统一捕获图。其他调用方先通过 `get_req_manager_max_seq_len()` 汇总需求、创建共享请求管理器并调用 `bind_req_manager()`，再按相同顺序复用 `prepare_attention_states()` 和 `init_runtime()`，共享分配和就绪标志仍由调用方协调。 +- allocator 所有索引；State 所有运行态 Tensor 和页池 Tensor；ReqManager 所有请求映射；缓存流程所有引用关系。 +- 任何索引只能由其 allocator 归还一次。 +- IPC 浅拷贝不得复制 allocator；DP shared 场景必须提供页池重新注册到当前设备的显式流程。 +- 失败路径必须由调用方释放已分配索引并阻止未完成数据进入可见 State。 -`create_full_att_states()` 由具体模型直接构造 `GQAAttentionConfig`、`MLAConfig` 或 `PackedMLALogicalConfig`，并调用对应的 `create_gqa_state()`、`create_mla_state()` 或 `create_packed_mla_state()`，返回 State 列表；不存在 full 层时返回空列表。模型自行决定 head 数、维度、精度、参与层及共享分组，同族相同方案可以继承工厂。bounded 工厂同样直接调用算法 create；GLM main/draft 分别选择池化窗口的参与层并读取自己的 K-pool 配置，不依赖其他工厂的副作用。 +## 6. 模型初始化和装配 -`AttentionStateManager.initialize_storage()` 由基类提供通用实现，在显式完成创建标记及静态标定绑定后，保留两个阶段：估算内存并直接计算完整页容量、调用 `allocate_storage()` 初始化实际存储。PD 模式在容量规划阶段按需构建 meta State 的 PD 布局，得到传输 buffer 的预留大小。`token_capacity` 为指定的 token 容量，传入 `None` 时使用 `mem_fraction` 评估可用显存并协调各 TP rank 的容量；其余页大小及传输选项读取 `args`。该方法只操作管理器内部资源，不接收模型列表，也不遍历模型回填容量。调用方负责共享管理器检查、KV 内存分区和后续各模型的 `init_runtime()`；模型进入 `_check_mem_size()` 时再从 allocator 取得最终容量。CPU 布局构建使用相同的模型编号和独立 State 定义，保留每个 draft 的全局布局；容量统计通过共享管理器完成。 +每种逻辑 attention 算法提供独立创建入口，量化和编码变体由参数选择或在同一逻辑目录中提供专用创建函数。模型基类只负责统一生命周期，不解析具体 State 的内部字段。 -自动容量规划使用所有模型实际加载后的剩余显存，不再按 main 权重比例预估 draft 占用。静态量化标定须在规划前加载和校验；KV、动态 scale、packed padding 及 HOLD 页计入预算，静态 KV/Q 标定参数占比较小，不在容量规划中单独扣减。按完整物理页选择可容纳的最大 token 容量，各 rank 再统一取最小容量。共享 allocator 的不同 State 存储分别计算，不能因共享索引而漏算 draft buffer。 +推荐流程： -固定配置下，`get_full_att_token_bytes()` 返回所有 full-attention State 的单 token 开销，页开销直接计算为 `page_bytes = token_bytes * page_size`。可分配页数为 `budget // page_bytes - 1`，其中扣除的一页用于 HOLD，再乘以 page size 得到 token 容量；容量查询接口不接收 token 容量或页数，也不通过两次总占用查询求差值。索引范围由 `FullAttKVAllocator` 的已有断言校验。 - -PD 容量规划使用 `manager.get_pd_page_bytes() * page_count` 计算预留空间，依赖方向是先布局、后页大小。FullAttentionState 的 `get_pd_transfer_regions()` 根据全局 `pd_kv_page_size` 返回实际字段区域，包含各 field 的 dtype 对齐及实现要求的内部 padding。manager 同时平移字段区间的起点与终点，以 `max(region.byte_end)` 确定当前 State 的数据结束位置，向上对齐到 8 字节后记录为 layout 的 `byte_end`，并作为下一个 State 的 `byte_start`。所有 State 排列完成后，`max(layout.byte_end)` 就是 KV 页大小。GPU packed 行中未被任何传输字段覆盖的 padding 不进入 PD 页。 - -bounded（包括压缩窗口） 的运行态布局使用同一规则：按字段的实际最大结束位置确定 State 占用，不另外维护纯数据字节数估算，也不累加多个映射共用的同一区域。它始终返回各运行态 State 自身的起止区间，不依赖 full-attention 布局，运行态总占用取 `max(layout.byte_end)`。`get_pd_page_bytes()` 返回 KV 页与完整运行态页占用的最大值；没有参与资源时返回 0。有 KV 且 KV 页小于运行态页时记录英文 warning，自动扩大实际传输页，保持 `pd_kv_page_size` 的 token 数和两类布局偏移不变，不再因两类页大小关系而中止初始化。PD 调用方按这个返回值分配和选取实际传输页，再以该页起点执行对应任务的布局；例如运行态布局为 80 字节、KV 页为 128 字节时，实际传输块为 128 字节，运行态仅操作自己的 80 字节范围。 - -PD master 负责校验连接参数，长度通过传输任务的 `end_kv_index` 传递；资源层不附加 PD 页布局指纹及长度、版本字段。两个布局 getter 与页大小查询均使用 `lru_cache`，存储初始化后一起清除缓存。 - -P/D 传输进程为每次 KV 任务创建普通 CPU token 索引 Tensor，在 `copy_cuda_stream` 中调用 `cuda(non_blocking=True)` 上传，后续 State 拷贝沿用同一流；预热使用相同路径。不预分配或保留 pinned 索引页缓冲，普通 CPU 内存也不保证 H2D 全程异步。`att_state` 任务按请求槽位搬运，不创建或上传 token 索引。 - -CPU 使用相同的依赖方向：`get_cpu_page_bytes()` → `get_cpu_page_layouts()` → 各 State 的 `get_cpu_transfer_regions()`。布局 getter 读取全局 `cpu_cache_token_page_size`，在管理器内部直接展开聚合：依次取得 full-attention 和 bounded 的字段区间，平移为页内位置，以 `max(region.byte_end)` 向上对齐到 8 字节作为当前 State 的终点。页大小查询再取 `max(layout.byte_end)` 并按 16 字节对齐。字段之间的 dtype 对齐及 State 实现要求的内部 padding 均计入容量；未参与传输的 GPU packed 行 padding、静态标定参数不占 CPU 页。bounded 每页只保存一份快照，不乘页内 token 数，也不包含 MTP 候选区、游标或其他页池槽位。无重叠窗口不参与 CPU 联合页；有重叠窗口保存必要历史，PD 则传输完整窗口。 - -各具体 State 的 `get_cpu_transfer_regions()` 使用 `@lru_cache(maxsize=None)` 缓存不可变区域元组， -并在其外层添加 `@deepcopy_result`，每次先取得缓存结果、再返回深拷贝，调用方不持有缓存原对象。 -装饰器保留 `cache_clear`、`cache_info` 和 `cache_parameters`，管理器仍可清理缓存。 -无需 CPU 快照时返回空元组。首次查询后 State 配置及全局 CPU 页 token 数保持不变； -当前各实现的区域不依赖存储容量或设备；为兼容定制布局受实际物理页规格影响的情况， -管理器完成存储分配后仍统一清除 State 与聚合布局缓存。子类覆盖此 getter 时须保留缓存能力。 - -State 的全局存储字节数查询继续描述存储开销，不用于推导 CPU/PD 传输容量。CPU 联合页不预留页头、长度或版本字段。布局及其派生页大小均按需缓存，实际存储分配后一起清理，避免继续使用 meta 准备阶段的布局尺寸。 - -State config 保存未切分的全局 head 数、当前 `tp_rank` / `tp_world_size` 及具体分布规则，本地数量使用只读属性计算，不保存第二份可独立修改的数值。CPU/PD 全局页和 DP 运行态映射直接读取对应 State 的 config，管理器不维护 `_full_head_counts`、`_bounded_head_counts`。每个 draft 的 State 独立保存完整 config，本地形状相同也不意味着全局布局相同，容量规划和传输均读取各自配置。 - -| 配置 | 保存的信息 | 计算属性 | -| --- | --- | --- | -| `GQAAttentionConfig` | 逻辑层分组、全局 KV head 数、`head_dim`、KV 浮点 `dtype`、TP 坐标与 `replicated` | `head_num` 返回当前 rank 的单份 K 或 V 的 head 数 | -| `GroupQuantizedAttentionConfig` | GQA 几何与 TP 信息、`bits`、`group_size`、动态 `scale_dtype` | `head_num`；KV 编码固定为 INT4/INT8，scale 精度独立描述 | -| `StaticFP8AttentionConfig` | GQA 几何与 TP 信息、`scale_mode`，不传物理 dtype 或标定 Tensor | `head_num`；FP8 编码固定，标定通过独立流程绑定 | -| `MLAConfig` | 普通 MLA 的逻辑层分组、`latent_dim`、未压缩 latent 的 `latent_dtype` 和 TP 坐标 | `create_mla_state()` 只创建普通 latent State | -| `PackedMLALogicalConfig` | 带 indexer 的逻辑层分组、`latent_dim`、packed `format`、indexer 宽度和 TP 坐标 | `create_packed_mla_state()` 再转换为具体物理布局；浮点精度由 `format` 决定 | -| `PackedMLAAttentionConfig` | 持有 `PackedMLALogicalConfig`，只保留量化分组等物理布局参数 | `kv_bytes`、`token_bytes`、`indexer_offset` 从逻辑格式、NoPE/RoPE 宽度和 scale 数量推导；FP8 行使用内部 uint8 buffer 承载混合字段，不另设 `storage_dtype` 属性 | -| `ConvSSMStateConfig` | 层、`global_k_heads`、`global_v_heads`、TP 坐标及 Conv/SSM 维度/dtype 信息；不保存资源身份 | `num_k_heads`、`num_v_heads` 分别返回本地 head 数，`conv_channels` 返回本地卷积通道数 | - -GQA 的 `replicated=True` 表示每个 rank 保存全部 head；MLA 固定复制单份 latent，不设置该参数。GQA 的 `replicated=False` 时按 head 均匀切分，全局 head 数小于 TP 时按组复制。KDA/GDN 的 K、V head 分别应用切分或分组复制规则。布局 helper 根据全局数量和本地数量计算 head 起点及副本写入者。只改变 config 的 TP 坐标即可重新计算目标分片，无需依赖来源 manager 的额外映射；实际数据迁移仍通过传输接口完成。 - -模型类统一暴露 `normalize_model_config(model_config)`：不创建实例、不加载权重或初始化 CUDA，返回新的有效配置，重复调用保持结果不变。模型启动读取配置后、CPU 服务创建 meta State 前，以及各模型 State 工厂，都使用这个入口。DFlash 的嵌套覆盖、原生 MTP 的自有层数等差异由子类实现，不再保留 `_resource_config` 这条单独的转换路径。`basemodel_utils.py` 只处理通用 text_config 展开和字段别名，不反射调用模型方法，也不选择 attention 算法。 - -DFlash 先合并 `dflash_config`，再同步 `num_hidden_layers` / `n_layer` 等别名，避免运行时与 CPU 页布局使用不同层数。GLM 原生 MTP 的有效层数描述 draft 自身；`mtp_checkpoint_layer_start` 单独保留 checkpoint 中接在 main 层后的权重起点，用于 transformer 和 pre/post 权重加载，重复规范化不会改变这个起点。 - -V4 同样用 `num_hidden_layers` / `n_layer` 表示当前模型自有层数，DSpark 不另存一份层数，也不重复覆盖运行时的 `layers_num`。DSpark 的 `draft_layer_start_id` 表示整体执行序列中的 draft 起始逻辑层号；`stage_id` 由当前逻辑层号减去该起点得到，权重前缀使用 `mtp.{stage_id}`。普通 V4 MTP 的 `draft_weight_layer_id` 则用于权重构造和量化配置索引，各 draft 实例的执行逻辑层号仍由启动流程独立确定。两种起点语义分开命名，重复规范化不改变已记录的值。V4 模型组合配置统一使用 `compress_ratios` 表示逐层压缩倍率。 - -```mermaid -flowchart LR - A[模型启动读取配置] --> C[模型类 normalize_model_config] - B[CPU cache 读取配置] --> C - C --> D[模型专有覆盖与通用别名同步] - D --> E[模型初始化及权重加载] - D --> F[模型 State 工厂] - F --> G[实际资源分配或 CPU meta 布局] +```text +模型读取配置 + -> 创建本模型拥有的 full/bounded State 列表 + -> manager.register_state(role, state) + -> 所有 main/draft 完成装配 + -> manager.finish_state_registration() + -> 绑定标定或其他外部参数 + -> manager.initialize_storage(...) + -> 模型和 backend 通过 manager.get_state(...) 获取视图 ``` +模型可以让一个 State 服务多个逻辑层，也可以为同一模型创建多个逻辑方案。层映射必须由 State 配置完整描述，推理层不得用物理存储下标推断逻辑层归属。 -GQA、MLA、packed MLA 的选择以及全局逻辑层分组直接写在模型的 `create_full_att_states()` 中。普通模型按 `model_logical_layer_ids` 创建独立层，Qwen/GLM 选择混合模型的 full 层，Gemma4 根据共享关系将多个逻辑层放入同一物理分组。各模型不再提供 `_full_attention_options`。算法级 `create_gqa_state()`、`create_mla_state()` 和 `create_packed_mla_state()` 分别负责具体存储布局及算法专有的 TP 副本规则。 +## 7. 目录和代码组织 -CPU 服务读取 main 和各 draft 自己的配置，直接调用模型类的 full/bounded 工厂，不构造模型实例。返回的 State 列表注册到临时管理器，再由 `get_cpu_page_layouts()` 聚合联合页。CPU 页可以包含同一模型的多个 State；静态标定不参与布局计算。无 attention 的 draft 由模型准备流程和 CPU 布局调用方跳过。`lightllm/utils/cpu_cache_layout.py` 负责 CPU 服务布局构建；标定加载由管理器组织，具体格式由静态 FP8 State 解释。 - -资源容量确定后，buffer 地址及已绑定视图保持稳定。初始化 helper 不持有推理请求映射，也不引入请求级生命周期接口。后续模型扩展通过准备自己的 State 和实现所需 backend 完成；不需要将 draft 限定为 full attention。 - -### 7.2 缓存保存与恢复 - -没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_bounded_att_state(req_indices)`，初始化全部 bounded 运行态（包括窗口压缩器）。全部初始化完成后再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 - -调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_bounded_att_runtime_to_pages`。`AttentionStateManager` 将同一组槽位传给全部参与 State 保存；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 - -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 - -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 - -MTP 接受位置由模型执行流程选择，供后续计算使用；当前大小页保存和 DP/PD 导出发生在 prefill 基准状态阶段。KDA/GDN 只保存实际捕获的状态；跨越但未保存的中间位置不可作为有效页，也不能发布为 CPU 缓存页；执行流程需在待缓存页边界安排真实状态捕获。管理器的数据操作不承担一次 forward 的提交或请求长度更新。 - -### 7.3 同步、引用与回收 - -请求流程负责请求暂停、结束及请求槽位回收；缓存功能流程负责节点引用和淘汰时机。归还资源索引前，调用方应停止相关使用、满足异步完成条件并解除相应引用。ReqManager 清理请求映射；对于状态页，调用方直接向对应 allocator 归还一次。 - -有缓存、请求或传输引用的数据必须保持有效。释放一份缓存数据不等于结束某个请求；请求结束也不表示其已共享缓存应被释放。调用方负责将具备回收条件的索引归还 allocator，管理器内部保留的页内容不参与引用计数。 - -保存、导入或复制失败时，不发布本次不完整数据；调用方在操作完成且引用解除后，直接回收本次不再使用的分配，已完成的共享缓存保持有效。CPU/DP/PD 操作遵守 stream/event 依赖，调用方在继续计算、复用 buffer 或回收索引前确认完成。 - -## 8. CPU、DP 与 PD 传输设计 - -### 8.1 具体 State 实现传输，`AttentionStateManager` 组织组合 - -CPU cache 的请求入口位于 `MultiLevelKvCacheModule`，不再单设混合缓存辅助文件： - -- `read_req_from_cpu(req, cpu_page_indices, page_lengths, *, start_token_len)` 为新增 KV 分配索引，补齐首个 CPU 页已有命中部分，为完整 bounded 页分配大页槽位，再调用资源管理器 `read_req_from_cpu`。成功提交后登记请求的大页引用并更新 KV 长度；CPU cache 加载按必须成功的流程执行，异常直接向上传播，不在此处捕获、同步回收或继续执行。 -- `write_req_to_cpu(req, cpu_kv_cache_stream)` 先按 KV 长度和碎页规则申请 CPU 页；实际页数确定后，合并 radix 共享大页和请求自有页，仅查找本批页面所需槽位。完整大页必须存在，直接索引获取，缺失时暴露错误，不截断缓存前缀。将整批页面、完整 token 索引、大小页槽位和 `cpu_page_ready_flags` 一次传给资源管理器 `write_req_to_cpu`，由各 State 跳过 ready 页面。返回含完成事件的 `TransTask`，由已有任务完成流程发布 ready 并释放引用。 - -写出入口仅通过 `cpu_kv_cache_stream.wait_stream(producer)` 等待当前生产流已提交的写入，不额外等待 overlap 流，也不对 hybrid 模型执行主机同步。TODO：bounded CPU→CPU 快照搬运仍需改为当前流上的 kernel；当前主机 copy 不受流等待约束，kernel 落地前源快照必须预先就绪。 - -命中边界、容量准入和 CPU 页引用释放仍由外层缓存调度负责。以上入口不把 `InferReq` 传入 `AttentionStateManager`；资源层只接收显式索引与槽位，并组合各 State 的数据操作。 - -外部通过 6.2 节的管理器接口执行 CPU/DP/PD 操作。`AttentionStateManager` 将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 - -- full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 -- `KDAState`、`GDNState` 实现 `BoundedAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移 prefill 基准请求状态；管理器组织参与对象及位置映射。 -- 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttentionStateManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 - -跨进程导入的管理器持有借用的运行态存储和索引容量描述，不持有 allocator 或 CPU 状态池。请求映射由 server 层作为 ReqManager 所有的独立存储导出，不能附加到资源管理器。CUDA IPC 句柄按实际读取进程数发布，每个读取者消费独立引用；原进程在全部读取者结束使用前保持存储存活。 - -### 8.2 PD - -PD 任务分为 `kv` 和 `att_state` 两类；涉及 main/draft 时按模型角色组合所需资源： - -| 任务 | 组成 | 处理方式 | -| --- | --- | --- | -| `kv` | full-attention 历史 KV，包含布局所需的 packed indexer 数据 | 按 token 页复制 | -| `att_state` | Conv/SSM、SWA 和窗口压缩器运行态 | 模型专属 `AttentionStateManager` 按请求恢复位置调用对应 State 方法 | - -P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，索引池化窗口 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将窗口写入目标 TP 各 rank 的本地请求槽位。 - -模型专属 `AttentionStateManager` 实现状态页的布局组织和容量校验，调用 linear State 与 窗口压缩器 State 的直接传输方法。各具体 State 负责自身区域的读写和边界校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 - -网络层只接收固定容量的字节页，不按 KV 的 head 数或 dtype 推断组合布局。两端在创建传输句柄前核对页容量，PD master 在连接时校验任务类别、资源身份、算法与编码、层映射和全局字段形状；页内不附加格式指纹、长度或版本。布局使用全局字段描述以支持不同 TP 大小，窗口布局只描述固定逻辑窗口，不包含本地 MTP 候选容量。运行态的前缀长度由传输任务提供。 - -P/D 使用相同模型和 main/draft 资源集合，允许不同 TP 与 `mtp_step`。full-attention 传输有效 KV 和量化字段；bounded 状态按 6.7 节传输 prefill 基准 Conv/SSM，恢复到本地基准位置并将游标归零；窗口压缩器按 6.8 节传输固定逻辑窗口并按本地 ring 容量重排。MTP 候选容量不进入 PD 字段形状，单个 task 仍使用既有的前缀位置描述。 - -### 8.3 CPU/Disk 与 DP fetch - -CPU 联合页使用字节存储，从偏移 0 开始按具体 State 的字段布局保存数据，不预留格式页头或状态长度、版本字段。各资源区域由模型配置与 State 字段确定，TP 副本由明确的单个写入者更新共享区域。CPU 服务绑定时核对页大小，读回长度由任务参数提供，前缀与槽位的关联由调用方维护。磁盘缓存使用配置目录下的固定文件名 `cache_file`。 - -GPU→CPU 保存状态页完成后，才能在 CPU 上重排或复制这些快照；主机侧读取必须等待保存完成，不能只向另一个 CUDA stream 排入依赖。CPU 页的 ready 发布仍由调用流程在全部参与资源与 TP 分片传输完成后执行。 +```text +lightllm/common/attention_state/ +├── manager.py +├── types.py +├── utils.py +├── allocators/ +├── full_att/ +│ ├── base_state.py +│ ├── dp_transfer.py +│ ├── gqa/ +│ ├── mla/ +│ ├── packed_mla/ +│ └── compressed_att/ +└── bounded_att/ + ├── base_state.py + ├── copy_kernel.py + ├── transfer_utils.py + ├── gdn/ + ├── kda/ + ├── sliding_window/ + └── window_compressor/ +``` -缓存恢复由各 bounded State 根据绝对前缀长度处理：有重叠窗口恢复快照，无重叠窗口在对齐边界重置；无需额外的分类恢复接口。 +每个逻辑 State 目录包含其配置、State、创建入口和专用传输实现；跨算法共享的基类、布局类型、allocator 和通用 kernel 保留在类别根目录。 -CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大页 pool 分别提供存储。各参与方的暂存数据消费完成后，后续请求才能复用该槽位。磁盘存储使用 CPU 页格式；main/draft 异构组合需定义各自资源段的位置、字段形状及 TP 切片。 +代码编写约束： -DP fetch 同时恢复缓存前缀需要的历史 KV 和 bounded 运行态，并按各 State 的规则恢复窗口历史。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 +1. 公共接口放在私有辅助函数之前；类内依次放置初始化、公共属性/方法、私有辅助方法。 +2. 新增或大幅修改的函数，参数数量不少于 4 个时，声明和调用均使用逐行 `key=value` 格式；不为格式要求引入 `*args`、`**kwargs` 或转发字典。 +3. 配置字段必须有中文注释，说明含义、单位、Tensor 轴、全局/TP 范围和默认值语义。 +4. 公共接口 docstring 必须说明操作对象、索引含义、长度语义、数据范围和调用方责任。 +5. 具体 State 之间不继承实现；只继承对应类别的抽象基类，共享逻辑通过独立工具函数或明确的公共 kernel 提供。 +6. 不创建只有一次调用、只做参数转发或只包装断言的小函数；特殊算法行为放在对应 State 中。 +7. 已知输入约束在边界处使用简洁断言；不在公共接口中做隐式设备、dtype、索引或请求映射转换。 +8. 传输 kernel 的输入输出布局、stride、有效长度和循环映射必须在调用前明确；kernel 不负责外部资源分配和跨流同步。 +9. 修改现有基础流程时优先保持调用顺序和数据所有权不变；目录整理、命名调整和功能重构分开提交。 +10. 代码注释说明设计原因和调用契约，不重复逐行解释显而易见的 Python 语句，也不记录短期性能实验参数。 -### 8.1 DeepSeek V4：压缩 KV 与有界窗口 +## 8. 测试和评审要求 -V4 沿用共享 `AttentionStateManager`，不引入专用 ReqManager 或 MemoryManager。各模型通过 -`create_full_att_states` / `create_bounded_att_states` 创建自己的定义，注册完成后统一分配容量。 -main 与 draft0..draft35 使用相同 token 分配域，但分别持有存储；SWA prefill workspace 也按模型角色区分。 -P/D 的 decode 节点仍需在初始化时对 HOLD 请求执行一次 prefill，因此保留每层一个 SWA HOLD 页； -其普通 prefill workspace 容量为 0，不分配长输入临时区，实际请求的前缀由 prefill 节点传入。 +新增或修改 State 至少验证： -模型组合配置保留在 `models/deepseek_v4/config.py`，不放入原子 State 目录。 -组合配置的 `create_swa_config()` / `create_compressed_config()` 负责选择参与层及算法参数， -State 创建、backend 和提前预留 workspace 共用这些配置入口。编码字节数、packed 页布局及 -`get_attention_view()` 均由对应原子配置实现，backend 不再维护一套物理格式计算。 -SWA 使用独立的 `FP8E8M0SlidingWindowConfig`，只记录逻辑层、近期窗口、KV 编码与 TP 信息。 -压缩器使用独立的 `WindowCompressorStateConfig`，明确记录 `window_size`、`compress_stride`、 -`candidate_dim`、`target`、`dtype` 和 TP 信息；没有 C4/C128 层列表或 SWA 字段。 -`window_size` 是每次聚合覆盖的原始 token 数，`compress_stride` 是每隔多少原始 token 产出一条， -大小页需保留的跨步历史由 `overlap_tokens = window_size - compress_stride` 推导，不额外存储重复配置。 +- 配置校验、逻辑层归属和共享层映射。 +- 存储容量、页池启用/禁用和 allocator 生命周期。 +- 运行态重置、保存、恢复及绝对状态长度语义。 +- CPU、DP、PD 的布局、索引方向和量化/编码字段完整性。 +- main/draft、多 TP、跨设备或 IPC 场景下的资源所有权。 +- 异步流依赖和失败路径，不以主机同步掩盖错误。 -| V4 创建的运行态实例 | 窗口 / 产出步长 | 每 token 候选宽度 | 大小页快照窗口 | PD 窗口 | -| --- | --- | --- | --- | --- | -| SWA | 配置的 sliding window / 每 token 更新 | packed KV 字节 | 完整 sliding window | 完整 sliding window | -| CSA KV compressor | W8 / S4 | `2 * head_dim`，gate 另占同宽 | 上一组 4 个位置 | 8 个位置 | -| CSA indexer compressor | W8 / S4 | `2 * index_head_dim`，gate 另占同宽 | 上一组 4 个位置 | 8 个位置 | -| HCA KV compressor | W128 / S128 | `head_dim`，gate 另占同宽 | 组边界无跨步历史，不占页池和 CPU 页 | 128 个位置 | +目录或导入路径调整必须同时运行： -规范名称包含用途、窗口、步长及精度，例如 `main_compressor_kv_w8_s4_fp32_state`、 -`main_compressor_indexer_w8_s4_fp32_state`、`main_compressor_kv_w128_s128_fp32_state`， -SWA 为 `main_swa_w128_fp8_group64_e8m0_state`。同一个压缩器实现也可以保存 W12/S4 等其他窗口， -压缩计算仍由对应模型算子负责，不在 State 中根据窗口或比例猜测算法。 +- `compileall` 和包导入检查。 +- 受影响 State 的单元测试。 +- 相关模型创建和资源布局测试。 +- `git diff --check`。 -各 State 分别拥有自己的大小页池，共享同一组 slot ID 和 allocator；不保留 V4 聚合 State 或联合 pool。 -无快照的压缩器返回空 CPU regions，manager 不为其分配联合页区域，但保存/恢复仍校验已完成产出步长。 -缺少其他参与 State 的 CPU layout 仍会显式失败。标准 V4 MTP draft 只创建 SWA State。 +评审时优先检查：接口契约是否与基类一致、State 是否越权访问请求流程、布局是否可逆、索引是否可能重复归还、同步是否由正确的生产/消费流建立，以及量化变体是否仍归属于正确的逻辑算法目录。 -| 组件 | 归属 | 数据与索引 | -| --- | --- | --- | -| `CompressedAttentionState` | full-attention | 分别保存 C4/C128 层的压缩 KV；C4 同块保存 indexer key/scale，外部索引仍以原始 token 为单位 | -| `FP8E8M0SlidingWindowState` | bounded-attention | 只保存近期原始 KV 的 SWA ring；不依赖压缩器或 V4 模型配置 | -| `WindowCompressorState` | bounded-attention | 每实例只保存一种窗口/步长/用途的候选与 gate；不保存压缩历史或 SWA | -| `WindowPagesPool` | bounded 大小页池 | 每个参与 State 各自保存单一窗口快照；共享 manager 的槽位分配域，不维护请求或 radix 引用 | -| `ReqManager` | 请求管理 | 只保存请求槽位和逻辑 token 到物理 token 的映射 | +## 9. 非目标 -压缩历史通过独立的 `create_compressed_att_state()` 创建，不经过 MLA 创建入口。 -`CompressedAttentionConfig` 是独立的原子布局定义，不继承其他 full-attention 配置，也不持有 -`DSV4AttentionConfig`。它直接保存逻辑层分组 `logical_layer_groups`、压缩比例 `compress_ratio`、 -KV/RoPE 维度 `kv_dim` / `rope_dim`、每块覆盖的原始 token 数 `block_token_size`、可选 -`indexer_dim` 和当前 TP 信息。`indexer_dim=0` 表示不保存 indexer，不根据压缩比例推断。 -V4 模型负责将自己的组合配置转换为这些字段；State 不读取 SWA 窗口、全模型层列表或压缩器配置。 -当前实现固定使用一个完整复制的共享 KV head，使用 uint8 保存 FP8/BF16/scale 的原始编码， -不再要求调用方传入 `global_head_num=1`、`replicated=True`、`dtype=torch.uint8` 等固定参数。 -配置的 `token_slot_bytes` 按压缩比例推导每个原始 token 槽位摊销的物理字节数，包含 indexer 和 padding， -不借用 `head_dim` 表示字节数。State 的名称为 `compressed_c{压缩比例}_fp8_group64_e8m0`， -例如 `main_compressed_c4_fp8_group64_e8m0_state`。CSA/HCA 描述计算与历史选择流程， -不作为存储类的固定身份；同一实现按原子字段分别承载它们的压缩历史。 -它保存已完成的序列压缩结果，不保存 SWA 或 compressor 运行态；State 仍直接继承 `FullAttentionState`， -CPU/DP/PD 继续复用适用于该 packed 字节布局的索引传输 kernel，不改变物理格式。 +本文不规定： -压缩 KV 以 256 个原始 token 为一个物理块，块内数据与 scale 分区存放；C4 还包含 indexer。 -PD 的最后一个碎页必须携带整个已分配压缩块，避免遗漏有效 token 对应的 scale/indexer。 -CPU 页面、KV 分配页和前缀边界须遵守相应的压缩块对齐约束。 +- 具体模型的 hidden size、head 数、压缩比例、窗口大小或页大小。 +- 某个 GPU、通信库或 Triton 版本的调优参数。 +- radix cache 的匹配策略、淘汰策略和外部磁盘格式。 +- 某个模型的 attention 数学公式或 kernel 内部实现。 +- 历史兼容路径和临时实验代码。 -大小页保存/恢复接收绝对 `state_token_lens`，由 State 将请求 ring 与时间有序快照相互转换。 -CPU 读回仍先写入大页池，再恢复请求运行态；最后一个碎页使用管理器的 LOAD 暂存槽位。 -C128 在缓存边界已经关闭当前压缩组，无重叠历史需要保存；V4 的状态页边界因此必须是 128 的倍数。 -C4 保存上一组的四个位置，供下一组重叠压缩使用。运行态与大小页之间直接调用 V4 kernel， -不使用外部 `TransferRegion`；仅 CPU/PD 联合页使用 region 描述字段与 TP 写入权限。 - -PD 运行态采用与本地 MTP 容量无关的逻辑窗口：SWA 为配置窗口，C4/indexer 为八个位置， -C128 为 128 个位置。收发两端根据传入的绝对前缀长度各自计算 ring 位置，因此允许不同的 MTP 配置。 -SWA、压缩 KV 和 compressor 在 TP 内完整复制，仅 rank 0 写出共享外部区域，各 rank 可读回。 -窗口是否需要真实快照由重叠历史决定，GLM K-pool 和 V4 压缩器使用相同的 bounded State 接口。 - -checkpoint 中的权重层号与 draft 自有逻辑层数分别记录；多个复用相同 MTP 权重的 draft 仍注册独立 State。 -模型和 backend 根据逻辑层查询 State 视图，不在 ReqManager 中附加 ring，不在 backend 中保存运行态副本。 - -## 9. 待细化的设计问题 - -| 问题 | 设计约定 | 待确定内容 | -| --- | --- | --- | -| main/draft 管理器装配 | 模型组合实现自己的管理器；各角色和具体方案拥有唯一 State，同一分配域共享唯一 allocator 实例 | 新模型的 main/draft 装配时序及独立存储分配；调用方覆盖全部参与资源的引用解除及单次归还责任 | -| AttentionStateManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | -| 内部 attention 资源的暴露 | 具体模型在初始化阶段绑定 State 或层视图，不要求统一模型视图获取接口 | 各模型/backend 的具体绑定方式与按层访问结构 | -| FullAttentionState 接口与继承 | 必需语义：容量计算、初始化、层视图、KV 写入、完整 token 复制及固定字段传输布局；CPU/DP/PD 由子类按支持路径直接实现读写方法 | 新布局的 TensorField 覆盖范围、StateTransferLayout 映射约束及具体搬运实现 | -| BoundedAttentionState 接口与继承 | KDAState、GDNState 实现统一抽象操作；已接受运行态、checkpoint 和传输分别定义，allocator 仍归管理器 | 按 6.7 节细化各子类的描述字段和 kernel，核对 MTP 状态选择、跨 TP/CPU 格式兼容性 | -| allocator 能力描述 | 公共入口按需初始化，不需要的 allocator 为 None | 恢复边界和传输能力使用哪些最小描述，避免重复配置导致不一致 | -| checkpoint 引用形式 | 调用方维护大小页类别、共享槽位 ID、前缀长度及参与模型资源；各侧 pool 只保存快照数据 | 引用对象或参数组合的具体类型，以及参与资源集合的表达 | -| 配置类型拆分 | main/draft 分别产生实现和布局参数，共享 allocator 接收统一容量与索引配置 | 各类 State 的配置字段、跨资源共享字段与校验责任 | -| PD 窗口压缩器的布局 | 固定 P 个逻辑位置，按任务前缀长度映射两端 ring；兼容不同 MTP 步数 | 保持相同模型资源集合与 prefill 状态交接约束 | -| CPU 缓存页的组合方式 | 每页包含全部有效 linear 状态；管理器组织读写，并在全部页面恢复后从末页恢复请求运行态 | 不同 main/draft 实现的资源段布局及恢复边界兼容 | - -以上待细化事项须遵守本文的资源归属、共享分配域和恢复语义约定。 +这些内容应在对应模型、State、kernel 或测试文件中维护，并通过公共接口满足本文契约。 From 597080c08afc3f5c03c5583168a40f10a6a3b999 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 10:32:07 +0000 Subject: [PATCH 39/56] fix packed MLA test state helpers --- unit_tests/common/att_resources/model_helpers.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index ec2cf24cb0..5adedc5408 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -129,7 +129,16 @@ def create_full_attention_roles( if start == end: states.append(None) continue - config = replace(template.config, logical_layer_groups=layers[start:end]) + if isinstance(template.config, PackedMLAAttentionConfig): + config = replace( + template.config, + logical_config=replace( + template.config.logical_config, + logical_layer_groups=layers[start:end], + ), + ) + else: + config = replace(template.config, logical_layer_groups=layers[start:end]) if isinstance(template, StaticFP8FullAttentionState): state = rebuild_full_state( template=template, From 52e0d6bf11a4d52a974817dac0732043eb2d5764 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 10:50:52 +0000 Subject: [PATCH 40/56] avoid dynamic SWA kernel specializations --- .../models/deepseek_v4/triton_kernel/swa_cache.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py index 9742e7cae8..0f9365b1b1 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py +++ b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py @@ -53,22 +53,24 @@ def _prepare_swa_prefill_workspace_kernel( b_req_idx, b_ready_cache_len, b_seq_len, - BATCH: tl.constexpr, + batch_size, SLIDING_WINDOW: tl.constexpr, SWA_PAGE_TOKEN_NUM: tl.constexpr, HOLD_REQUEST_ID: tl.constexpr, - SWA_WORKSPACE_BASE: tl.constexpr, + swa_workspace_base, BLOCK: tl.constexpr, ): batch = tl.arange(0, BLOCK) - valid = batch < BATCH + valid = batch < batch_size req_idx = tl.load(b_req_idx + batch, valid, HOLD_REQUEST_ID) ready = tl.load(b_ready_cache_len + batch, valid, 0) end = tl.load(b_seq_len + batch, valid, 0) first = tl.maximum(0, ready - SLIDING_WINDOW + 1) // SWA_PAGE_TOKEN_NUM * SWA_PAGE_TOKEN_NUM length = tl.cdiv(end - first, SWA_PAGE_TOKEN_NUM) * SWA_PAGE_TOKEN_NUM tokens = tl.where((req_idx != HOLD_REQUEST_ID) & valid, length, 0) - workspace = SWA_WORKSPACE_BASE + tl.cumsum(tokens) - tokens + # 请求数和 workspace 起点随每个 prefill microbatch 变化，保持为运行时参数， + # 避免按 batch 大小或 microbatch 索引重复编译同一份布局 kernel。 + workspace = swa_workspace_base + tl.cumsum(tokens) - tokens tl.store(prefill_token_ranges + req_idx * 2, ready, valid) tl.store(prefill_token_ranges + req_idx * 2 + 1, end, valid) tl.store(swa_prefill_workspace_offsets + req_idx, workspace, valid) @@ -119,11 +121,11 @@ def prepare_swa_prefill_workspace(swa_state, prefill_state, b_req_idx, b_ready_c b_req_idx, b_ready_cache_len, b_seq_len, - BATCH=b_req_idx.numel(), + batch_size=b_req_idx.numel(), SLIDING_WINDOW=config.window_size, SWA_PAGE_TOKEN_NUM=config.page_token_size, HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, - SWA_WORKSPACE_BASE=workspace_base, + swa_workspace_base=workspace_base, BLOCK=triton.next_power_of_2(b_req_idx.numel()), num_warps=4, ) From 4bf0441fa19578f502805d3ce0a8cb661daf7286 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 10:52:54 +0000 Subject: [PATCH 41/56] avoid dynamic DSpark SWA specializations --- .../triton_kernel/build_dspark_swa_indices.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py index 5e0c6dc492..6a48b024cd 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py +++ b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py @@ -15,7 +15,7 @@ def _build_dspark_swa_indices_kernel( HOLD_REQUEST_ID: tl.constexpr, HOLD_SWA_SLOT: tl.constexpr, SLIDING_WINDOW: tl.constexpr, - BLOCK_SIZE: tl.constexpr, + block_size, WIDTH: tl.constexpr, BLOCK_W: tl.constexpr, ): @@ -24,14 +24,16 @@ def _build_dspark_swa_indices_kernel( position = tl.load(position_ids + token_idx).to(tl.int64) is_hold = req_idx == HOLD_REQUEST_ID - block_offset = token_idx % BLOCK_SIZE + # draft proposal 长度随请求变化，不参与编译期分支；保持运行时参数避免 + # 每种 MTP 接受长度都生成一份新的 kernel。 + block_offset = token_idx % block_size block_start = tl.maximum(position - block_offset, 0) history_len = tl.minimum(block_start, SLIDING_WINDOW) column = tl.arange(0, BLOCK_W) in_width = column < WIDTH is_history = column < history_len - is_block = (column >= history_len) & (column < history_len + BLOCK_SIZE) + is_block = (column >= history_len) & (column < history_len + block_size) history_position = block_start - 1 - column block_position = block_start + column - history_len source_position = tl.where(is_history, history_position, block_position) @@ -43,7 +45,7 @@ def _build_dspark_swa_indices_kernel( output = tl.where(is_hold, HOLD_SWA_SLOT, output).to(tl.int32) tl.store(swa_indices + token_idx * swa_indices_stride0 + column, output, mask=in_width) - length = tl.where(is_hold, 1, history_len + BLOCK_SIZE).to(tl.int32) + length = tl.where(is_hold, 1, history_len + block_size).to(tl.int32) tl.store(swa_lengths + token_idx, length) write_slot = req_idx * SWA_CAPACITY_TOKENS + position % SWA_CAPACITY_TOKENS write_slot = tl.where(is_hold, HOLD_SWA_SLOT, write_slot).to(tl.int32) @@ -77,7 +79,7 @@ def build_dspark_swa_indices( HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, HOLD_SWA_SLOT=swa_state.HOLD_SWA_SLOT, SLIDING_WINDOW=swa_state.config.window_size, - BLOCK_SIZE=block_size, + block_size=block_size, WIDTH=width, BLOCK_W=triton.next_power_of_2(width), num_warps=4, From f8067823a704239ec3b929d96ed3f24dd154fea0 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 12:01:16 +0000 Subject: [PATCH 42/56] bind startup barriers to current CUDA device --- lightllm/common/basemodel/basemodel.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e0c439cd8a..1f478a71e8 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1052,7 +1052,8 @@ def _check_max_len_infer(self): return # 做一次 同步 - torch.distributed.barrier() + # DP/TP 混合启动时显式指定当前卡，避免 NCCL 根据全局 rank 猜错 barrier 设备。 + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) # 模拟最大长度进行 prefill，观察是否出现 OOM try: @@ -1119,7 +1120,7 @@ def autotune_layers(self): @post_empty_cache def _autotune_warmup(self): Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096] @@ -1187,7 +1188,7 @@ def _autotune_warmup(self): gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) Autotuner.end_autotune_warmup() @final @@ -1197,7 +1198,7 @@ def _init_padded_req(self): 对 padded 所使用的req 进行初始化， 目前有非常多的地方需要使用，所以将其初始化固定为固定流程 """ # 做一次 同步 - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) # prefill init padding req. prefill_input_len = 1 From 5c3d7a9d876edc07eab0e7fcc814a0fba04bbd07 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 12:29:16 +0000 Subject: [PATCH 43/56] Revert "bind startup barriers to current CUDA device" This reverts commit f8067823a704239ec3b929d96ed3f24dd154fea0. --- lightllm/common/basemodel/basemodel.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 1f478a71e8..e0c439cd8a 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1052,8 +1052,7 @@ def _check_max_len_infer(self): return # 做一次 同步 - # DP/TP 混合启动时显式指定当前卡，避免 NCCL 根据全局 rank 猜错 barrier 设备。 - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() # 模拟最大长度进行 prefill，观察是否出现 OOM try: @@ -1120,7 +1119,7 @@ def autotune_layers(self): @post_empty_cache def _autotune_warmup(self): Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096] @@ -1188,7 +1187,7 @@ def _autotune_warmup(self): gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() Autotuner.end_autotune_warmup() @final @@ -1198,7 +1197,7 @@ def _init_padded_req(self): 对 padded 所使用的req 进行初始化， 目前有非常多的地方需要使用，所以将其初始化固定为固定流程 """ # 做一次 同步 - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() # prefill init padding req. prefill_input_len = 1 From 98ad42072eebb0669823e07819628e22fede9166 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 12:33:37 +0000 Subject: [PATCH 44/56] Revert "Revert "bind startup barriers to current CUDA device"" This reverts commit 5c3d7a9d876edc07eab0e7fcc814a0fba04bbd07. --- lightllm/common/basemodel/basemodel.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e0c439cd8a..1f478a71e8 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1052,7 +1052,8 @@ def _check_max_len_infer(self): return # 做一次 同步 - torch.distributed.barrier() + # DP/TP 混合启动时显式指定当前卡，避免 NCCL 根据全局 rank 猜错 barrier 设备。 + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) # 模拟最大长度进行 prefill，观察是否出现 OOM try: @@ -1119,7 +1120,7 @@ def autotune_layers(self): @post_empty_cache def _autotune_warmup(self): Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096] @@ -1187,7 +1188,7 @@ def _autotune_warmup(self): gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) Autotuner.end_autotune_warmup() @final @@ -1197,7 +1198,7 @@ def _init_padded_req(self): 对 padded 所使用的req 进行初始化， 目前有非常多的地方需要使用，所以将其初始化固定为固定流程 """ # 做一次 同步 - torch.distributed.barrier() + torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) # prefill init padding req. prefill_input_len = 1 From f25a8ec8e3401d865c71b78bd3314b75d651e15f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 12:48:40 +0000 Subject: [PATCH 45/56] Revert "Revert "Revert "bind startup barriers to current CUDA device""" This reverts commit 98ad42072eebb0669823e07819628e22fede9166. --- lightllm/common/basemodel/basemodel.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 1f478a71e8..e0c439cd8a 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1052,8 +1052,7 @@ def _check_max_len_infer(self): return # 做一次 同步 - # DP/TP 混合启动时显式指定当前卡，避免 NCCL 根据全局 rank 猜错 barrier 设备。 - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() # 模拟最大长度进行 prefill，观察是否出现 OOM try: @@ -1120,7 +1119,7 @@ def autotune_layers(self): @post_empty_cache def _autotune_warmup(self): Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() warmup_lengths = [1, 4, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096] @@ -1188,7 +1187,7 @@ def _autotune_warmup(self): gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() Autotuner.end_autotune_warmup() @final @@ -1198,7 +1197,7 @@ def _init_padded_req(self): 对 padded 所使用的req 进行初始化， 目前有非常多的地方需要使用，所以将其初始化固定为固定流程 """ # 做一次 同步 - torch.distributed.barrier(device_ids=[torch.cuda.current_device()]) + torch.distributed.barrier() # prefill init padding req. prefill_input_len = 1 From 5d323c6e9150d53cef758c38cf8f159794f850b4 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sun, 1 Nov 2026 14:46:34 +0000 Subject: [PATCH 46/56] fix flashinfer allreduce subgroup rendezvous --- lightllm/distributed/flashinfer_all_reduce.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/lightllm/distributed/flashinfer_all_reduce.py b/lightllm/distributed/flashinfer_all_reduce.py index 35cb0aaedf..1029795fd5 100644 --- a/lightllm/distributed/flashinfer_all_reduce.py +++ b/lightllm/distributed/flashinfer_all_reduce.py @@ -102,6 +102,9 @@ def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool: hidden_dim=hidden_dim, dtype=dtype, comm_backend=TorchDistBackend(group=self.group), + # TRT-LLM symmetric-memory rendezvous must use the same TP + # subgroup as world_size/rank; otherwise it defaults to WORLD. + group=self.group, ) except Exception as e: logger.warning("FlashInferAllReduce workspace init failed: %s. Disabling.", e) From d998618cdca64ae794838f7cd20af920cb13c29a Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 2 Nov 2026 00:36:17 +0000 Subject: [PATCH 47/56] refactor attention transfer operators --- .../attention_state/cpu_cache_operator.py | 183 +++++++ .../attention_state/dp_shared_operator.py | 76 +++ lightllm/common/attention_state/manager.py | 474 ++---------------- .../common/attention_state/pd_operator.py | 170 +++++++ lightllm/common/attention_state/types.py | 2 +- .../model_infer/mode_backend/base_backend.py | 2 +- .../dp_backend/dp_shared_kv_trans.py | 2 +- .../mode_backend/multi_level_kv_cache.py | 12 +- .../mode_backend/pd/resource_io.py | 18 +- .../model_infer/shared_att_resources_utils.py | 4 +- lightllm/utils/cpu_cache_layout.py | 2 +- test/att_resources/validate_model_cache.py | 28 +- test/cpu_cache_kernel/test_speed.py | 12 +- .../common/att_resources/manager_fixtures.py | 13 - .../common/att_resources/model_helpers.py | 13 + .../att_resources/test_backend_startup.py | 6 +- .../common/att_resources/test_capacity.py | 131 ++--- .../att_resources/test_compute_binding.py | 3 +- .../common/att_resources/test_dsv4_state.py | 14 +- .../att_resources/test_fp8_compute_binding.py | 3 +- .../test_full_storage_initialization.py | 4 +- .../common/att_resources/test_layer_groups.py | 29 +- .../test_layer_state_selection.py | 5 +- .../att_resources/test_manager_helpers.py | 12 +- .../test_model_initialization.py | 6 +- .../att_resources/test_model_resources.py | 12 +- .../test_model_storage_definition.py | 10 +- .../att_resources/test_role_transfers.py | 26 +- .../att_resources/test_runtime_transport.py | 26 +- .../att_resources/test_state_identity.py | 21 +- .../att_resources/test_state_queries.py | 10 +- .../common/att_resources/test_tp_config.py | 13 +- .../test_window_compressor_state.py | 4 +- .../attention/test_fp8_logical_layer.py | 3 +- unit_tests/common/dsv4_test_utils.py | 24 +- .../common/test_deepseek4_paged_cache.py | 12 +- .../deepseek_v4/test_vision_integration.py | 6 +- unit_tests/models/glm5_next/test_fp8_kv.py | 10 +- unit_tests/models/glm5_next/test_pd_cache.py | 8 +- .../test_multimodal_cache_matching.py | 8 +- .../mode_backend/test_dp_big_pages.py | 2 +- .../mode_backend/test_multi_level_kv_cache.py | 8 +- .../test_paged_side_allocations.py | 4 +- .../mode_backend/test_pd_token_indices.py | 2 +- .../model_infer/test_cpu_resource_layout.py | 41 +- .../test_hybrid_cpu_state_cache.py | 4 +- .../model_infer/test_multimodal_chunking.py | 10 +- 47 files changed, 777 insertions(+), 711 deletions(-) create mode 100644 lightllm/common/attention_state/cpu_cache_operator.py create mode 100644 lightllm/common/attention_state/dp_shared_operator.py create mode 100644 lightllm/common/attention_state/pd_operator.py diff --git a/lightllm/common/attention_state/cpu_cache_operator.py b/lightllm/common/attention_state/cpu_cache_operator.py new file mode 100644 index 0000000000..9f07ec0356 --- /dev/null +++ b/lightllm/common/attention_state/cpu_cache_operator.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import replace +from functools import lru_cache +from typing import TYPE_CHECKING, List + +import torch + +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +class CpuCacheOperator: + """实现 CPU cache 的联合布局和 State 搬运。 + + AttentionStateManager 通过该 Operator 提供 CPU 页操作；这里集中处理 CPU 页如何组织， + 以及 full-attention token 与 bounded 运行态快照如何写回/读回。 + """ + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + @lru_cache(maxsize=None) + def get_page_layouts(self) -> Mapping[str, StateTransferLayout]: + page_tokens = get_env_start_args().cpu_cache_token_page_size + assert page_tokens > 0, "CPU pages require a positive token capacity" + states = (*self.owner.get_full_attention_states(), *self.owner.get_bounded_attention_states()) + assert states, "a CPU page requires participating storage" + + layouts, offset = {}, 0 + for state in states: + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_cpu_transfer_regions() + ) + # 没有跨步历史的 State 不占 CPU 页；其余 State 仍按联合页顺序排列。 + if not regions: + continue + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_page_bytes(self) -> int: + layouts = self.get_page_layouts() + page_bytes = max((layout.byte_end for layout in layouts.values()), default=0) + return (page_bytes + 15) // 16 * 16 + + def clear_layout_cache(self) -> None: + self.get_page_layouts.cache_clear() + self.get_page_bytes.cache_clear() + + def write_req( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layouts: Mapping[str, StateTransferLayout], + cpu_page_ready_flags: List[bool] | None = None, + big_page_slot_ids: List[int] | None = None, + small_page_slot_ids: List[int] | None = None, + ) -> None: + """写出单个请求的 full-attention token 和 bounded 页快照。""" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + if all(cpu_page_ready_flags): + return + + full_states = self.owner.get_full_attention_states() + bounded_states = self.owner.get_bounded_attention_states() + if bounded_states: + big_page_slot_ids = [] if big_page_slot_ids is None else big_page_slot_ids + small_page_slot_ids = [] if small_page_slot_ids is None else small_page_slot_ids + big_page_count = len(big_page_slot_ids) + assert big_page_count + len(small_page_slot_ids) == len(cpu_page_indices), ( + "one state slot is required per CPU page" + ) + big_cpu_page_indices = cpu_page_indices[:big_page_count] + small_cpu_page_indices = cpu_page_indices[big_page_count:] + big_page_readies = ( + cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None + ) + small_page_readies = ( + cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None + ) + + for state in full_states: + state.write_to_cpu( + token_indices=token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + layout=layouts[state.state_name], + cpu_page_ready_flags=cpu_page_ready_flags, + ) + + for state in bounded_states: + layout = layouts.get(state.state_name) + if layout is None: + # 某些压缩窗口在合法缓存边界不需要快照，因此可以不进入 CPU cache。 + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue + if big_page_slot_ids: + state.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=big_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=big_cpu_page_indices, + layout=layout, + cpu_page_ready_flags=big_page_readies, + ) + if small_page_slot_ids: + state.write_state_pages_to_cpu( + page_kind="small", + page_slot_ids=small_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=small_cpu_page_indices, + layout=layout, + cpu_page_ready_flags=small_page_readies, + ) + + def read_req( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layouts: Mapping[str, StateTransferLayout], + big_page_slot_ids: List[int] | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + """读回请求页面，并从末页恢复 bounded 运行态。""" + full_states = self.owner.get_full_attention_states() + bounded_states = self.owner.get_bounded_attention_states() + if bounded_states: + assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" + assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" + page_tokens = get_env_start_args().cpu_cache_token_page_size + has_partial_page = token_indices.numel() % page_tokens != 0 + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), ( + "slots must cover complete CPU pages only" + ) + load_big_page_slot_ids = big_page_slot_ids.copy() + if has_partial_page: + load_big_page_slot_ids.append(self.owner.bounded_att_cpu_load_big_page_slot) + + for state in full_states: + state.read_from_cpu( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + token_indices=token_indices, + layout=layouts[state.state_name], + ) + + for state in bounded_states: + if state.state_name not in layouts: + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue + state.read_state_pages_from_cpu( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + big_page_slot_ids=load_big_page_slot_ids, + layout=layouts[state.state_name], + ) + + if bounded_states: + self.owner.restore_bounded_att_runtime_from_pages( + req_indices=[req_idx], + state_token_lens=[state_token_len], + page_kind="big", + page_slot_ids=[load_big_page_slot_ids[-1]], + ) diff --git a/lightllm/common/attention_state/dp_shared_operator.py b/lightllm/common/attention_state/dp_shared_operator.py new file mode 100644 index 0000000000..8aa90330cd --- /dev/null +++ b/lightllm/common/attention_state/dp_shared_operator.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from copy import copy +from typing import TYPE_CHECKING, List + +import torch + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +class DPSharedOperator: + """封装 DP shared 的跨副本复制、IPC 导出和共享页注册。""" + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + def create_ipc_copy(self) -> AttentionStateManager: + """保留共享运行态和共享 CPU 页池，排除本地 allocator。""" + for state in self.owner.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + for tensor in vars(pool).values(): + if isinstance(tensor, torch.Tensor): + assert tensor.device.type == "cpu" and tensor.is_shared(), ( + "DP IPC requires shared CPU page pools" + ) + ipc_copy = copy(self.owner) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + ipc_copy._reset_operators() + return ipc_copy + + def register_page_pools_pin_memory(self) -> None: + """在当前 GPU 映射所有共享 bounded 页池。""" + for state in self.owner.get_bounded_attention_states(): + if state.big_pages_pool is not None: + state.big_pages_pool.register_pin_memory() + if state.small_pages_pool is not None: + state.small_pages_pool.register_pin_memory() + + def copy_from_dp( + self, + remote_att_manager: AttentionStateManager, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + src_big_page_ids: List[int], + dst_big_page_ids: List[int], + src_req_idx: int, + dst_req_idx: int, + ) -> None: + """复制其他 DP 的 token、大小页和 prefill 基准运行态。""" + assert type(remote_att_manager) is type(self.owner), "DP replicas must use the same model resource manager" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + assert len(src_big_page_ids) == len(dst_big_page_ids), "source and destination page counts differ" + if src_token_indices.numel(): + for state in self.owner.get_full_attention_states(): + state.copy_tokens_from_dp( + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + remote_dp_full_att_state=getattr(remote_att_manager, state.state_name), + ) + if src_big_page_ids: + for state in self.owner.get_bounded_attention_states(): + state.copy_big_pages_from_dp( + src_page_ids=src_big_page_ids, + dst_page_ids=dst_big_page_ids, + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), + ) + for state in self.owner.get_bounded_attention_states(): + state.copy_runtime_state_from_dp( + src_req_indices=[src_req_idx], + dst_req_indices=[dst_req_idx], + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), + ) diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index ffd539d67e..ff55d7c34d 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -1,26 +1,23 @@ from __future__ import annotations from typing import TYPE_CHECKING -from collections.abc import Mapping -from copy import copy -from dataclasses import replace from functools import lru_cache -from typing import List, Literal +from typing import List import torch -from lightllm.common.attention_state.types import ( - StateTransferLayout, - PageKind, -) +from lightllm.common.attention_state.types import PageKind from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from .allocators import StatePageAllocator, FullAttKVAllocator +from .cpu_cache_operator import CpuCacheOperator +from .dp_shared_operator import DPSharedOperator from .full_att import FullAttentionState from .bounded_att import BoundedAttentionState +from .pd_operator import PDOperator from .utils import MODEL_ROLES if TYPE_CHECKING: @@ -36,6 +33,10 @@ class AttentionStateManager: 通过聚合查询操作各类 State；模型通过两类创建接口决定资源组合，main/draft 共享本管理器。 未支持的传输路径必须明确报错。 + + CPU cache、DP shared 和 PD 的实现分别由三个 Operator 持有；管理器只保留 + State 注册、查询和共享容量规划。Operator 只依赖本管理器提供的 State 查询 + 和 allocator，不复制 State 的算法逻辑。 """ # ======================================================================== @@ -54,6 +55,27 @@ def __init__(self, request_capacity: int): self.bounded_att_big_page_allocator = None self.bounded_att_small_page_allocator = None self._request_capacity = request_capacity + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + + def _reset_operators(self) -> None: + """让浅拷贝后的 Operator 指向新的资源管理器副本。""" + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + + def get_cpu_cache_operator(self) -> CpuCacheOperator: + """返回当前资源管理器唯一的 CPU cache 操作对象。""" + return self._cpu_cache_operator + + def get_dp_shared_operator(self) -> DPSharedOperator: + """返回当前资源管理器唯一的 DP shared 操作对象。""" + return self._dp_shared_operator + + def get_pd_operator(self) -> PDOperator: + """返回当前资源管理器唯一的 PD 操作对象。""" + return self._pd_operator def register_state( self, model_role: str, state: FullAttentionState | BoundedAttentionState @@ -176,7 +198,7 @@ def initialize_storage( assert self.get_full_attention_states(), "the shared token index domain requires a full-attention State" pd_reserved_bytes = 0 if args.run_mode in ("prefill", "decode"): - pd_reserved_bytes = self.get_pd_page_bytes() * args.pd_kv_page_num + pd_reserved_bytes = self.get_pd_operator().get_page_bytes() * args.pd_kv_page_num if token_capacity is None: # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。 torch.cuda.empty_cache() @@ -269,53 +291,9 @@ def allocate_storage( for state in (*full_states, *bounded_states): state.get_cpu_transfer_regions.cache_clear() state.get_pd_transfer_regions.cache_clear() - self.get_cpu_page_layouts.cache_clear() - self.get_cpu_page_bytes.cache_clear() - self.get_pd_token_page_layouts.cache_clear() - self.get_pd_runtime_page_layouts.cache_clear() - self.get_pd_page_bytes.cache_clear() - - def create_pd_ipc_copy(self) -> AttentionStateManager: - """创建 PD 导出副本，保留运行态 Tensor，排除 allocator 和全部大小页池。""" - ipc_copy = copy(self) - ipc_copy.full_att_allocator = None - ipc_copy.bounded_att_big_page_allocator = None - ipc_copy.bounded_att_small_page_allocator = None - - # PD 只传运行态；先复制 State 对象再移除页池，避免修改推理进程的资源。 - for state in self.get_bounded_attention_states(): - state_copy = copy(state) - state_copy.big_pages_pool = None - state_copy.small_pages_pool = None - setattr(ipc_copy, state.state_name, state_copy) - return ipc_copy - - def create_dp_ipc_copy(self) -> AttentionStateManager: - """创建 DP shared 导出副本，保留运行态和共享大小页池，仅排除 allocator。""" - # DP fetch 初始化已选择共享 CPU 页池；禁止导出私有页池时隐式分配、复制整块存储。 - for state in self.get_bounded_attention_states(): - for pool in (state.big_pages_pool, state.small_pages_pool): - if pool is not None: - for tensor in vars(pool).values(): - if isinstance(tensor, torch.Tensor): - assert tensor.device.type == "cpu" and tensor.is_shared(), "DP IPC requires shared CPU page pools" - ipc_copy = copy(self) - ipc_copy.full_att_allocator = None - ipc_copy.bounded_att_big_page_allocator = None - ipc_copy.bounded_att_small_page_allocator = None - return ipc_copy - - def register_page_pools_pin_memory(self) -> None: - """在当前 GPU 上统一注册所有 bounded State 的大小页池，供 DP 接收端调用。 - - 调用方先选择本进程设备；注册状态不随 IPC 传递，重复调用会跳过已注册存储。 - 本接口只注册本地映射，不申请槽位或改变原进程的 allocator 所有权。 - """ - for state in self.get_bounded_attention_states(): - if state.big_pages_pool is not None: - state.big_pages_pool.register_pin_memory() - if state.small_pages_pool is not None: - state.small_pages_pool.register_pin_memory() + # Layout cache 归属于各 Operator；State 存储规格变化后统一失效。 + self.get_cpu_cache_operator().clear_layout_cache() + self.get_pd_operator().clear_layout_cache() def get_full_attention_states(self, role: str | None = None) -> tuple[FullAttentionState, ...]: """返回全部或指定角色的 full-attention State；只缓存对象引用。""" @@ -332,36 +310,6 @@ def get_full_att_token_bytes(self) -> int: total += state.get_token_bytes() return total - @lru_cache(maxsize=None) - def get_cpu_page_bytes(self) -> int: - """从 CPU 联合页布局的最大终点计算物理页大小，包含字段及 State 间的对齐空间。""" - page_bytes = 0 - for layout in self.get_cpu_page_layouts().values(): - page_bytes = max(page_bytes, layout.byte_end) - # State 区间按 8 字节对齐，整个 CPU 物理页再按 16 字节对齐。 - return (page_bytes + 15) // 16 * 16 - - @lru_cache(maxsize=None) - def get_pd_page_bytes(self) -> int: - """取 KV 页与完整运行态页占用的最大值，保证两类 PD 任务都能容纳。""" - full_layouts = self.get_pd_token_page_layouts() - runtime_layouts = self.get_pd_runtime_page_layouts() - full_att_page_bytes = 0 - for layout in full_layouts.values(): - full_att_page_bytes = max(full_att_page_bytes, layout.byte_end) - runtime_page_bytes = 0 - for layout in runtime_layouts.values(): - runtime_page_bytes = max(runtime_page_bytes, layout.byte_end) - - # KV 页不足时扩大物理传输页，不改变 KV 页的 token 数或各 State 的布局。 - if full_layouts and full_att_page_bytes < runtime_page_bytes: - logger.warning( - f"PD KV page size ({full_att_page_bytes} bytes) is smaller than the complete attention runtime " - f"({runtime_page_bytes} bytes); using {runtime_page_bytes} bytes per transfer page" - ) - # 纯 KV 或纯运行态取各自占用；没有任何参与资源时返回 0。 - return max(full_att_page_bytes, runtime_page_bytes) - # ======================================================================== # 二、Full attention：token 数据复制 # ======================================================================== @@ -445,361 +393,9 @@ def commit_bounded_att_state( state.commit_bounded_att_state(req_indices, accepted_steps) # ======================================================================== - # 四、CPU Cache 数据交换 - # ======================================================================== - - @lru_cache(maxsize=None) - def get_cpu_page_layouts(self) -> Mapping[str, StateTransferLayout]: - """聚合各 State 的 CPU 字段区间，返回按 state_name 索引的联合页布局。 - - 使用全局 cpu_cache_token_page_size，按需缓存；不依赖页大小查询。 - """ - page_tokens = get_env_start_args().cpu_cache_token_page_size - assert page_tokens > 0, "CPU pages require a positive token capacity" - states = (*self.get_full_attention_states(), *self.get_bounded_attention_states()) - assert states, "a CPU page requires participating storage" - - # 按 full-attention、bounded 顺序组织全部角色；各类 State 内部按资源名排序。 - # 无重叠窗口在缓存组边界无需快照，由 State 返回空 CPU 区域。 - layouts, offset = {}, 0 - for state in states: - regions = tuple( - replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) - for region in state.get_cpu_transfer_regions() - ) - # 无跨步重叠的压缩器在组边界不需要快照，不占用 CPU 页。 - if not regions: - continue - # 以实际字段终点确定 State 占用，保留内部 padding，并对齐下一个 State 的起点。 - state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.state_name] = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=page_tokens, - byte_start=offset, - byte_end=state_end, - regions=regions, - ) - offset = state_end - return layouts - - # CPU cache：按 layouts 组织联合页；混合模型的每页都包含全部参与 bounded State 的快照。 - # 纯 full-attention 模型只传 KV；具体实现负责校验格式、长度及各资源的区域映射。 - def write_req_to_cpu( - self, - token_indices: torch.Tensor, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layouts: Mapping[str, StateTransferLayout], - cpu_page_ready_flags: List[bool] | None = None, - big_page_slot_ids: List[int] | None = None, - small_page_slot_ids: List[int] | None = None, - ) -> None: - """将单个请求本次指定前缀页面的 KV 和 bounded 快照写入 CPU cache。 - - cpu_page_indices 前段对应 big_page_slot_ids，后段对应 small_page_slot_ids， - 各段按列表顺序一一对应；未参与的页池可传 None 或空列表。 - cpu_page_ready_flags 与 CPU 页一一对应，True 跳过整页；None 表示全部写出。 - """ - if cpu_page_ready_flags is not None: - assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" - if all(cpu_page_ready_flags): - return - full_states = self.get_full_attention_states() - bounded_states = self.get_bounded_attention_states() - if bounded_states: - big_page_slot_ids = [] if big_page_slot_ids is None else big_page_slot_ids - small_page_slot_ids = [] if small_page_slot_ids is None else small_page_slot_ids - big_page_count = len(big_page_slot_ids) - assert big_page_count + len(small_page_slot_ids) == len(cpu_page_indices), "one state slot is required per CPU page" - big_cpu_page_indices = cpu_page_indices[:big_page_count] - small_cpu_page_indices = cpu_page_indices[big_page_count:] - big_page_readies = cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None - small_page_readies = cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None - - # 阶段一：统一写出本批 full-attention token 数据。 - for state in full_states: - state.write_to_cpu( - token_indices=token_indices, - cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=cpu_page_indices, - layout=layouts[state.state_name], - cpu_page_ready_flags=cpu_page_ready_flags, - ) - - # 阶段二：分别从大小页池取快照，写入对应 CPU 页面，不修改请求运行态。 - for state in bounded_states: - layout = layouts.get(state.state_name) - if layout is None: - # 部分 bounded 状态在缓存边界没有需保留的历史（如 W128/S128 压缩器），不写入 CPU cache。 - # 仅允许未声明 CPU 传输区域的 State 跳过，避免必要快照因布局遗漏而丢失。 - assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" - continue - if big_page_slot_ids: - state.write_state_pages_to_cpu( - page_kind="big", - page_slot_ids=big_page_slot_ids, - cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=big_cpu_page_indices, - layout=layout, - cpu_page_ready_flags=big_page_readies, - ) - if small_page_slot_ids: - state.write_state_pages_to_cpu( - page_kind="small", - page_slot_ids=small_page_slot_ids, - cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=small_cpu_page_indices, - layout=layout, - cpu_page_ready_flags=small_page_readies, - ) - - def read_req_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - token_indices: torch.Tensor, - layouts: Mapping[str, StateTransferLayout], - big_page_slot_ids: List[int] | None = None, - req_idx: int | None = None, - state_token_len: int | None = None, - ) -> None: - """读回单个请求本次指定的有序前缀页面，并从末页恢复 req_idx 的 bounded 运行态。 - - big_page_slot_ids 只包含完整 CPU 页对应的普通槽位；末尾碎页使用内部预留的 LOAD 大页。 - 存在 bounded 时须提供 state_token_len，即恢复后的绝对前缀长度，不是本批 token 数。 - 存在 bounded 时必须提供 req_idx；最终前缀须满足各 State 的快照或重置边界。 - """ - full_states = self.get_full_attention_states() - bounded_states = self.get_bounded_attention_states() - if bounded_states: - assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" - if bounded_states: - assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" - # CPU 页的 token 数与 bounded 大页边界在启动时统一，直接读取全局配置。 - page_tokens = get_env_start_args().cpu_cache_token_page_size - # token 索引只用于判断本批末页是否完整，不读取 GPU 上的索引值。 - has_partial_page = token_indices.numel() % page_tokens != 0 - assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), "slots must cover complete CPU pages only" - # 只扩展内部列表，预留槽位不写回调用方，也不作为普通大页发布。 - load_big_page_slot_ids = big_page_slot_ids.copy() - if has_partial_page: - load_big_page_slot_ids.append(self.bounded_att_cpu_load_big_page_slot) - - # 阶段一：恢复 full-attention token 数据；各 State 负责对应布局的搬运约束。 - for state in full_states: - state.read_from_cpu( - cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=cpu_page_indices, - token_indices=token_indices, - layout=layouts[state.state_name], - ) - - # 阶段二：完整页和末尾碎页统一读入大页池，碎页使用追加的 LOAD 槽位。 - for state in bounded_states: - if state.state_name not in layouts: - # 无跨步历史的 bounded 状态不占 CPU 页，因此无需读回快照；运行态恢复仍在下一阶段处理。 - # 有 CPU 传输区域的 State 必须具备布局，不能静默跳过。 - assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" - continue - state.read_state_pages_from_cpu( - cpu_page_buffers=cpu_page_buffers, - cpu_page_indices=cpu_page_indices, - big_page_slot_ids=load_big_page_slot_ids, - layout=layouts[state.state_name], - ) - - # 阶段三：从末页恢复请求运行态；纯碎页也走 LOAD 槽位。 - if bounded_states: - self.restore_bounded_att_runtime_from_pages( - req_indices=[req_idx], - state_token_lens=[state_token_len], - page_kind="big", - page_slot_ids=[load_big_page_slot_ids[-1]], - ) - - # ======================================================================== - # 五、DP 迁移 - # ======================================================================== - - def copy_from_dp( - self, - remote_att_manager: AttentionStateManager, - src_token_indices: torch.Tensor, - dst_token_indices: torch.Tensor, - src_big_page_ids: List[int], - dst_big_page_ids: List[int], - src_req_idx: int, - dst_req_idx: int, - ) -> None: - """复制其他 DP 同一 TP 分片的 token、大页和指定请求的 prefill 基准运行态。 - - token 索引由调用方准备在源、目标各自的 CUDA 设备上，目标索引为 int64。 - 不需要复制 token 或大页时传空 Tensor 或空列表；请求运行态仍按请求索引复制。 - 调用方保证源 MTP 游标为 0、目标前缀满足各 State 的缓存边界，以及源页就绪和引用有效。 - """ - assert type(remote_att_manager) is type(self), "DP replicas must use the same model resource manager" - assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" - assert len(src_big_page_ids) == len(dst_big_page_ids), "source and destination page counts differ" - if src_token_indices.numel(): - for state in self.get_full_attention_states(): - state.copy_tokens_from_dp( - src_token_indices=src_token_indices, - dst_token_indices=dst_token_indices, - remote_dp_full_att_state=getattr(remote_att_manager, state.state_name), - ) - if src_big_page_ids: - for state in self.get_bounded_attention_states(): - state.copy_big_pages_from_dp( - src_page_ids=src_big_page_ids, - dst_page_ids=dst_big_page_ids, - remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), - ) - # 运行态由具体 State 复制或按合法缓存边界重置，不由大页复制隐式恢复。 - for state in self.get_bounded_attention_states(): - state.copy_runtime_state_from_dp( - src_req_indices=[src_req_idx], - dst_req_indices=[dst_req_idx], - remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), - ) - - # ======================================================================== - # 六、PD 分离传输 + # 四、CPU/DP/PD 操作均由对应 Operator 提供。 # ======================================================================== - @lru_cache(maxsize=None) - def get_pd_token_page_layouts(self) -> Mapping[str, StateTransferLayout]: - """按实际字段区间生成 kv 布局，各 State 记录自身在页内的起点与终点。 - - 使用全局 pd_kv_page_size，按需缓存；不依赖页大小查询或运行态布局。 - """ - page_tokens = get_env_start_args().pd_kv_page_size - assert page_tokens > 0, "PD pages require a positive token capacity" - layouts, offset = {}, 0 - for state in self.get_full_attention_states(): - regions = tuple( - replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) - for region in state.get_pd_transfer_regions() - ) - # 字段保留自身数据区间，State 终点包含留给下一个资源的对齐空间。 - state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.state_name] = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=page_tokens, - byte_start=offset, - byte_end=state_end, - regions=regions, - ) - offset = state_end - return layouts - - @lru_cache(maxsize=None) - def get_pd_runtime_page_layouts(self) -> Mapping[str, StateTransferLayout]: - """独立生成 att_state 布局，记录各运行态 State 的完整字节区间。 - - 每个 State 保存一份运行态，tokens_per_page 为 None；无 State 时返回空字典。 - """ - layouts, offset = {}, 0 - for state in self.get_bounded_attention_states(): - regions = tuple( - replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) - for region in state.get_pd_transfer_regions() - ) - state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 - layouts[state.state_name] = StateTransferLayout( - state_name=state.state_name, - tokens_per_page=None, - byte_start=offset, - byte_end=state_end, - regions=regions, - ) - offset = state_end - return layouts - - # PD 分离：kind="kv" 传输 token 数据，kind="att_state" 传输 bounded 运行态。 - # pd_page_buffers 第一维为页，pd_page_index 选择物理页；网络收发与同步由传输流程组织。 - def write_to_pd_page( - self, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - kind: Literal["kv", "att_state"], - layouts: Mapping[str, StateTransferLayout], - token_indices: torch.Tensor | None = None, - req_idx: int | None = None, - state_token_len: int | None = None, - ) -> None: - """按任务布局写出 PD 页；att_state 必须传入已执行的绝对前缀长度 state_token_len。""" - match kind: - case "kv": - # KV 页只写出 full-attention 数据；各 State 负责自身布局与容量约束。 - assert token_indices is not None, "KV PD pages require token indices" - for state in self.get_full_attention_states(): - state.write_to_pd_page( - token_indices=token_indices, - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - layout=layouts[state.state_name], - ) - - case "att_state": - # 前缀长度由 PD 任务提供，各 State 据此定位自己的运行态窗口。 - assert req_idx is not None, "state PD pages require a request index" - assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" - - # 阶段一：写出 bounded 的 prefill 基准运行态。 - for state in self.get_bounded_attention_states(): - state.write_to_pd_page( - req_idx=req_idx, - pd_page_buffers=pd_page_buffers, - pd_page_index=pd_page_index, - state_token_len=state_token_len, - layout=layouts[state.state_name], - ) - - case _: - raise AssertionError(f"unknown PD page kind {kind}") - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - kind: Literal["kv", "att_state"], - layouts: Mapping[str, StateTransferLayout], - token_indices: torch.Tensor | None = None, - req_idx: int | None = None, - state_token_len: int | None = None, - ) -> None: - """按布局恢复 PD 页；att_state 使用绝对前缀长度 state_token_len 恢复运行态位置。""" - match kind: - case "kv": - # KV 页只恢复 full-attention 数据；各 State 负责自身布局与容量约束。 - assert token_indices is not None, "KV PD pages require token indices" - for state in self.get_full_attention_states(): - state.read_from_pd_page( - pd_page_buffers=pd_page_buffers, - token_indices=token_indices, - pd_page_index=pd_page_index, - layout=layouts[state.state_name], - ) - - case "att_state": - # 前缀长度由 PD 任务提供，各 State 据此定位自己的运行态窗口。 - assert req_idx is not None, "state PD pages require a request index" - assert state_token_len is not None and state_token_len >= 0, "PD state requires its prefill prefix length" - - # 阶段一：恢复 bounded 的 prefill 基准运行态。 - for state in self.get_bounded_attention_states(): - state.read_from_pd_page( - pd_page_buffers=pd_page_buffers, - req_idx=req_idx, - pd_page_index=pd_page_index, - state_token_len=state_token_len, - layout=layouts[state.state_name], - ) - - case _: - raise AssertionError(f"unknown PD page kind {kind}") - # ======================================================================== # 七、私有辅助实现 # ======================================================================== diff --git a/lightllm/common/attention_state/pd_operator.py b/lightllm/common/attention_state/pd_operator.py new file mode 100644 index 0000000000..576ce7e9ff --- /dev/null +++ b/lightllm/common/attention_state/pd_operator.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from collections.abc import Mapping +from copy import copy +from dataclasses import replace +from functools import lru_cache +from typing import TYPE_CHECKING, Literal + +import torch + +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +logger = init_logger(__name__) + + +class PDOperator: + """封装 PD 分离所需的联合布局、页大小和读写操作。""" + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + @lru_cache(maxsize=None) + def get_pd_full_att_page_layouts(self) -> Mapping[str, StateTransferLayout]: + page_tokens = get_env_start_args().pd_kv_page_size + assert page_tokens > 0, "PD pages require a positive token capacity" + layouts, offset = {}, 0 + for state in self.owner.get_full_attention_states(): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_pd_bounded_att_page_layouts(self) -> Mapping[str, StateTransferLayout]: + layouts, offset = {}, 0 + for state in self.owner.get_bounded_attention_states(): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_page_bytes(self) -> int: + token_layouts = self.get_pd_full_att_page_layouts() + runtime_layouts = self.get_pd_bounded_att_page_layouts() + full_bytes = max((layout.byte_end for layout in token_layouts.values()), default=0) + runtime_bytes = max((layout.byte_end for layout in runtime_layouts.values()), default=0) + if token_layouts and full_bytes < runtime_bytes: + logger.warning( + f"PD KV page size ({full_bytes} bytes) is smaller than the complete attention runtime " + f"({runtime_bytes} bytes); using {runtime_bytes} bytes per transfer page" + ) + return max(full_bytes, runtime_bytes) + + def clear_layout_cache(self) -> None: + self.get_pd_full_att_page_layouts.cache_clear() + self.get_pd_bounded_att_page_layouts.cache_clear() + self.get_page_bytes.cache_clear() + + def create_ipc_copy(self) -> AttentionStateManager: + """保留运行态 Tensor，移除 allocator 和 bounded 页池。""" + ipc_copy = copy(self.owner) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + for state in self.owner.get_bounded_attention_states(): + state_copy = copy(state) + state_copy.big_pages_pool = None + state_copy.small_pages_pool = None + setattr(ipc_copy, state.state_name, state_copy) + ipc_copy._reset_operators() + return ipc_copy + + def write_page( + self, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + match kind: + case "kv": + assert token_indices is not None, "KV PD pages require token indices" + for state in self.owner.get_full_attention_states(): + state.write_to_pd_page( + token_indices=token_indices, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + layout=layouts[state.state_name], + ) + case "att_state": + assert req_idx is not None, "state PD pages require a request index" + assert state_token_len is not None and state_token_len >= 0, ( + "PD state requires its prefill prefix length" + ) + for state in self.owner.get_bounded_attention_states(): + state.write_to_pd_page( + req_idx=req_idx, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + state_token_len=state_token_len, + layout=layouts[state.state_name], + ) + case _: + raise AssertionError(f"unknown PD page kind {kind}") + + def read_page( + self, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + match kind: + case "kv": + assert token_indices is not None, "KV PD pages require token indices" + for state in self.owner.get_full_attention_states(): + state.read_from_pd_page( + pd_page_buffers=pd_page_buffers, + token_indices=token_indices, + pd_page_index=pd_page_index, + layout=layouts[state.state_name], + ) + case "att_state": + assert req_idx is not None, "state PD pages require a request index" + assert state_token_len is not None and state_token_len >= 0, ( + "PD state requires its prefill prefix length" + ) + for state in self.owner.get_bounded_attention_states(): + state.read_from_pd_page( + pd_page_buffers=pd_page_buffers, + req_idx=req_idx, + pd_page_index=pd_page_index, + state_token_len=state_token_len, + layout=layouts[state.state_name], + ) + case _: + raise AssertionError(f"unknown PD page kind {kind}") diff --git a/lightllm/common/attention_state/types.py b/lightllm/common/attention_state/types.py index 8e20c6d6e0..38c1b95691 100644 --- a/lightllm/common/attention_state/types.py +++ b/lightllm/common/attention_state/types.py @@ -229,7 +229,7 @@ class StateTransferLayout: # 本 State 在外部页中的结束字节，不包含该位置；涵盖全部字段及本 State 的对齐、预留空间。 # byte_end - byte_start 是本 State 的占用，不是整张物理页的跨度。 # 同类布局的最大 byte_end 表示整组 State 的占用；CPU 物理页再按 16 字节对齐， - # PD 物理页由 manager.get_pd_page_bytes() 取 KV 与完整运行态布局占用的最大值。 + # PD 物理页由 PDOperator.get_page_bytes() 取 KV 与完整运行态布局占用的最大值。 byte_end: int # 当前 State 参与传输的字段区域；其他 State 的区域放在它们各自的 layout 中。 diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index c51668035b..9351a7739f 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -318,7 +318,7 @@ def init_dp_kv_shared(self): ): if rank_idx != self.rank_in_node: att_state_manager, req_to_token_indices = load_att_resources(rank_idx, kind="dp") - att_state_manager.register_page_pools_pin_memory() + att_state_manager.get_dp_shared_operator().register_page_pools_pin_memory() else: att_state_manager = self.model.att_state_manager req_to_token_indices = self.model.req_manager.req_to_token_indexs diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 4c693a1ac8..2fadb6b3b5 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -138,7 +138,7 @@ def kv_trans(self, trans_tasks: List["TransTask"]) -> None: source = self.backend.dp_att_state_managers[task.src_dp_rank] # 源索引来自远端 CUDA 请求表；allocator 返回 CPU 索引，在此一次性准备目标 CUDA int64 索引。 dst_token_indices = task.dst_token_indices.to(device="cuda", dtype=torch.long, non_blocking=True) - target.copy_from_dp( + target.get_dp_shared_operator().copy_from_dp( remote_att_manager=source, src_token_indices=task.src_token_indices, dst_token_indices=dst_token_indices, diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 54839a5e80..e7a1045074 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -35,8 +35,9 @@ def __init__(self, backend): self.cpu_cache_handle_queue: Deque[TransTask] = deque() self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False) resources = self.backend.model.att_state_manager + self.cpu_cache_operator = resources.get_cpu_cache_operator() meta = self.cpu_cache_client.kv_cache_tensor_meta - if meta.page_bytes != resources.get_cpu_page_bytes(): + if meta.page_bytes != self.cpu_cache_operator.get_page_bytes(): raise ValueError("CPU cache service and model resource page sizes differ") @lru_cache() @@ -148,6 +149,7 @@ def read_req_from_cpu( if not cpu_page_indices: return resources = self.backend.model.att_state_manager + cpu_cache_operator = resources.get_cpu_cache_operator() req_manager = self.backend.model.req_manager need_token_num = page_lengths[-1] - req.cur_kv_len @@ -177,11 +179,11 @@ def read_req_from_cpu( assert big_page_slot_ids is not None, "not enough bounded state slots for CPU cache load" # 阶段三：统一恢复各类 State；异常直接向上传播，不在此处回收或继续执行。 - resources.read_req_from_cpu( + cpu_cache_operator.read_req( cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, cpu_page_indices=cpu_page_indices, token_indices=token_indices, - layouts=resources.get_cpu_page_layouts(), + layouts=cpu_cache_operator.get_page_layouts(), big_page_slot_ids=big_page_slot_ids, req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, state_token_len=page_lengths[-1], @@ -335,11 +337,11 @@ def write_req_to_cpu( if includes_tail_page and page_len_list[-1] % page_tokens != 0: small_page_slot_ids.append(req.tail_small_page_buffer_id) # 整批提交并保留 token/page 对齐，由资源层跳过 ready 页。 - resources.write_req_to_cpu( + resources.get_cpu_cache_operator().write_req( token_indices=token_indexes, cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, cpu_page_indices=page_list, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), cpu_page_ready_flags=ready_list, big_page_slot_ids=big_page_slot_ids, small_page_slot_ids=small_page_slot_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index 370a17694a..54fd9bc119 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -14,7 +14,7 @@ def allocate_transfer_buffer(resources: AttentionStateManager, page_count: int, device: str | torch.device | int): if page_count <= 0: raise ValueError("PD buffer must contain pages") - return torch.empty((page_count, resources.get_pd_page_bytes()), dtype=torch.uint8, device=device) + return torch.empty((page_count, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device=device) def copy_page( @@ -45,23 +45,25 @@ def copy_page( ), "invalid PD group or page index" tp_transfer_layouts = [] for att_manager in tp_att_managers: + pd_operator = att_manager.get_pd_operator() match task_kind: case "kv": - state_layouts = att_manager.get_pd_token_page_layouts() + state_layouts = pd_operator.get_pd_full_att_page_layouts() case "att_state": - state_layouts = att_manager.get_pd_runtime_page_layouts() + state_layouts = pd_operator.get_pd_bounded_att_page_layouts() case _: assert False, f"unknown PD task kind {task_kind}" - assert state_layouts and att_manager.get_pd_page_bytes() == pd_page_buffers.shape[1], ( + assert state_layouts and pd_operator.get_page_bytes() == pd_page_buffers.shape[1], ( "PD page layout does not match this group" ) tp_transfer_layouts.append(state_layouts) for att_manager, state_layouts in zip(tp_att_managers, tp_transfer_layouts): + pd_operator = att_manager.get_pd_operator() match task_kind: case "kv": if write_to_pd_page: - att_manager.write_to_pd_page( + pd_operator.write_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -69,7 +71,7 @@ def copy_page( token_indices=token_indices, ) else: - att_manager.read_from_pd_page( + pd_operator.read_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -78,7 +80,7 @@ def copy_page( ) case "att_state": if write_to_pd_page: - att_manager.write_to_pd_page( + pd_operator.write_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, @@ -87,7 +89,7 @@ def copy_page( state_token_len=state_token_len, ) else: - att_manager.read_from_pd_page( + pd_operator.read_page( pd_page_buffers=pd_page_buffers, pd_page_index=pd_page_index, kind=task_kind, diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py index 312d524d50..1d4504015a 100644 --- a/lightllm/server/router/model_infer/shared_att_resources_utils.py +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -35,9 +35,9 @@ def publish_att_resources( mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ count = reader_count if kind == "pd": - ipc_resources = resources.create_pd_ipc_copy() + ipc_resources = resources.get_pd_operator().create_ipc_copy() else: - ipc_resources = resources.create_dp_ipc_copy() + ipc_resources = resources.get_dp_shared_operator().create_ipc_copy() with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): # 每个读者独立序列化 CUDA IPC 引用；计数偏移变化可能使各份 pickle 长度不同。 payloads = [ForkingPickler.dumps((ipc_resources, request_mapping)).tobytes() for _ in range(count)] diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index 942f8002b9..abcdbd0650 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -58,4 +58,4 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: assert ( args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num ), "hybrid CPU pages must match the large state-page token boundary" - return manager.get_cpu_page_layouts() + return manager.get_cpu_cache_operator().get_page_layouts() diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 52577d4ce1..d8beb27faf 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -105,22 +105,22 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP next_token = output.logits.argmax(dim=-1).cpu() source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() - cpu = torch.zeros((len(pages), resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - resources.write_req_to_cpu( + cpu = torch.zeros((len(pages), resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + resources.get_cpu_cache_operator().write_req( token_indices=source_indices, cpu_page_buffers=cpu, cpu_page_indices=pages, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=source_slots if hybrid else None, ) # Host-side checkpoint repacking consumes the CPU page on the next call. torch.cuda.synchronize() - full_att_layouts = resources.get_pd_token_page_layouts() - runtime_layouts = resources.get_pd_runtime_page_layouts() - pd = {"kv": torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda")} + full_att_layouts = resources.get_pd_operator().get_pd_full_att_page_layouts() + runtime_layouts = resources.get_pd_operator().get_pd_bounded_att_page_layouts() + pd = {"kv": torch.empty((1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda")} if runtime_layouts: - pd["att_state"] = torch.empty((1, resources.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - resources.write_to_pd_page( + pd["att_state"] = torch.empty((1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + resources.get_pd_operator().write_page( pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", @@ -128,7 +128,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=source_indices, ) if "att_state" in pd: - resources.write_to_pd_page( + resources.get_pd_operator().write_page( pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", @@ -155,17 +155,17 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP slots = big.alloc(len(pages)) if hybrid else None if hybrid: allocated_slots.extend(slots) - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=cpu, cpu_page_indices=pages, token_indices=indices, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=slots, req_idx=int(request[0]) if hybrid else None, state_token_len=prefix_len if hybrid else None, ) elif kind == "dp_copy": - resources.copy_from_dp( + resources.get_dp_shared_operator().copy_from_dp( remote_att_manager=resources, src_token_indices=source_indices, dst_token_indices=indices, @@ -175,7 +175,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP dst_req_idx=int(request[0]), ) else: - resources.read_from_pd_page( + resources.get_pd_operator().read_page( pd_page_buffers=pd["kv"], pd_page_index=0, kind="kv", @@ -183,7 +183,7 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP token_indices=indices, ) if "att_state" in pd: - resources.read_from_pd_page( + resources.get_pd_operator().read_page( pd_page_buffers=pd["att_state"], pd_page_index=0, kind="att_state", diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 8236c5ff36..5912d2bbda 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -88,28 +88,28 @@ def main(): resources.finish_state_registration() pages = list(range(args.pages)) slots = resources.bounded_att_big_page_allocator.alloc(args.pages) - cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) # Populate every TP region through the production API before measuring restores. full.k_storage.zero_() full.v_storage.zero_() resources.save_bounded_att_runtime_to_pages([0] * args.pages, page_kind="big", page_slot_ids=slots) torch.cuda.synchronize() token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") - resources.write_req_to_cpu( + resources.get_cpu_cache_operator().write_req( token_indices=token_indices, cpu_page_buffers=cpu_page_buffers, cpu_page_indices=pages, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=slots, ) torch.cuda.synchronize() def restore(): - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=cpu_page_buffers, cpu_page_indices=pages, token_indices=token_indices, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=slots, req_idx=0, ) @@ -132,7 +132,7 @@ def restore(): ).big_pages_pool state_bytes = args.pages * sum(t[0].numel() * t.element_size() for t in (pool.conv_state, pool.ssm_state)) median = statistics.median(times) - print(f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; CPU bytes/page: {resources.get_cpu_page_bytes():,}") + print(f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; CPU bytes/page: {resources.get_cpu_cache_operator().get_page_bytes():,}") print(f"Mean: {statistics.mean(times) * 1000:.3f} ms; median: {median * 1000:.3f} ms") print(f"KV + state payload throughput: {(token_bytes + state_bytes) / median / 1024**3:.3f} GiB/s") diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index 409f27fe80..36ff514ab6 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING from dataclasses import replace -from functools import lru_cache from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState @@ -43,15 +42,3 @@ def __init__(self, device: str | torch.device | int = "cpu"): self.bounded_att_big_page_allocator = StatePageAllocator(3) self.bounded_att_small_page_allocator = StatePageAllocator(2) self.finish_state_registration() - - @lru_cache(maxsize=None) - def get_cpu_page_layouts(self): - raise NotImplementedError("This test manager uses explicitly supplied CPU layouts") - - @lru_cache(maxsize=None) - def get_pd_token_page_layouts(self): - raise NotImplementedError("This test manager uses explicitly supplied PD layouts") - - @lru_cache(maxsize=None) - def get_pd_runtime_page_layouts(self): - raise NotImplementedError("This test manager uses explicitly supplied PD layouts") diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 5adedc5408..b3966eb62c 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -89,6 +89,19 @@ def rebuild_full_state( return prepared_full_state +def replace_logical_layer_groups(config, logical_layer_groups): + """替换测试配置的逻辑层分组，兼容 Packed MLA 的嵌套 logical_config。""" + if isinstance(config, PackedMLAAttentionConfig): + return replace( + config, + logical_config=replace( + config.logical_config, + logical_layer_groups=logical_layer_groups, + ), + ) + return replace(config, logical_layer_groups=logical_layer_groups) + + def owned_states(attributes: Mapping[str, object], state_class: type | tuple[type, ...]): """从显式成员选择指定类型的 State，按资源身份排序。""" return tuple( diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index d3d0705d1a..d5d5ace566 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -330,10 +330,10 @@ def start_runtime(): assert torch.equal(state.q_scales[:, 0].cpu(), scales[start:end] * 2) layouts = build_cpu_cache_layouts() - assert (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 == manager.get_cpu_page_bytes() - assert layouts.keys() == manager.get_cpu_page_layouts().keys() + assert (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 == manager.get_cpu_cache_operator().get_page_bytes() + assert layouts.keys() == manager.get_cpu_cache_operator().get_page_layouts().keys() for state_name, layout in layouts.items(): - runtime_layout = manager.get_cpu_page_layouts()[state_name] + runtime_layout = manager.get_cpu_cache_operator().get_page_layouts()[state_name] assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in layout.regions] == [ (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in runtime_layout.regions ] diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index 92f70a2ab6..5d4425663c 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -118,7 +118,7 @@ def test_budget_matches_token_stores_including_hold_and_padding( planning.register_state("main", rebuild_full_state(template, config=template.config)) planning.finish_state_registration() token_bytes = planning.get_full_att_token_bytes() - pd_page_bytes = planning.get_pd_page_bytes() + pd_page_bytes = planning.get_pd_operator().get_page_bytes() assert not hasattr(planning, "pd_page_bytes") assert not hasattr(planning, "pd_transfer_layouts") assert token_bytes * (20 + 4) == actual @@ -133,14 +133,14 @@ def test_budget_matches_token_stores_including_hold_and_padding( manager.full_att_allocator = FullAttKVAllocator(20, 4, pin_memory=False) manager.finish_state_registration() (state,) = manager.get_full_attention_states("main") - assert manager.get_pd_page_bytes() == pd_page_bytes + assert manager.get_pd_operator().get_page_bytes() == pd_page_bytes args = StartArgs( page_size=4, run_mode="prefill", pd_kv_page_size=8, pd_kv_page_num=3, ) - reserve = manager.get_pd_page_bytes() * args.pd_kv_page_num + reserve = manager.get_pd_operator().get_page_bytes() * args.pd_kv_page_num monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 0) @@ -176,7 +176,7 @@ def test_budget_matches_token_stores_including_hold_and_padding( device="cpu", ) assert resources.full_att_allocator.size == expected - assert resources.get_pd_page_bytes() == manager.get_pd_page_bytes() + assert resources.get_pd_operator().get_page_bytes() == manager.get_pd_operator().get_page_bytes() def test_shared_allocator_accounts_for_distinct_state_stores(): @@ -244,14 +244,14 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( manager.register_state("main", tail) manager.finish_state_registration() # PD 页只包含运行态数据，使用第二页验证 8 字节页跨度下的读写。 - assert manager.get_pd_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 + assert manager.get_pd_operator().get_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 reference.conv_state.fill_(2) reference.ssm_state.fill_(3) tail.buffer.fill_(4) - pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - layouts = manager.get_pd_runtime_page_layouts() + pages = torch.zeros((2, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + layouts = manager.get_pd_operator().get_pd_bounded_att_page_layouts() assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "compressor"} - manager.write_to_pd_page( + manager.get_pd_operator().write_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -262,7 +262,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( reference.conv_state.zero_() reference.ssm_state.zero_() tail.buffer.zero_() - manager.read_from_pd_page( + manager.get_pd_operator().read_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -281,7 +281,7 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( monkeypatch.setattr( state, "get_pd_transfer_regions", lambda *args, **kwargs: pytest.fail("cached page size must not rebuild regions") ) - assert manager.get_pd_page_bytes() == pages.shape[1] + assert manager.get_pd_operator().get_page_bytes() == pages.shape[1] @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -307,7 +307,7 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( manager.register_state(role, state_cls(config, page_size=1, device="meta")) manager.finish_state_registration() # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 - assert manager.get_pd_page_bytes() == 72 + assert manager.get_pd_operator().get_page_bytes() == 72 monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) @@ -317,16 +317,16 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( big_page_tokens=4, device="cuda", ) - assert manager.get_pd_page_bytes() == 72 + assert manager.get_pd_operator().get_page_bytes() == 72 states = manager.get_full_attention_states() for i, state in enumerate(states, 1): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.fill_(i) - pages = torch.zeros((2, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + pages = torch.zeros((2, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") page = pages[1] - layouts = manager.get_pd_token_page_layouts() + layouts = manager.get_pd_operator().get_pd_full_att_page_layouts() assert [min(region.byte_start for region in layouts[state.state_name].regions) for state in states] == [0, 8, 48] - manager.write_to_pd_page( + manager.get_pd_operator().write_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind="kv", @@ -337,7 +337,7 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( for state in states: for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() - manager.read_from_pd_page( + manager.get_pd_operator().read_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind="kv", @@ -390,11 +390,12 @@ def _get_transfer_regions(self, page_tokens: int): for state in manager.get_full_attention_states(): assert state.get_token_bytes() * 3 == 792 monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("capacity must use field layouts")) - get_layouts = manager.get_cpu_page_layouts if transfer == "cpu" else manager.get_pd_token_page_layouts - get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes + get_layouts = manager.get_cpu_cache_operator().get_page_layouts if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts + get_page_bytes = manager.get_cpu_cache_operator().get_page_bytes if transfer == "cpu" else manager.get_pd_operator().get_page_bytes if layout_first: with monkeypatch.context() as patch: - patch.setattr(manager, f"get_{transfer}_page_bytes", lambda: pytest.fail("layout must not query page bytes")) + operator = manager.get_cpu_cache_operator() if transfer == "cpu" else manager.get_pd_operator() + patch.setattr(operator, "get_page_bytes", lambda: pytest.fail("layout must not query page bytes")) layouts = get_layouts() assert get_page_bytes() == 1600 layouts = get_layouts() @@ -415,27 +416,27 @@ def _get_transfer_regions(self, page_tokens: int): field.tensor.zero_() field.tensor[:, :3].fill_(index * 4 + field_index + 1) if transfer == "cpu": - manager.write_req_to_cpu( + manager.get_cpu_cache_operator().write_req( token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), cpu_page_buffers=pages, cpu_page_indices=[1], layouts=layouts, ) - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), layouts=layouts, ) else: - manager.write_to_pd_page( + manager.get_pd_operator().write_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", layouts=layouts, token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), ) - manager.read_from_pd_page( + manager.get_pd_operator().read_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", @@ -495,22 +496,30 @@ def get_pd_transfer_regions(self): # Conv 为 6 字节，SSM 从偏移 72 开始，占 4 字节；完整运行态需要 80 字节。 # 即使 KV 页不足，布局仍可独立查询；最终块大小查询选择两类占用的最大值。 with monkeypatch.context() as patch: - patch.setattr(manager, "get_pd_token_page_layouts", lambda: pytest.fail("runtime layout must be independent")) - patch.setattr(manager, "get_pd_page_bytes", lambda: pytest.fail("layout must not query page bytes")) - layouts = manager.get_pd_runtime_page_layouts() + patch.setattr( + manager.get_pd_operator(), + "get_pd_full_att_page_layouts", + lambda: pytest.fail("runtime layout must be independent"), + ) + patch.setattr( + manager.get_pd_operator(), + "get_page_bytes", + lambda: pytest.fail("layout must not query page bytes"), + ) + layouts = manager.get_pd_operator().get_pd_bounded_att_page_layouts() layout = layouts[state.state_name] assert layout.byte_end == 80 and layout.tokens_per_page is None assert layout.regions[-1].byte_start == 72 warnings = [] - monkeypatch.setattr("lightllm.common.attention_state.manager.logger.warning", warnings.append) + monkeypatch.setattr("lightllm.common.attention_state.pd_operator.logger.warning", warnings.append) expected = max(80, (full_page_tokens or 0) * 8) - assert manager.get_pd_page_bytes() == expected - assert manager.get_pd_runtime_page_layouts() is layouts + assert manager.get_pd_operator().get_page_bytes() == expected + assert manager.get_pd_operator().get_pd_bounded_att_page_layouts() is layouts # 实际传输页可大于运行态布局，选中第二页后仍从该页起点搬运。 pages = torch.full((3, expected), 91, dtype=torch.uint8, device="cuda") state.conv_state.fill_(3) state.ssm_state.fill_(5) - manager.write_to_pd_page( + manager.get_pd_operator().write_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -520,7 +529,7 @@ def get_pd_transfer_regions(self): ) state.conv_state.zero_() state.ssm_state.zero_() - manager.read_from_pd_page( + manager.get_pd_operator().read_page( pd_page_buffers=pages, pd_page_index=1, kind="att_state", @@ -534,7 +543,7 @@ def get_pd_transfer_regions(self): if warnings: assert "64 bytes" in warnings[0] and "using 80 bytes per transfer page" in warnings[0] # 相同 manager 的容量查询命中缓存，不重复打印 warning。 - assert manager.get_pd_page_bytes() == expected + assert manager.get_pd_operator().get_page_bytes() == expected assert len(warnings) == int(full_page_tokens == 8) @@ -565,7 +574,7 @@ def _get_transfer_regions(self, page_tokens: int): ) manager.register_state("main", PaddedState(config, page_size=1, device="meta")) manager.finish_state_registration() - get_page_bytes = manager.get_cpu_page_bytes if transfer == "cpu" else manager.get_pd_page_bytes + get_page_bytes = manager.get_cpu_cache_operator().get_page_bytes if transfer == "cpu" else manager.get_pd_operator().get_page_bytes assert get_page_bytes() == 48 monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) @@ -577,7 +586,7 @@ def _get_transfer_regions(self, page_tokens: int): device="cpu", ) assert get_page_bytes() == (80 if transfer == "cpu" else 72) - layouts = manager.get_cpu_page_layouts() if transfer == "cpu" else manager.get_pd_token_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts() assert layouts["main_gqa_bf16_state"].byte_end == 72 @@ -628,7 +637,7 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( page_tokens = runtime_page_bytes // 8 + extra_tokens pd_args.pd_kv_page_size = page_tokens manager.finish_state_registration() - assert manager.get_pd_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) + assert manager.get_pd_operator().get_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) assert manager.full_att_allocator is None assert read_kv_storage(manager.main_gqa_fp32_state) is None assert not hasattr(manager, "pd_transfer_layouts") @@ -696,7 +705,7 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( if isinstance(state, FullAttentionState): monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("CPU capacity must use field layouts")) monkeypatch.setattr(state, "get_pd_transfer_regions", lambda: pytest.fail("CPU layout must be independent")) - layouts = manager.get_cpu_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() assert min(region.byte_start for layout in layouts.values() for region in layout.regions) == 0 assert not {"state_token_len", "format_version"} & { region.field for layout in layouts.values() for region in layout.regions @@ -705,8 +714,8 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( assert all(layout.tokens_per_page == page_tokens for layout in layouts.values()) for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): monkeypatch.setattr(state, "get_cpu_transfer_regions", lambda: pytest.fail("cached layout must be reused")) - assert manager.get_cpu_page_bytes() == expected_page_bytes - assert manager.get_cpu_page_layouts() is layouts + assert manager.get_cpu_cache_operator().get_page_bytes() == expected_page_bytes + assert manager.get_cpu_cache_operator().get_page_layouts() is layouts def test_cpu_bounded_capacity_includes_state_field_padding(cpu_args: EasyDict): @@ -737,8 +746,8 @@ def get_cpu_transfer_regions(self): manager.register_state("main", state) manager.finish_state_registration() # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 - assert manager.get_cpu_page_bytes() == 80 - layout = manager.get_cpu_page_layouts()[state.state_name] + assert manager.get_cpu_cache_operator().get_page_bytes() == 80 + layout = manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] assert (layout.byte_start, layout.byte_end) == (0, 80) assert layout.regions[-1].byte_end == 76 @@ -765,7 +774,7 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea manager.register_state("draft0", states[1]) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) manager.finish_state_registration() - layouts = manager.get_cpu_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() packed, gqa = manager.get_full_attention_states() packed_layout, gqa_layout = layouts[packed.state_name], layouts[gqa.state_name] data_end = max(r.byte_end for r in packed_layout.regions) @@ -778,12 +787,12 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea field.tensor[:, :3].fill_(value) field.tensor[:, 4:7].zero_() pages = torch.full( - size=(3, manager.get_cpu_page_bytes()), + size=(3, manager.get_cpu_cache_operator().get_page_bytes()), fill_value=165, dtype=torch.uint8, pin_memory=True, ) - manager.write_req_to_cpu( + manager.get_cpu_cache_operator().write_req( token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), cpu_page_buffers=pages, cpu_page_indices=[1], @@ -792,7 +801,7 @@ def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: Ea torch.cuda.synchronize() assert (pages[1, data_end:gqa_start] == 165).all() assert (pages[[0, 2]] == 165).all() - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), @@ -831,10 +840,10 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas manager.register_state(role, state) manager.finish_state_registration() states = manager.get_bounded_attention_states() - layouts = manager.get_cpu_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() # 两份快照分别含 10/22 字节数据，预留 16/24 字节；页 token 数不放大快照。 assert [min(r.byte_start for r in layouts[state.state_name].regions) for state in states] == [0, 16] - assert manager.get_cpu_page_bytes() == 48 + assert manager.get_cpu_cache_operator().get_page_bytes() == 48 assert [(layout.byte_start, layout.byte_end) for layout in layouts.values()] == [(0, 16), (16, 40)] for i, state in enumerate(states): state.conv_state.fill_(i + 3) @@ -846,8 +855,8 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas page_kind="big", page_slot_ids=source_slots, ) - pages = torch.full((2, manager.get_cpu_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) - manager.write_req_to_cpu( + pages = torch.full((2, manager.get_cpu_cache_operator().get_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) + manager.get_cpu_cache_operator().write_req( token_indices=torch.tensor([], dtype=torch.long), cpu_page_buffers=pages, cpu_page_indices=[1], @@ -859,7 +868,7 @@ def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: Eas assert (pages[1, 12:16] == 87).all() manager.reset_bounded_att_state([0]) target_slots = manager.bounded_att_big_page_allocator.alloc(1) - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=torch.tensor([], dtype=torch.long), @@ -903,7 +912,15 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation request_capacity=1, device="cpu", )) manager.finish_state_registration() - getter = getattr(manager, f"get_{transfer}_page_layouts") + if transfer == "cpu": + operator = manager.get_cpu_cache_operator() + getter = operator.get_page_layouts + elif transfer == "pd_token": + operator = manager.get_pd_operator() + getter = operator.get_pd_full_att_page_layouts + else: + operator = manager.get_pd_operator() + getter = operator.get_pd_bounded_att_page_layouts before = getter() assert getter() is before state = manager.main_gdn_conv_bf16_ssm_bf16_state if transfer == "pd_runtime" else manager.main_gqa_bf16_state @@ -923,7 +940,7 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation assert after_layout is not before_layout full_state = manager.main_gqa_bf16_state assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta - page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() + page_bytes = manager.get_cpu_cache_operator().get_page_bytes() if transfer == "cpu" else manager.get_pd_operator().get_page_bytes() if transfer == "pd_runtime": assert after_layout.byte_end == 8 and page_bytes == 40 elif transfer == "cpu": @@ -966,7 +983,7 @@ def get_cpu_transfer_regions(self): manager.register_state("draft0", prepared_full_state) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) manager.finish_state_registration() - layouts = manager.get_cpu_page_layouts() if transfer == "cpu" else manager.get_pd_token_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts() for index, state in enumerate(manager.get_full_attention_states()): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() @@ -980,7 +997,7 @@ def get_cpu_transfer_regions(self): assert k_region.byte_start == index * state_bytes + k_offset assert v_region.byte_start == index * state_bytes + state_bytes // 2 - k_offset - page_bytes = manager.get_cpu_page_bytes() if transfer == "cpu" else manager.get_pd_page_bytes() + page_bytes = manager.get_cpu_cache_operator().get_page_bytes() if transfer == "cpu" else manager.get_pd_operator().get_page_bytes() pages = torch.zeros( size=(1, page_bytes), dtype=torch.uint8, @@ -988,13 +1005,13 @@ def get_cpu_transfer_regions(self): pin_memory=transfer == "cpu", ) if transfer == "cpu": - manager.write_req_to_cpu( + manager.get_cpu_cache_operator().write_req( token_indices=torch.arange(4, dtype=torch.long, device="cuda"), cpu_page_buffers=pages, cpu_page_indices=[0], layouts=layouts, ) - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[0], token_indices=torch.arange( @@ -1006,14 +1023,14 @@ def get_cpu_transfer_regions(self): layouts=layouts, ) else: - manager.write_to_pd_page( + manager.get_pd_operator().write_page( pd_page_buffers=pages, pd_page_index=0, kind="kv", layouts=layouts, token_indices=torch.arange(4, dtype=torch.long, device="cuda"), ) - manager.read_from_pd_page( + manager.get_pd_operator().read_page( pd_page_buffers=pages, pd_page_index=0, kind="kv", diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py index 031bc33f40..85988d5edd 100644 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -5,6 +5,7 @@ from unit_tests.common.att_resources.model_helpers import read_kv_storage from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from dataclasses import replace from types import SimpleNamespace @@ -32,7 +33,7 @@ def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, template = definition(kind) state = rebuild_full_state( template=template, - config=replace(template.config, logical_layer_groups=((3,), (8,))), + config=replace_logical_layer_groups(template.config, ((3,), (8,))), token_capacity=8, device="cuda", ) diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py index 5e5a8a236c..84fa1f19d0 100644 --- a/unit_tests/common/att_resources/test_dsv4_state.py +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -455,17 +455,17 @@ def test_manager_joint_cpu_page_roundtrip(monkeypatch: pytest.MonkeyPatch) -> No runtime_buffer(state).random_(0, 127) manager.save_bounded_att_runtime_to_pages([0], [256], "big", [0]) expected = gather_swa_reference(bounded[0], 0, 256) - layouts = manager.get_cpu_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() assert bounded[-1].state_name not in layouts - pages = torch.zeros((2, manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - manager.write_req_to_cpu( + pages = torch.zeros((2, manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + manager.get_cpu_cache_operator().write_req( token_indices=torch.arange(256, dtype=torch.int64, device="cuda"), cpu_page_buffers=pages, cpu_page_indices=[1], layouts=layouts, big_page_slot_ids=[0], ) - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=torch.arange(256, 512, dtype=torch.int64, device="cuda"), @@ -479,7 +479,7 @@ def test_manager_joint_cpu_page_roundtrip(monkeypatch: pytest.MonkeyPatch) -> No assert torch.equal(full.storage[:, :256], full.storage[:, 256:512]) assert torch.equal(expected, gather_swa_reference(bounded[0], 1, 256)) assert not pages[0].any() - ipc = manager.create_pd_ipc_copy() + ipc = manager.get_pd_operator().create_ipc_copy() assert all(state.big_pages_pool is None for state in ipc.get_bounded_attention_states()) assert bounded[0].big_pages_pool is not None @@ -499,11 +499,11 @@ def test_atomic_dp_ipc_preserves_shared_pools(monkeypatch: pytest.MonkeyPatch) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1) manager.register_state("main", state) manager.finish_state_registration() - exported = manager.create_dp_ipc_copy() + exported = manager.get_dp_shared_operator().create_ipc_copy() for state in exported.get_bounded_attention_states(): for pool in (state.big_pages_pool, state.small_pages_pool): if pool is not None: assert pool.buffer.is_shared() - exported.register_page_pools_pin_memory() + exported.get_dp_shared_operator().register_page_pools_pin_memory() finally: get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 8c79ea9d36..6bb3f32774 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -5,6 +5,7 @@ from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from typing import TYPE_CHECKING from dataclasses import replace @@ -46,7 +47,7 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( q_scales = None if kind == "static_tensor" else template.q_scales + value states.append(rebuild_full_state( template=template, - config=replace(template.config, logical_layer_groups=((3,), (8,))), + config=replace_logical_layer_groups(template.config, ((3,), (8,))), token_capacity=8, device="cuda", scales=scales, diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index a239b823df..29a1fa0265 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -94,7 +94,7 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( (state,) = reference.get_full_attention_states("main") args = get_env_start_args() args.run_mode, args.pd_kv_page_size, args.pd_kv_page_num = mode, 8, 3 - transfer_bytes = reference.get_pd_page_bytes() * 3 if mode == "prefill" else 0 + transfer_bytes = reference.get_pd_operator().get_page_bytes() * 3 if mode == "prefill" else 0 storage_bytes = allocated_bytes(state) budget = storage_bytes + transfer_bytes monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (budget, 1 << 30)) @@ -114,4 +114,4 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( assert state.capacity == 16 assert allocated_bytes(state) == storage_bytes if mode == "prefill": - assert model.att_state_manager.get_pd_page_bytes() == reference.get_pd_page_bytes() + assert model.att_state_manager.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 596682785d..6576708961 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -13,6 +13,7 @@ full_state_test_buffers, full_state_test_fields, rebuild_full_state, + replace_logical_layer_groups, ) from unit_tests.common.att_resources.test_capacity import definition, allocated_bytes @@ -26,7 +27,7 @@ @pytest.mark.parametrize("kind", SCHEMES) def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") -> None: template = definition(kind) - config = replace(template.config, logical_layer_groups=((0, 4), (8, 11))) + config = replace_logical_layer_groups(template.config, ((0, 4), (8, 11))) state = rebuild_full_state( template=template, config=config, @@ -60,7 +61,7 @@ def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") - @pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) def test_invalid_layer_groups_are_rejected(kind: str, groups: tuple) -> None: with pytest.raises(AssertionError): - replace(definition(kind).config, logical_layer_groups=groups) + replace_logical_layer_groups(definition(kind).config, groups) def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "EasyDict") -> None: @@ -95,13 +96,13 @@ def test_shared_layers_transfer_each_physical_layer_once( template = definition(kind) source = rebuild_full_state( template=template, - config=replace(template.config, logical_layer_groups=((0, 4), (8,))), + config=replace_logical_layer_groups(template.config, ((0, 4), (8,))), token_capacity=8, device="cuda", ) target = rebuild_full_state( template=template, - config=replace(template.config, logical_layer_groups=((20, 24), (28,))), + config=replace_logical_layer_groups(template.config, ((20, 24), (28,))), token_capacity=8, device="cuda", ) @@ -120,34 +121,34 @@ def test_shared_layers_transfer_each_physical_layer_once( src_ids = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") dst_ids = torch.tensor([1, 4, 7], dtype=torch.int32, device="cuda") if transfer == "cpu": - pages = torch.zeros((2, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - managers[0].write_req_to_cpu( + pages = torch.zeros((2, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + managers[0].get_cpu_cache_operator().write_req( token_indices=src_ids, cpu_page_buffers=pages, cpu_page_indices=[1], - layouts=managers[0].get_cpu_page_layouts(), + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), ) - managers[1].read_req_from_cpu( + managers[1].get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=[1], token_indices=dst_ids, - layouts=managers[1].get_cpu_page_layouts(), + layouts=managers[1].get_cpu_cache_operator().get_page_layouts(), ) else: - pages = torch.zeros((2, managers[0].get_pd_page_bytes()), dtype=torch.uint8, device="cuda") - managers[0].write_to_pd_page( + pages = torch.zeros((2, managers[0].get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + managers[0].get_pd_operator().write_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=src_ids, - layouts=managers[0].get_pd_token_page_layouts(), + layouts=managers[0].get_pd_operator().get_pd_full_att_page_layouts(), ) - managers[1].read_from_pd_page( + managers[1].get_pd_operator().read_page( pd_page_buffers=pages, pd_page_index=1, kind="kv", token_indices=dst_ids, - layouts=managers[1].get_pd_token_page_layouts(), + layouts=managers[1].get_pd_operator().get_pd_full_att_page_layouts(), ) torch.cuda.synchronize() source_fields = full_state_test_fields(source, source_buffers) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 3d7fc0c3d2..5ad569dd87 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -12,6 +12,7 @@ from unit_tests.common.att_resources.test_bounded_state import make_state from unit_tests.common.att_resources.test_window_compressor_state import make_window from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer @@ -89,7 +90,7 @@ def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> Non plain = definition("gqa") quantized = definition("int8") quantized = type(quantized)( - layout_config=replace(quantized.config, logical_layer_groups=((10, 12),)), device="meta", + layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), device="meta", ) manager.register_state("main", plain) manager.register_state("main", quantized) @@ -151,7 +152,7 @@ def test_mla_selection_excludes_same_layer_tail() -> None: tail = make_window() template = definition("mla") state = type(template)( - layout_config=replace(template.config, logical_layer_groups=((3, 7),)), device="meta", + layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), device="meta", ) manager.register_state("main", state) manager.register_state("main", tail) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index 9da30911eb..b0e6670ba6 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -157,7 +157,7 @@ def record_restore( monkeypatch.setattr(manager, "restore_bounded_att_runtime_from_pages", record_restore) # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), cpu_page_indices=[0, 1], token_indices=torch.arange(6, device="cuda"), @@ -186,7 +186,7 @@ def record_length( for state in manager.get_bounded_attention_states(): monkeypatch.setattr(state, operation, record_length) - getattr(manager, operation)( + getattr(manager.get_pd_operator(), "write_page" if operation == "write_to_pd_page" else "read_page")( pd_page_buffers=torch.empty((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), pd_page_index=0, kind="att_state", @@ -220,7 +220,7 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i page_slot_ids=small_slots, ) torch.cuda.synchronize() - source.write_req_to_cpu( + source.get_cpu_cache_operator().write_req( token_indices=torch.arange(token_count, device="cuda"), cpu_page_buffers=pages, cpu_page_indices=page_ids, @@ -237,7 +237,7 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i slots = allocator.alloc(token_count // 4) original_slots = slots.copy() free_count = allocator.free_count - target.read_req_from_cpu( + target.get_cpu_cache_operator().read_req( cpu_page_buffers=pages, cpu_page_indices=page_ids, token_indices=torch.arange(4, 4 + token_count, device="cuda"), @@ -296,7 +296,7 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo pin_memory=True, ) pages = torch.full_like(baseline, 139, pin_memory=True) - source.write_req_to_cpu( + source.get_cpu_cache_operator().write_req( token_indices=token_ids, cpu_page_buffers=baseline, cpu_page_indices=page_ids, @@ -312,7 +312,7 @@ def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bo # ready 页不能访问其页池快照；负的越界槽位可检测是否在构造视图前跳过。 masked_big_slots = [-999 if readies[i] else slot for i, slot in enumerate(big_slots)] masked_small_slots = [-999] if readies[2] else small_slots - source.write_req_to_cpu( + source.get_cpu_cache_operator().write_req( token_indices=token_ids, cpu_page_buffers=pages, cpu_page_indices=page_ids, diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 7745ef9a32..af8dbfcde0 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -109,7 +109,7 @@ def check_runtime_before_profile(device: str | torch.device | int): observed.append(main_linear_state(resources).conv_state.data_ptr()) reserve = 0 if mode == "prefill": - reserve = resources.get_pd_page_bytes() * args.pd_kv_page_num + reserve = resources.get_pd_operator().get_page_bytes() * args.pd_kv_page_num return resources.get_full_att_token_bytes() * (32 + args.page_size) + reserve, 0 monkeypatch.setattr(torch.cuda, "mem_get_info", check_runtime_before_profile) @@ -164,8 +164,8 @@ def check_runtime_before_profile(device: str | torch.device | int): assert not hasattr(model, "mtp_state_cursors") assert (main_linear_state(resources).mtp_state_index is not None) == bool(mtp) if mode == "prefill": - assert resources.get_pd_token_page_layouts() - assert resources.get_pd_runtime_page_layouts() + assert resources.get_pd_operator().get_pd_full_att_page_layouts() + assert resources.get_pd_operator().get_pd_bounded_att_page_layouts() if mtp: draft = drafts[0] assert draft.req_manager is model.req_manager diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 8c0dd81779..46c28789cf 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -145,14 +145,14 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( ) torch.cuda.synchronize() expected.append((state.big_pages_pool.conv_state[slot].clone(), state.big_pages_pool.ssm_state[slot].clone())) - size = sources[0][0].get_cpu_page_bytes() + size = sources[0][0].get_cpu_cache_operator().get_page_bytes() buffer = torch.zeros((2, size), dtype=torch.uint8, pin_memory=True) for resources, full in sources: - resources.write_req_to_cpu( + resources.get_cpu_cache_operator().write_req( token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=buffer, cpu_page_indices=[1], - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=[0], ) torch.cuda.synchronize() @@ -160,7 +160,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( if glm: resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(13) slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=buffer, cpu_page_indices=[1], token_indices=torch.arange( @@ -170,7 +170,7 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( device="cuda", ), state_token_len=8, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=[slot], req_idx=1, ) @@ -200,5 +200,5 @@ def test_cpu_layout_query_does_not_require_cpu_cache_enabled(monkeypatch: pytest rank=0, mtp=0, ) - assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_page_layouts().values()) + assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_cache_operator().get_page_layouts().values()) get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index 067063166c..e85b92aad2 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -136,11 +136,11 @@ def test_model_and_cpu_service_use_identical_storage_geometry( assert state.get_layer_view(logical_layer_id=0)[0].data_ptr() == state.get_layer_view(logical_layer_id=1)[0].data_ptr() assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) assert resources.full_att_allocator.size == 16 - layout = resources.get_cpu_page_layouts()[state.state_name] + layout = resources.get_cpu_cache_operator().get_page_layouts()[state.state_name] assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) # 按物理层数计算传输页；共享消费者不重复增加 KV 或传输容量。 expected_pd_bytes = storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() - assert resources.get_pd_page_bytes() == expected_pd_bytes + assert resources.get_pd_operator().get_page_bytes() == expected_pd_bytes fields = layout.regions assert all(region.shape[2] == global_heads for region in fields) assert not hasattr(model, "mem_manager") @@ -163,7 +163,7 @@ def test_gpu_profile_counts_token_fields_hold_and_transfer_reserve( pd_kv_page_size=8, pd_kv_page_num=3, ) - reserve = reference.get_pd_page_bytes() * args.pd_kv_page_num + reserve = reference.get_pd_operator().get_page_bytes() * args.pd_kv_page_num required = state.get_token_bytes() * (32 + args.page_size) # One byte below the full 32-token budget must drop a complete page. total = 1 << 30 @@ -197,7 +197,7 @@ def memory_info(device: str | torch.device | int): device="cpu", ) assert resources.full_att_allocator.size == 28 - assert resources.get_pd_page_bytes() == reference.get_pd_page_bytes() + assert resources.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() assert observed_devices == [torch.device("cuda", 3)] @@ -251,6 +251,6 @@ def create_full_att_states( model.tp_rank_, model.tp_world_size_, model.max_req_num = 0, 1, 1 prepare_model_states(model) model.att_state_manager.finish_state_registration() - runtime_layouts = model.att_state_manager.get_cpu_page_layouts() + runtime_layouts = model.att_state_manager.get_cpu_cache_operator().get_page_layouts() assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_state", "main_gqa_fp16_state"} assert cpu_layouts == runtime_layouts diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index 186565bb69..f81210c842 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -45,23 +45,23 @@ def test_main_and_draft_transfers_copy_pd_and_dp_windows( getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 2].fill_(19) main_linear_state(source).conv_state.fill_(3) main_linear_state(source).ssm_state.fill_(5) - page = torch.empty(source.get_pd_page_bytes(), dtype=torch.uint8, device="cuda") + page = torch.empty(source.get_pd_operator().get_page_bytes(), dtype=torch.uint8, device="cuda") for kind, getter in ( - ("kv", "get_pd_token_page_layouts"), - ("att_state", "get_pd_runtime_page_layouts"), + ("kv", "get_pd_full_att_page_layouts"), + ("att_state", "get_pd_bounded_att_page_layouts"), ): - source.write_to_pd_page( + source.get_pd_operator().write_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, - layouts=getattr(source, getter)(), + layouts=getattr(source.get_pd_operator(), getter)(), **({"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 0, "state_token_len": 8}), ) - target.read_from_pd_page( + target.get_pd_operator().read_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, - layouts=getattr(target, getter)(), + layouts=getattr(target.get_pd_operator(), getter)(), **({"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 1, "state_token_len": 8}), ) src_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") @@ -72,7 +72,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: with monkeypatch.context() as patch: patch.setattr(torch.Tensor, "tolist", reject_host_list) - target.copy_from_dp( + target.get_dp_shared_operator().copy_from_dp( remote_att_manager=source, src_token_indices=src_indices, dst_token_indices=dst_indices, @@ -102,17 +102,17 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: page_slot_ids=[slot], ) torch.cuda.synchronize() - cpu = torch.empty((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - source.write_req_to_cpu( + cpu = torch.empty((1, source.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.get_cpu_cache_operator().write_req( token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=cpu, cpu_page_indices=[0], - layouts=source.get_cpu_page_layouts(), + layouts=source.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=[slot], ) torch.cuda.synchronize() dst_slot = target.bounded_att_big_page_allocator.alloc(1)[0] - target.read_req_from_cpu( + target.get_cpu_cache_operator().read_req( cpu_page_buffers=cpu, cpu_page_indices=[0], token_indices=torch.arange( @@ -122,7 +122,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: device="cuda", ), state_token_len=8, - layouts=target.get_cpu_page_layouts(), + layouts=target.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=[dst_slot], req_idx=1, ) diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 1fd8f3e21e..5c34e35e28 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -102,29 +102,29 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: target_ids=[4, 6], ) - cpu_pages = torch.empty((1, managers[0].get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") + cpu_pages = torch.empty((1, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") if device == "cpu": with pytest.raises(AssertionError, match="CUDA State storage"): - managers[0].write_req_to_cpu( + managers[0].get_cpu_cache_operator().write_req( token_indices=source_indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], - layouts=managers[0].get_cpu_page_layouts(), + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), ) else: - managers[0].write_req_to_cpu( + managers[0].get_cpu_cache_operator().write_req( token_indices=source_indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], - layouts=managers[0].get_cpu_page_layouts(), + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), ) if device == "cuda": torch.cuda.synchronize() - managers[1].read_req_from_cpu( + managers[1].get_cpu_cache_operator().read_req( cpu_page_buffers=cpu_pages, cpu_page_indices=[0], token_indices=cpu_indices, - layouts=managers[1].get_cpu_page_layouts(), + layouts=managers[1].get_cpu_cache_operator().get_page_layouts(), ) assert_token_fields_equal( source=src, @@ -135,7 +135,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): - managers[1].copy_from_dp( + managers[1].get_dp_shared_operator().copy_from_dp( remote_att_manager=managers[0], src_token_indices=source_indices, dst_token_indices=dp_indices, @@ -145,7 +145,7 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: dst_req_idx=0, ) return - managers[1].copy_from_dp( + managers[1].get_dp_shared_operator().copy_from_dp( remote_att_manager=managers[0], src_token_indices=source_indices, dst_token_indices=dp_indices, @@ -468,8 +468,8 @@ def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, targe def test_ipc_readers_accept_different_pickle_lengths(kind: str, monkeypatch: pytest.MonkeyPatch) -> None: """各读者独立导出；模拟 IPC 计数变化引起的长度差异，并验证补齐后仍可逐份读取。""" resources = SimpleNamespace( - create_pd_ipc_copy=lambda: {"kind": "pd"}, - create_dp_ipc_copy=lambda: {"kind": "dp"}, + get_pd_operator=lambda: SimpleNamespace(create_ipc_copy=lambda: {"kind": "pd"}), + get_dp_shared_operator=lambda: SimpleNamespace(create_ipc_copy=lambda: {"kind": "dp"}), ) mapping = [[0, 1, 2]] payloads = [] @@ -522,8 +522,8 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: for state in linear_states(resources): assert state.big_pages_pool.conv_state.is_shared() assert not state.big_pages_pool.conv_state.is_pinned() - resources.register_page_pools_pin_memory() - resources.register_page_pools_pin_memory() + resources.get_dp_shared_operator().register_page_pools_pin_memory() + resources.get_dp_shared_operator().register_page_pools_pin_memory() for state in linear_states(resources): for pool in (state.big_pages_pool, state.small_pages_pool): if kind == "dp": diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 787d3ef04f..5b383bba4e 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -7,6 +7,7 @@ from unit_tests.common.att_resources.model_helpers import read_kv_storage from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from typing import TYPE_CHECKING from dataclasses import replace @@ -219,25 +220,25 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp target.register_state("draft0", target_draft) target.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) target.finish_state_registration() - page = torch.zeros((1, source.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) - source.write_req_to_cpu( + page = torch.zeros((1, source.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.get_cpu_cache_operator().write_req( token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), cpu_page_buffers=page, cpu_page_indices=[0], - layouts=source.get_cpu_page_layouts(), + layouts=source.get_cpu_cache_operator().get_page_layouts(), ) - target.read_req_from_cpu( + target.get_cpu_cache_operator().read_req( cpu_page_buffers=page, cpu_page_indices=[0], token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), - layouts=source.get_cpu_page_layouts(), + layouts=source.get_cpu_cache_operator().get_page_layouts(), ) for src, dst in ((main, target_main), (draft, target_draft)): fields = {name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items()} for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): if "token" in field.axes: assert torch.equal(field.tensor[:, :2], fields[name][:, 6:8]) - assert set(source.get_cpu_page_layouts()) == {main.state_name, draft.state_name} + assert set(source.get_cpu_cache_operator().get_page_layouts()) == {main.state_name, draft.state_name} @pytest.mark.parametrize("overlap", [False, True]) @@ -249,7 +250,7 @@ def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_la manager.register_state("draft0", draft) manager.full_att_allocator = FullAttKVAllocator(main.capacity, main.page_size, pin_memory=False) template = allocated_state("int8") - quant = rebuild_full_state(template, config=replace(template.config, logical_layer_groups=((4,), (6,)))) + quant = rebuild_full_state(template, config=replace_logical_layer_groups(template.config, ((4,), (6,)))) manager.register_state("main", quant) if overlap: overlapping = allocated_state("static_head") @@ -285,9 +286,9 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): assert (field.tensor[:, 4:6] == 5).all() - layouts = manager.get_cpu_page_layouts() + layouts = manager.get_cpu_cache_operator().get_page_layouts() page = torch.zeros((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) - manager.write_req_to_cpu( + manager.get_cpu_cache_operator().write_req( token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), cpu_page_buffers=page, cpu_page_indices=[0], @@ -296,7 +297,7 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): field.tensor[:, 6:8].zero_() - manager.read_req_from_cpu( + manager.get_cpu_cache_operator().read_req( cpu_page_buffers=page, cpu_page_indices=[0], token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 33c1098a94..7424d5be90 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -101,7 +101,7 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): source_states = source.get_bounded_attention_states() owner_pools = [(state.big_pages_pool, state.small_pages_pool) for state in source_states] - exported = source.create_pd_ipc_copy() + exported = source.get_pd_operator().create_ipc_copy() assert type(exported) is type(source) and exported is not source for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): @@ -131,16 +131,16 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share_pools: bool) -> None: cpu_args.enable_dp_prompt_cache_fetch = share_pools source = MainDraftTestResourceManager() - pd_copy = source.create_pd_ipc_copy() + pd_copy = source.get_pd_operator().create_ipc_copy() for state in source.get_bounded_attention_states(): exported = getattr(pd_copy, state.state_name) assert exported.big_pages_pool is None and exported.small_pages_pool is None assert state.big_pages_pool is not None and state.small_pages_pool is not None if not share_pools: with pytest.raises(AssertionError, match="shared CPU page pools"): - source.create_dp_ipc_copy() + source.get_dp_shared_operator().create_ipc_copy() return - dp_copy = source.create_dp_ipc_copy() + dp_copy = source.get_dp_shared_operator().create_ipc_copy() for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): assert getattr(dp_copy, name) is None assert getattr(source, name) is not None @@ -155,7 +155,7 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): source = MainDraftTestResourceManager() source.get_full_attention_states() source.get_bounded_attention_states() - imported = pickle.loads(pickle.dumps(source.create_pd_ipc_copy())) + imported = pickle.loads(pickle.dumps(source.get_pd_operator().create_ipc_copy())) assert imported._state_registration_finished assert imported.full_att_allocator is None assert imported.bounded_att_big_page_allocator is None diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index 1db1627b07..aa26a9252f 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -7,6 +7,7 @@ from unit_tests.common.att_resources.model_helpers import full_state_test_buffers from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from typing import TYPE_CHECKING import pickle @@ -35,7 +36,7 @@ def cpu_layout(state: FullAttentionState | BoundedAttentionState): manager = AttentionStateManager(request_capacity=1) manager.register_state("main", state) manager.finish_state_registration() - return manager.get_cpu_page_layouts()[state.state_name] + return manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] @pytest.mark.parametrize( @@ -190,22 +191,22 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( assert query("draft1") == (other,) assert query() == (source, other) - page_bytes = manager.get_pd_page_bytes() + page_bytes = manager.get_pd_operator().get_page_bytes() assert not hasattr(manager, "pd_page_bytes") assert not hasattr(manager, "pd_transfer_layouts") if not bounded: # 两个 draft 共 6 个全局 KV head：K/V、head_dim=4、FP32，无额外页头。 assert page_bytes == page_tokens * 192 - assert manager.get_pd_page_bytes() == page_bytes + assert manager.get_pd_operator().get_page_bytes() == page_bytes getter = ( - manager.get_pd_runtime_page_layouts if bounded else manager.get_pd_token_page_layouts + manager.get_pd_operator().get_pd_bounded_att_page_layouts if bounded else manager.get_pd_operator().get_pd_full_att_page_layouts ) layouts = getter() assert getter() is layouts assert set(layouts) == {source.state_name, other.state_name} assert all(layout.tokens_per_page == (None if bounded else page_tokens) for layout in layouts.values()) unused_getter = ( - manager.get_pd_token_page_layouts if bounded else manager.get_pd_runtime_page_layouts + manager.get_pd_operator().get_pd_full_att_page_layouts if bounded else manager.get_pd_operator().get_pd_bounded_att_page_layouts ) assert unused_getter() == {} - assert manager.get_pd_page_bytes() == page_bytes + assert manager.get_pd_operator().get_page_bytes() == page_bytes diff --git a/unit_tests/common/att_resources/test_window_compressor_state.py b/unit_tests/common/att_resources/test_window_compressor_state.py index 8aa3c6cf06..ad907852a6 100644 --- a/unit_tests/common/att_resources/test_window_compressor_state.py +++ b/unit_tests/common/att_resources/test_window_compressor_state.py @@ -143,12 +143,12 @@ def test_manager_allocates_page_tensors_only_for_snapshot_states( if window_size == compress_stride: initialize_pools.assert_not_called() assert state.big_pages_pool is state.small_pages_pool is None - assert state.state_name not in manager.get_cpu_page_layouts() + assert state.state_name not in manager.get_cpu_cache_operator().get_page_layouts() else: initialize_pools.assert_called_once() assert state.big_pages_pool.buffer.shape == (4, 2, window_size - compress_stride, 4) assert state.small_pages_pool.buffer.shape == (3, 2, window_size - compress_stride, 4) - assert state.state_name in manager.get_cpu_page_layouts() + assert state.state_name in manager.get_cpu_cache_operator().get_page_layouts() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index c814494d6e..a861c945fb 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -13,6 +13,7 @@ from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups from unit_tests.common.att_resources.test_capacity import definition @@ -31,7 +32,7 @@ def test_fp8_backend_selects_calibration_from_logical_layer( template = definition("static_head" if backend_name == "fa3" else "static_tensor") state = rebuild_full_state( template=template, - config=replace(template.config, logical_layer_groups=((3, 7), (8, 11))), + config=replace_logical_layer_groups(template.config, ((3, 7), (8, 11))), token_capacity=8, device="cpu", ) diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py index 18311d021d..a38e67bb5f 100644 --- a/unit_tests/common/dsv4_test_utils.py +++ b/unit_tests/common/dsv4_test_utils.py @@ -273,17 +273,17 @@ def alloc_paged_kv_move_buffer(self, count: int, page_size: int): setattr(self.manager, state.state_name, state) self.manager._get_states.cache_clear() get_env_start_args().pd_kv_page_size = page_size - self.manager.get_pd_token_page_layouts.cache_clear() - self.manager.get_pd_runtime_page_layouts.cache_clear() - self.manager.get_pd_page_bytes.cache_clear() - self.kv_move_buffer = torch.empty((count, self.manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + self.manager.get_pd_operator().get_pd_full_att_page_layouts.cache_clear() + self.manager.get_pd_operator().get_pd_bounded_att_page_layouts.cache_clear() + self.manager.get_pd_operator().get_page_bytes.cache_clear() + self.kv_move_buffer = torch.empty((count, self.manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") return self.kv_move_buffer def get_pd_kv_move_buffer_size(self): - return self.manager.get_pd_page_bytes() * get_env_start_args().pd_kv_page_num + return self.manager.get_pd_operator().get_page_bytes() * get_env_start_args().pd_kv_page_num def get_paged_kv_move_buffer_shape(self, count: int, page_size: int): - return (count, self.manager.get_pd_page_bytes()) + return (count, self.manager.get_pd_operator().get_page_bytes()) def write_mem_to_page_kv_move_buffer( self, @@ -296,13 +296,13 @@ def write_mem_to_page_kv_move_buffer( req_idx: int | None = None, ): for manager in managers[dp_index * tp_size : dp_index * tp_size + 1]: - manager.manager.write_to_pd_page( + manager.manager.get_pd_operator().write_page( pd_page_buffers=self.kv_move_buffer, pd_page_index=page, kind=kind, - layouts=manager.manager.get_pd_token_page_layouts() + layouts=manager.manager.get_pd_operator().get_pd_full_att_page_layouts() if kind == "kv" - else manager.manager.get_pd_runtime_page_layouts(), + else manager.manager.get_pd_operator().get_pd_bounded_att_page_layouts(), token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), req_idx=req_idx, state_token_len=getattr(manager, "state_token_len", 0), @@ -319,13 +319,13 @@ def read_page_kv_move_buffer_to_mem( req_idx: int | None = None, ): for manager in managers[dp_index * tp_size : (dp_index + 1) * tp_size]: - manager.manager.read_from_pd_page( + manager.manager.get_pd_operator().read_page( pd_page_buffers=self.kv_move_buffer, pd_page_index=page, kind=kind, - layouts=manager.manager.get_pd_token_page_layouts() + layouts=manager.manager.get_pd_operator().get_pd_full_att_page_layouts() if kind == "kv" - else manager.manager.get_pd_runtime_page_layouts(), + else manager.manager.get_pd_operator().get_pd_bounded_att_page_layouts(), token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), req_idx=req_idx, state_token_len=getattr(manager, "state_token_len", 0), diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index a4549dc078..983198f835 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -135,8 +135,8 @@ def test_cache_io_copies_individual_token_slots(cache): for buffer in manager.kv_buffer: buffer.zero_() buffer[source] = torch.arange(1, 5, device="cuda", dtype=torch.uint8)[:, None, None] - layouts = manager.manager.get_cpu_page_layouts() - pages = torch.empty((1, manager.manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + layouts = manager.manager.get_cpu_cache_operator().get_page_layouts() + pages = torch.empty((1, manager.manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) for state in manager.full_states.values(): state.write_to_cpu(source, pages, [0], layout=layouts[state.state_name]) state.read_from_cpu(pages, [0], destination, layout=layouts[state.state_name]) @@ -1743,17 +1743,17 @@ def test_cpu_cache_preserves_already_published_pages(cache): state_id = manager.big_page_buffers.alloc_one_state_cache() manager.big_page_buffers.buffer[state_id].fill_(7) checkpoint_ids = [state_id] - manager.manager.write_req_to_cpu( + manager.manager.get_cpu_cache_operator().write_req( token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_page_layouts(), cpu_page_ready_flags=readies, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), cpu_page_ready_flags=readies, big_page_slot_ids=checkpoint_ids, ) torch.cuda.synchronize() assert pages.eq(221).all() readies[0] = False - manager.manager.write_req_to_cpu( + manager.manager.get_cpu_cache_operator().write_req( token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_page_layouts(), cpu_page_ready_flags=readies, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), cpu_page_ready_flags=readies, big_page_slot_ids=checkpoint_ids, ) torch.cuda.synchronize() diff --git a/unit_tests/models/deepseek_v4/test_vision_integration.py b/unit_tests/models/deepseek_v4/test_vision_integration.py index 9ffd9f3962..8290ab7d89 100644 --- a/unit_tests/models/deepseek_v4/test_vision_integration.py +++ b/unit_tests/models/deepseek_v4/test_vision_integration.py @@ -281,8 +281,10 @@ def allocate(req, count): model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), att_state_manager=SimpleNamespace( - get_cpu_page_layouts=lambda: {}, - read_req_from_cpu=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), + ), ), ), ) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 6e9241a55b..86eebff928 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -1283,6 +1283,8 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.cpu_cache_operator.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.pd_operator.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) config = PackedMLALogicalConfig( logical_layer_groups=((3,),), latent_dim=512, format="packed_group64_e8m0", @@ -1300,12 +1302,12 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): write_fp8_e8m0_packed_kv(kv, indices, expected) torch.testing.assert_close(state.get_layer_view(logical_layer_id=3).kv[:4], expected, rtol=0, atol=0) state.get_layer_view(logical_layer_id=3).indexer_k[:4].random_(0, 256) - cpu_layout = manager.get_cpu_page_layouts()[state.state_name] - cpu_pages = torch.empty((1, manager.get_cpu_page_bytes()), dtype=torch.uint8, pin_memory=True) + cpu_layout = manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] + cpu_pages = torch.empty((1, manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) state.write_to_cpu(token_indices=indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], layout=cpu_layout) state.read_from_cpu(cpu_page_buffers=cpu_pages, cpu_page_indices=[0], token_indices=indices + 4, layout=cpu_layout) - pd_layout = manager.get_pd_token_page_layouts()[state.state_name] - pd_pages = torch.empty((1, manager.get_pd_page_bytes()), dtype=torch.uint8, device="cuda") + pd_layout = manager.get_pd_operator().get_pd_full_att_page_layouts()[state.state_name] + pd_pages = torch.empty((1, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") state.write_to_pd_page(token_indices=indices, pd_page_buffers=pd_pages, pd_page_index=0, layout=pd_layout) state.read_from_pd_page(pd_page_buffers=pd_pages, token_indices=indices + 8, pd_page_index=0, layout=pd_layout) torch.cuda.synchronize() diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index b163d0d972..7ed6fb487b 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -191,7 +191,7 @@ def test_pd_kv_and_runtime_state_roundtrip( source_buffer = allocate_transfer_buffer(source_resources[0], 1, device="cuda") dest_buffer = allocate_transfer_buffer(dest_resources[0], 1, device="cuda") assert source_buffer.shape == dest_buffer.shape - assert source_resources[0].get_pd_runtime_page_layouts().keys() == dest_resources[0].get_pd_runtime_page_layouts().keys() + assert source_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() == dest_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() for kind in ("kv", "att_state"): copy_page( node_att_managers=source_resources, @@ -284,12 +284,12 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_token_page_layouts().values()) + full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_full_att_page_layouts().values()) runtime_bytes = max( - layout.byte_end for layout in mem.resources.get_pd_runtime_page_layouts().values() + layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_bounded_att_page_layouts().values() ) assert runtime_bytes > full_bytes - assert mem.resources.get_pd_page_bytes() == runtime_bytes + assert mem.resources.get_pd_operator().get_page_bytes() == runtime_bytes assert "smaller than the complete attention runtime" in caplog.text diff --git a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py index fd75b23e95..df68cc14bd 100644 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py @@ -80,9 +80,11 @@ def allocate(req, count): model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(1024).view(1, -1)), att_state_manager=SimpleNamespace( - get_cpu_page_layouts=lambda: {}, - read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: copied_pages.extend( - cpu_page_indices + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: copied_pages.extend( + cpu_page_indices + ), ), ), ), diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py index 2a90c7f943..2b30c127b2 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -128,7 +128,7 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( barriers = [] manager = SimpleNamespace( bounded_att_big_page_allocator=SimpleNamespace(alloc=allocate_pages), - copy_from_dp=copy_resources, + get_dp_shared_operator=lambda: SimpleNamespace(copy_from_dp=copy_resources), ) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index 5f71ea7dad..aba9ff3906 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -182,9 +182,11 @@ def alloc_req_kv_mem(req, alloc_token_num): return mem_indexes resources = SimpleNamespace( - get_cpu_page_layouts=lambda: {}, - read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded_indexes.extend( - token_indices.tolist() + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded_indexes.extend( + token_indices.tolist() + ), ), ) backend = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py index bef86b47b7..1da16313fd 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -174,7 +174,9 @@ def test_dp_imports_only_matching_tp_shards( tables = {} for rank in expected_ranks: managers[rank] = SimpleNamespace( - register_page_pools_pin_memory=lambda rank=rank: registered_ranks.append(rank), + get_dp_shared_operator=lambda rank=rank: SimpleNamespace( + register_page_pools_pin_memory=lambda: registered_ranks.append(rank), + ), ) tables[rank] = object() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py index a6f9c623e3..6d7018f561 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -21,7 +21,7 @@ def test_pd_uploads_cpu_indices_on_copy_stream_and_skips_runtime_pages( module_class = implementation._PrefillTransModule if prefill else implementation._DecodeTransModule device_id = torch.cuda.current_device() args = SimpleNamespace(tp=2, dp=2, nnodes=1, pd_kv_page_num=2, pd_kv_page_size=4, pd_node_id=1) - resource = SimpleNamespace(get_pd_page_bytes=lambda: 64) + resource = SimpleNamespace(get_pd_operator=lambda: SimpleNamespace(get_page_bytes=lambda: 64)) uploads = [] copies = [] original_cuda = torch.Tensor.cuda diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index ff1c29add4..e0c429419e 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -225,10 +225,10 @@ def test_cpu_service_describes_main_and_draft_as_distinct_resources( layer for group in draft.config.logical_layer_groups for layer in group ) assert main.capacity == draft.capacity == resources.full_att_allocator.size - assert {main.state_name, draft.state_name} <= resources.get_cpu_page_layouts().keys() - layout_end = max(layout.byte_end for layout in resources.get_cpu_page_layouts().values()) + assert {main.state_name, draft.state_name} <= resources.get_cpu_cache_operator().get_page_layouts().keys() + layout_end = max(layout.byte_end for layout in resources.get_cpu_cache_operator().get_page_layouts().values()) assert (layout_end + 15) // 16 * 16 == meta.page_bytes - assert resources.get_cpu_page_bytes() == meta.page_bytes + assert resources.get_cpu_cache_operator().get_page_bytes() == meta.page_bytes @pytest.mark.parametrize("kind,encoding,model_cls", CASES) @@ -254,9 +254,9 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( # No byte-size formulas duplicated in the test: check the service's actual # meta-device descriptors against every real model's layout. for resources in managers: - layout = next(iter(resources.get_cpu_page_layouts().values())) + layout = next(iter(resources.get_cpu_cache_operator().get_page_layouts().values())) assert 0 <= layout.byte_start < layout.byte_end <= meta.page_bytes - assert meta.page_bytes == resources.get_cpu_page_bytes() + assert meta.page_bytes == resources.get_cpu_cache_operator().get_page_bytes() meta.page_num = 3 client = cpu_cache_factory(meta) buffer = client.cpu_kv_cache_tensor @@ -287,11 +287,11 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( if "token" in f.axes } ) - resources.write_req_to_cpu( + resources.get_cpu_cache_operator().write_req( token_indices=torch.arange(8, dtype=torch.long, device="cuda"), cpu_page_buffers=buffer, cpu_page_indices=[2], - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), **kwargs, ) torch.cuda.synchronize() @@ -306,7 +306,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( big_page_slot_ids=[slot], req_idx=1, ) - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=buffer, cpu_page_indices=[2], token_indices=torch.arange( @@ -316,7 +316,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( device="cuda", ), state_token_len=8, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), **kwargs, ) if resources.bounded_att_big_page_allocator is not None: @@ -330,7 +330,7 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( @pytest.mark.parametrize("hybrid", [False, True]) -def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( +def test_production_offload_preserves_ready_pages_without_host_sync( runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -357,13 +357,13 @@ def test_production_offload_uses_stream_dependency_and_preserves_ready_pages( client.update_pages_status_to_ready([pages[0]], deref=False) slots = resources.bounded_att_big_page_allocator.alloc(2) if hybrid else None write_batches = [] - write_to_cpu = resources.write_req_to_cpu + write_to_cpu = resources.get_cpu_cache_operator().write_req def capture_write(*args: object, **kwargs: object) -> None: write_batches.append((kwargs["cpu_page_indices"], kwargs["cpu_page_ready_flags"])) write_to_cpu(*args, **kwargs) - monkeypatch.setattr(resources, "write_req_to_cpu", capture_write) + monkeypatch.setattr(resources.get_cpu_cache_operator(), "write_req", capture_write) if hybrid: # 当前 bounded 页仍由主机复制，测试先准备好源快照。 # TODO：CPU→CPU kernel 落地后，再覆盖尚未完成 D2H 的快照依赖。 @@ -374,7 +374,8 @@ def capture_write(*args: object, **kwargs: object) -> None: producer = torch.cuda.Stream() producer.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(producer): - # 延迟当前生产流的 KV 写入，验证搬运流通过 wait_stream 等待。 + # 延迟当前生产流的 KV 写入；当前实现不在此处隐式插入 wait_stream， + # 生产者与搬运流的依赖由调用方负责组织。 torch.cuda._sleep(10_000_000) (full,) = resources.get_full_attention_states("main") for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): @@ -425,7 +426,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: patch.setattr(torch.cuda.Stream, "synchronize", reject_host_sync) patch.setattr(torch.cuda.Stream, "wait_stream", record_wait) task = module.write_req_to_cpu(req, copy_stream) - assert waited_streams == [producer] + assert waited_streams == [] task.sync_event.synchronize() assert write_batches == [(pages, [True, False])] assert req.cpu_cache_task_status.is_running() @@ -439,12 +440,12 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: resources.bounded_att_big_page_allocator.free(slots) destination = resources.bounded_att_big_page_allocator.alloc(1)[0] kwargs = dict(big_page_slot_ids=[destination], req_idx=1) - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=[pages[1]], token_indices=torch.arange(8, dtype=torch.long, device="cuda"), state_token_len=8, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), **kwargs, ) for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): @@ -480,11 +481,11 @@ def test_disk_preserves_resource_pages( (full,) = resources.get_full_attention_states("main") for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.random_(0, 256) - resources.write_req_to_cpu( + resources.get_cpu_cache_operator().write_req( token_indices=torch.arange(16, dtype=torch.long, device="cuda"), cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=pages, - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), ) torch.cuda.synchronize() expected = client.cpu_kv_cache_tensor.clone() @@ -500,10 +501,10 @@ def test_disk_preserves_resource_pages( assert torch.equal(client.cpu_kv_cache_tensor, expected) for storage_buffer in full_state_test_buffers(full).values(): storage_buffer.zero_() - resources.read_req_from_cpu( + resources.get_cpu_cache_operator().read_req( cpu_page_buffers=client.cpu_kv_cache_tensor, cpu_page_indices=pages, token_indices=torch.arange(16, dtype=torch.long, device="cuda"), - layouts=resources.get_cpu_page_layouts(), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), ) assert any((tmp_path / "disk").glob("cache_file*")) diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 291edfddc2..d4861792c2 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -140,7 +140,7 @@ def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( expected_k = full.k_storage[:, :token_num].clone() expected_v = full.v_storage[:, :token_num].clone() client = cpu_cache_factory( - CpuKVCacheMeta(3, resources.get_cpu_page_bytes()), + CpuKVCacheMeta(3, resources.get_cpu_cache_operator().get_page_bytes()), ) indexes = torch.arange(token_num, dtype=torch.int32, device="cuda") req = SimpleNamespace( @@ -313,7 +313,7 @@ def write_to_cpu( lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), allocate_pages=allocate_pages, cpu_kv_cache_tensor=torch.empty((2, 1), dtype=torch.uint8), ) - monkeypatch.setattr(resources, "write_req_to_cpu", write_to_cpu) + monkeypatch.setattr(resources.get_cpu_cache_operator(), "write_req", write_to_cpu) broadcasts = [] monkeypatch.setattr(common.dist, "broadcast_object_list", lambda data, **kwargs: broadcasts.append(data)) task = module.write_req_to_cpu(req, torch.cuda.Stream()) diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 24714e9b71..f24d3dcdae 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -331,11 +331,13 @@ def allocate(req, count): model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), att_state_manager=SimpleNamespace( - get_cpu_page_layouts=lambda: {}, - bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), - read_req_from_cpu=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded.append( - req.hold_kv_len + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded.append( + req.hold_kv_len + ), ), + bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), ), ), ) From 227d8ac6a9a4ed07f50e6db26b51edc437656bdc Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 2 Nov 2026 06:56:28 +0000 Subject: [PATCH 48/56] refactor: extract attention state interfaces --- .../attention_state/bounded_att/base_state.py | 220 +---------------- .../attention_state/bounded_att/interface.py | 232 ++++++++++++++++++ .../attention_state/full_att/base_state.py | 184 +------------- .../attention_state/full_att/interface.py | 197 +++++++++++++++ 4 files changed, 433 insertions(+), 400 deletions(-) create mode 100644 lightllm/common/attention_state/bounded_att/interface.py create mode 100644 lightllm/common/attention_state/full_att/interface.py diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 9796d49c0a..6a96e76dc9 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -2,224 +2,8 @@ from __future__ import annotations -from functools import lru_cache +from .interface import CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface -from typing import Any -from abc import ABC, abstractmethod -from typing import List - -import torch - -from ..utils import deepcopy_result -from lightllm.common.attention_state.types import ( - StateTransferLayout, - PageKind, - TransferRegion, -) - - -class BoundedAttentionState(ABC): +class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - @abstractmethod - def get_state_type(self) -> str: - """返回算法与存储方案标记，用于生成资源成员名。 - - 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， - 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 - 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 - 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 - 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 - - KDA/GDN 格式为 {kda|gdn}_conv_{Conv精度}_ssm_{SSM精度}。 - 两份状态精度分别标明，不能只写一个 dtype，例如 kda_conv_bf16_ssm_fp32， - 对应 draft0_kda_conv_bf16_ssm_fp32_state。不同精度方案可以在同一角色下共存。 - 其他 bounded 算法使用自己的算法名及方案参数，不复用 kda/gdn 标记。 - """ - ... - - @abstractmethod - def initialize_page_pools( - self, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, - ) -> None: - """容量规划完成后创建大小页池；容量为零表示不启用该页池。""" - ... - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - - @abstractmethod - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" - ... - - @abstractmethod - def get_layer_view(self, logical_layer_id: int) -> Any: - """按逻辑层号返回运行态和算法控制信息的计算视图。""" - ... - - @abstractmethod - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" - ... - - @abstractmethod - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 - - 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 - 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 - """ - ... - - @abstractmethod - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 - - req_indices、state_token_lens、page_slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， - 不表示本批搬运的 token 数，调用方保证与所选快照的前缀一致。 - """ - ... - - @abstractmethod - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" - ... - - # ======================================================================== - # 三、CPU Cache 数据交换：读写大小页快照 - # ======================================================================== - - @abstractmethod - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回一份 CPU 快照的字段区间、全局形状与 TP 切片。 - - byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； - 参与 CPU 页时，管理器按字段最大终点计算占用，不乘页内 token 数。 - 在合法缓存边界无需保留历史的 State 返回空元组，不占用 CPU 页面。 - 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 - 管理器分配后统一清理，子类覆盖时须保留缓存和副本语义。 - """ - ... - - @abstractmethod - def write_state_pages_to_cpu( - self, - page_kind: PageKind, - page_slot_ids: List[int], - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_ready_flags 中为 True 的页面。""" - ... - - @abstractmethod - def read_state_pages_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], - layout: StateTransferLayout, - ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" - ... - - # ======================================================================== - # 四、DP 迁移：大页快照与基准运行态 - # ======================================================================== - - @abstractmethod - def copy_big_pages_from_dp( - self, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state: BoundedAttentionState, - ) -> None: - """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" - ... - - @abstractmethod - def copy_runtime_state_from_dp( - self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: BoundedAttentionState, - ) -> None: - """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 - - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ - ... - - # ======================================================================== - # 五、PD 分离传输：读写基准运行态 - # ======================================================================== - - @abstractmethod - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 - - 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 - 管理器分配存储后统一清理，子类覆盖时须保留缓存和副本语义。 - """ - ... - - @abstractmethod - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 - - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ - ... - - @abstractmethod - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """按 state_token_len 恢复请求运行态及 ring 位置等控制信息，不写入大小页池。""" - ... diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py new file mode 100644 index 0000000000..fe1fceb3bd --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -0,0 +1,232 @@ +"""有界 attention 状态的功能接口；具体算法实现运行态与大小页存储。""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from functools import lru_cache +from typing import Any, List + +import torch + +from ..utils import deepcopy_result +from lightllm.common.attention_state.types import PageKind, StateTransferLayout, TransferRegion + + +class CoreInterface(ABC): + """有界 attention State 的资源、层视图和运行态核心能力。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + @abstractmethod + def get_state_type(self) -> str: + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + KDA/GDN 格式为 {kda|gdn}_conv_{Conv精度}_ssm_{SSM精度}。 + 两份状态精度分别标明，不能只写一个 dtype，例如 kda_conv_bf16_ssm_fp32， + 对应 draft0_kda_conv_bf16_ssm_fp32_state。不同精度方案可以在同一角色下共存。 + 其他 bounded 算法使用自己的算法名及方案参数，不复用 kda/gdn 标记。 + """ + ... + + @abstractmethod + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建大小页池；容量为零表示不启用该页池。""" + ... + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + @abstractmethod + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" + ... + + @abstractmethod + def get_layer_view(self, logical_layer_id: int) -> Any: + """按逻辑层号返回运行态和算法控制信息的计算视图。""" + ... + + @abstractmethod + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" + ... + + @abstractmethod + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 + + 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 + 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 + """ + ... + + @abstractmethod + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 + + req_indices、state_token_lens、page_slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， + 不表示本批搬运的 token 数，调用方保证与所选快照的前缀一致。 + """ + ... + + @abstractmethod + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + ... + + +class CpuCacheInterface(ABC): + # ======================================================================== + # 三、CPU Cache 数据交换：读写大小页快照 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回一份 CPU 快照的字段区间、全局形状与 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 参与 CPU 页时，管理器按字段最大终点计算占用，不乘页内 token 数。 + 在合法缓存边界无需保留历史的 State 返回空元组，不占用 CPU 页面。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 管理器分配后统一清理，子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_ready_flags 中为 True 的页面。""" + ... + + @abstractmethod + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + ... + + +class DpCopyInterface(ABC): + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + @abstractmethod + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: CoreInterface, + ) -> None: + """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" + ... + + @abstractmethod + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: CoreInterface, + ) -> None: + """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + +class PdInterface(ABC): + # ======================================================================== + # 五、PD 分离传输：读写基准运行态 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 管理器分配存储后统一清理，子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + @abstractmethod + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按 state_token_len 恢复请求运行态及 ring 位置等控制信息，不写入大小页池。""" + ... + + +__all__ = [ + "CoreInterface", + "CpuCacheInterface", + "DpCopyInterface", + "PdInterface", +] diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index e591927ac9..d3f00430b0 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -1,192 +1,12 @@ from __future__ import annotations -from functools import lru_cache -from typing import List -from abc import ABC, abstractmethod -from typing import Any +from .interface import CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface -import torch -from ..utils import deepcopy_result -from lightllm.common.attention_state.types import ( - StateTransferLayout, - TransferRegion, -) - - -class FullAttentionState(ABC): +class FullAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 所有表示 token 索引的参数都须为调用方准备的 CUDA Tensor，位于对应存储设备上； State 不转换设备、dtype 或连续布局，int32/int64 由具体算子支持，DP 目标索引须为 int64。 CPU 页位于主机内存不改变 token 索引的 CUDA 契约。 """ - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - # 具体存储和量化方案标识，用于 main/draft State 成员命名。 - @abstractmethod - def get_state_type(self) -> str: - """返回算法与存储方案标记，用于生成资源成员名。 - - 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， - 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 - 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 - 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 - 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 - 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 - - - GQA 浮点 KV：gqa_{bf16|fp16|fp32}。 - - GQA 分组量化：gqa_int{4|8}_group{分组大小}_{scale精度}_scale。 - - GQA 静态 FP8：gqa_fp8_static_{per_tensor|per_head}。 - - MLA 浮点：mla_{精度}；KV/indexer 联合存储：mla_packed_{精度}。 - - MLA packed FP8：mla_packed_fp8_group128，或 mla_packed_fp8_group64_e8m0。 - - 序列压缩历史：compressed_c{压缩比例}_fp8_group64_e8m0。 - 例如 gqa_int8_group128_bf16_scale 对应 main_gqa_int8_group128_bf16_scale_state。 - 新量化实现必须区分影响方案选择的位宽、分组、scale 方式及精度，不编码标定数值。 - """ - ... - - @abstractmethod - def allocate_storage(self, token_capacity: int) -> None: - """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" - raise NotImplementedError() - - @abstractmethod - def get_token_bytes(self) -> int: - """累计当前 TP rank 各物理层的单 token 字节数，包含动态量化参数和行 padding；共享逻辑层不重复计数。""" - raise NotImplementedError() - - # ======================================================================== - # 二、Full attention：计算视图、KV 写入与完整 token 复制 - # ======================================================================== - - @abstractmethod - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" - ... - - @abstractmethod - def get_storage_layer_index(self, logical_layer_id: int) -> int: - """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" - ... - - @abstractmethod - def get_layer_view(self, logical_layer_id: int) -> Any: - """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" - raise NotImplementedError() - - @abstractmethod - def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: - """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。 - - 融合压缩等专用写入路径可直接使用计算视图，不调用本接口。 - """ - raise NotImplementedError() - - @abstractmethod - def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: - """复制全部所辖层的完整 token 数据，包含动态 scale、indexer 字段，保持原存储表示。 - - 源和目标索引为存储设备上的等长一维 CUDA 整数 Tensor，位置不重叠，直接使用原 dtype 和 stride。""" - raise NotImplementedError() - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - - @abstractmethod - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """按全局 cpu_cache_token_page_size 返回 CPU 字段、全局形状和 TP 切片。 - - byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； - 管理器按字段最大终点计算该 State 的 CPU 页占用。 - 本次有效 token 数由读写流程处理，不改变布局。 - 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 - 首次查询后配置与全局 CPU 页规格保持不变。 - 管理器实际分配存储后统一清除缓存，以兼容物理页规格变化；子类覆盖时须保留缓存能力。 - """ - ... - - @abstractmethod - def write_to_cpu( - self, - token_indices: torch.Tensor, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """按布局写出 token 数据及配套字段；cpu_page_ready_flags 中为 True 的页跳过写入。""" - raise NotImplementedError() - - @abstractmethod - def read_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - token_indices: torch.Tensor, - layout: StateTransferLayout, - ) -> None: - """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" - raise NotImplementedError() - - # ======================================================================== - # 四、DP 迁移 - # ======================================================================== - - @abstractmethod - def copy_tokens_from_dp( - self, - src_token_indices: torch.Tensor, - dst_token_indices: torch.Tensor, - remote_dp_full_att_state: FullAttentionState, - ) -> None: - """复制其他 DP 同一 TP 分片的完整 token 数据，包含量化参数及 indexer 字段。 - - 源、目标索引须已位于各自存储的 CUDA 设备，可为 int32/int64；接口不转换设备或 dtype。 - 调用方保证 peer/IPC 可访问、源数据就绪、目标索引唯一，且源目标数据区域不重叠。 - """ - raise NotImplementedError() - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - - @abstractmethod - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 - - 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 - 首次查询后配置与全局 PD 页规格保持不变；管理器分配存储后统一清除缓存。 - 子类覆盖时须保留缓存和副本语义。 - """ - ... - - @abstractmethod - def write_to_pd_page( - self, - token_indices: torch.Tensor, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - layout: StateTransferLayout, - ) -> None: - """将 token 数据写入PD 页集合中的指定发送页，网络发送与同步由调用方负责。""" - raise NotImplementedError() - - @abstractmethod - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - token_indices: torch.Tensor, - pd_page_index: int, - layout: StateTransferLayout, - ) -> None: - """从PD 页集合中的指定接收页还原 token 数据，不分配索引或修改请求映射。""" - raise NotImplementedError() diff --git a/lightllm/common/attention_state/full_att/interface.py b/lightllm/common/attention_state/full_att/interface.py new file mode 100644 index 0000000000..7240bf9f79 --- /dev/null +++ b/lightllm/common/attention_state/full_att/interface.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from functools import lru_cache +from typing import Any, List + +import torch + +from ..utils import deepcopy_result +from lightllm.common.attention_state.types import StateTransferLayout, TransferRegion + + +class CoreInterface(ABC): + """Full-attention 的资源定义、层视图和 KV 存储核心接口。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + # 具体存储和量化方案标识，用于 main/draft State 成员命名。 + @abstractmethod + def get_state_type(self) -> str: + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + - GQA 浮点 KV：gqa_{bf16|fp16|fp32}。 + - GQA 分组量化：gqa_int{4|8}_group{分组大小}_{scale精度}_scale。 + - GQA 静态 FP8：gqa_fp8_static_{per_tensor|per_head}。 + - MLA 浮点：mla_{精度}；KV/indexer 联合存储：mla_packed_{精度}。 + - MLA packed FP8：mla_packed_fp8_group128，或 mla_packed_fp8_group64_e8m0。 + - 序列压缩历史：compressed_c{压缩比例}_fp8_group64_e8m0。 + 例如 gqa_int8_group128_bf16_scale 对应 main_gqa_int8_group128_bf16_scale_state。 + 新量化实现必须区分影响方案选择的位宽、分组、scale 方式及精度，不编码标定数值。 + """ + ... + + @abstractmethod + def allocate_storage(self, token_capacity: int) -> None: + """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" + raise NotImplementedError() + + @abstractmethod + def get_token_bytes(self) -> int: + """累计当前 TP rank 各物理层的单 token 字节数，包含动态量化参数和行 padding；共享逻辑层不重复计数。""" + raise NotImplementedError() + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + @abstractmethod + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" + ... + + @abstractmethod + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + ... + + @abstractmethod + def get_layer_view(self, logical_layer_id: int) -> Any: + """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" + raise NotImplementedError() + + @abstractmethod + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: + """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。 + + 融合压缩等专用写入路径可直接使用计算视图，不调用本接口。 + """ + raise NotImplementedError() + + @abstractmethod + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制全部所辖层的完整 token 数据，包含动态 scale、indexer 字段，保持原存储表示。 + + 源和目标索引为存储设备上的等长一维 CUDA 整数 Tensor，位置不重叠，直接使用原 dtype 和 stride。""" + raise NotImplementedError() + + +class CpuCacheInterface(ABC): + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 cpu_cache_token_page_size 返回 CPU 字段、全局形状和 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 管理器按字段最大终点计算该 State 的 CPU 页占用。 + 本次有效 token 数由读写流程处理，不改变布局。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 CPU 页规格保持不变。 + 管理器实际分配存储后统一清除缓存，以兼容物理页规格变化；子类覆盖时须保留缓存能力。 + """ + ... + + @abstractmethod + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局写出 token 数据及配套字段；cpu_page_ready_flags 中为 True 的页跳过写入。""" + raise NotImplementedError() + + @abstractmethod + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" + raise NotImplementedError() + + +class DpCopyInterface(ABC): + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + @abstractmethod + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: CoreInterface, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，包含量化参数及 indexer 字段。 + + 源、目标索引须已位于各自存储的 CUDA 设备，可为 int32/int64；接口不转换设备或 dtype。 + 调用方保证 peer/IPC 可访问、源数据就绪、目标索引唯一，且源目标数据区域不重叠。 + """ + raise NotImplementedError() + + +class PdInterface(ABC): + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 PD 页规格保持不变；管理器分配存储后统一清除缓存。 + 子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将 token 数据写入 PD 页集合中的指定发送页，网络发送与同步由调用方负责。""" + raise NotImplementedError() + + @abstractmethod + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """从 PD 页集合中的指定接收页还原 token 数据，不分配索引或修改请求映射。""" + raise NotImplementedError() + + +__all__ = [ + "CoreInterface", + "CpuCacheInterface", + "DpCopyInterface", + "PdInterface", +] From 3debd74778502fc761845a56a1dc8f40a545bdcd Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:11:35 +0800 Subject: [PATCH 49/56] fix: pass pre-commit checks --- .../allocators/full_att_kv_allocator.py | 6 +- .../bounded_att/copy_kernel.py | 26 +-- .../bounded_att/gdn/gdn_state.py | 56 ++++--- .../bounded_att/kda/kda_pages_pool.py | 6 +- .../bounded_att/kda/kda_state.py | 56 ++++--- .../bounded_att/transfer_utils.py | 7 +- .../attention_state/bounded_att/utils.py | 2 + .../window_compressor_state.py | 15 +- .../attention_state/cpu_cache_operator.py | 20 +-- .../attention_state/dp_shared_operator.py | 6 +- .../attention_state/full_att/__init__.py | 6 +- .../attention_state/full_att/gqa/gqa_state.py | 36 ++-- .../full_att/gqa/gqa_transfer.py | 10 +- .../full_att/gqa/group_quantized_state.py | 67 +++++--- .../full_att/gqa/group_quantized_transfer.py | 10 +- .../full_att/gqa/static_fp8_state.py | 64 +++---- .../full_att/gqa/static_fp8_transfer.py | 10 +- .../attention_state/full_att/mla/mla_state.py | 23 +-- .../packed_mla/create_packed_mla_state.py | 4 +- .../full_att/packed_mla/packed_mla_state.py | 45 +++-- .../packed_mla/packed_mla_transfer.py | 12 +- .../full_att/packed_mla/packed_transfer.py | 10 +- lightllm/common/attention_state/manager.py | 40 +++-- .../common/attention_state/pd_operator.py | 12 +- lightllm/common/attention_state/utils.py | 7 +- .../common/basemodel/attention/linear/gdn.py | 15 +- .../common/basemodel/attention/linear/kda.py | 15 +- .../attention/nsa/dsv4_fp8_flashmla_sparse.py | 3 +- lightllm/common/basemodel/basemodel.py | 40 +++-- .../triton_kernel/destindex_copy_kv_fp8.py | 7 +- lightllm/models/bloom/model.py | 16 +- lightllm/models/deepseek2/model.py | 15 +- .../triton_kernel/gqa_flash_decoding_fp8.py | 8 +- lightllm/models/deepseek3_2/model.py | 15 +- .../models/deepseek_v4/layer_infer/indexer.py | 8 +- lightllm/models/deepseek_v4/model.py | 66 ++++---- .../deepseek_v4/triton_kernel/swa_cache.py | 10 +- lightllm/models/gemma3/model.py | 16 +- lightllm/models/gemma4/model.py | 50 ++++-- lightllm/models/gemma_2b/model.py | 16 +- lightllm/models/glm5_next/model.py | 90 +++++----- lightllm/models/glm5_next_mtp/model.py | 11 +- .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/llama/model.py | 16 +- lightllm/models/mistral/model.py | 16 +- lightllm/models/mixtral/model.py | 16 +- lightllm/models/qwen3next/model.py | 8 +- lightllm/models/starcoder/model.py | 16 +- lightllm/models/starcoder2/model.py | 16 +- .../router/dynamic_prompt/radix_cache.py | 4 +- lightllm/server/router/manager.py | 3 +- .../server/router/model_infer/infer_batch.py | 8 +- .../model_infer/mode_backend/base_backend.py | 16 +- .../dp_backend/dp_shared_kv_trans.py | 11 +- .../mode_backend/dp_backend/impl.py | 4 +- .../mode_backend/multi_level_kv_cache.py | 24 +-- .../mode_backend/pd/resource_io.py | 6 +- .../model_infer/shared_att_resources_utils.py | 7 +- lightllm/utils/cpu_cache_layout.py | 4 +- lightllm/utils/kv_cache_utils.py | 4 +- test/att_resources/validate_model_cache.py | 14 +- test/cpu_cache_kernel/test_speed.py | 17 +- unit_tests/common/att_resources/conftest.py | 4 +- .../common/att_resources/manager_fixtures.py | 3 +- .../common/att_resources/model_helpers.py | 80 ++++++--- .../att_resources/test_backend_startup.py | 33 ++-- .../att_resources/test_bounded_copy_kernel.py | 50 ++++-- .../att_resources/test_bounded_state.py | 19 ++- .../common/att_resources/test_capacity.py | 156 +++++++++++++----- .../att_resources/test_dsv4_integration.py | 6 +- .../common/att_resources/test_dsv4_state.py | 8 +- .../att_resources/test_fp8_compute_binding.py | 39 +++-- .../test_fp8_e8m0_packed_mla_state.py | 14 +- .../att_resources/test_full_att_cpu_kernel.py | 41 ++++- .../test_full_att_dp_transfer.py | 10 +- .../att_resources/test_full_att_transfer.py | 50 ++++-- .../common/att_resources/test_full_state.py | 20 ++- .../test_full_storage_initialization.py | 11 +- .../att_resources/test_gqa_pd_kernel.py | 49 +++--- .../common/att_resources/test_gqa_transfer.py | 32 ++-- .../test_hybrid_layer_mapping.py | 9 +- .../common/att_resources/test_layer_groups.py | 11 +- .../test_layer_state_selection.py | 53 +++--- .../att_resources/test_manager_helpers.py | 31 ++-- .../att_resources/test_model_calibration.py | 38 +++-- .../test_model_config_normalization.py | 34 ++-- .../att_resources/test_model_resources.py | 40 +++-- .../test_model_storage_definition.py | 38 ++++- .../att_resources/test_network_transport.py | 8 +- .../att_resources/test_page_pool_storage.py | 9 +- .../test_quantized_full_state.py | 48 ++++-- .../test_quantized_paged_storage.py | 3 +- .../att_resources/test_role_transfers.py | 16 +- .../att_resources/test_runtime_transport.py | 62 ++++--- .../att_resources/test_state_creation.py | 44 +++-- .../att_resources/test_state_identity.py | 19 ++- .../att_resources/test_state_queries.py | 21 ++- .../common/att_resources/test_tp_config.py | 21 ++- .../att_resources/test_transfer_views.py | 21 ++- .../test_window_compressor_state.py | 21 ++- .../basemodel/attention/linear/test_gdn.py | 14 +- .../basemodel/attention/test_dsv4_flashmla.py | 3 +- .../attention/test_fp8_logical_layer.py | 4 +- unit_tests/common/dsv4_test_utils.py | 8 +- .../common/test_deepseek4_paged_cache.py | 133 +++++++++++---- unit_tests/models/glm5_next/test_cache.py | 9 +- unit_tests/models/glm5_next/test_fp8_kv.py | 12 +- unit_tests/models/glm5_next/test_kernels.py | 19 ++- unit_tests/models/glm5_next/test_layers.py | 5 +- unit_tests/models/glm5_next/test_mtp.py | 10 +- unit_tests/models/glm5_next/test_pd_cache.py | 20 ++- unit_tests/models/sensenova_loop/test_loop.py | 13 +- .../mode_backend/test_dp_big_pages.py | 14 +- .../mode_backend/test_nccl_control_channel.py | 3 +- .../mode_backend/test_pd_token_indices.py | 4 +- .../model_infer/test_cpu_resource_layout.py | 2 +- .../test_hybrid_cpu_state_cache.py | 27 ++- .../model_infer/test_multimodal_chunking.py | 4 +- unit_tests/server/router/test_kv_capacity.py | 4 +- unit_tests/server/test_pd_start_args.py | 12 +- unit_tests/server/test_tokenizer.py | 11 +- 121 files changed, 1732 insertions(+), 944 deletions(-) diff --git a/lightllm/common/attention_state/allocators/full_att_kv_allocator.py b/lightllm/common/attention_state/allocators/full_att_kv_allocator.py index 545ed600b6..601b748553 100644 --- a/lightllm/common/attention_state/allocators/full_att_kv_allocator.py +++ b/lightllm/common/attention_state/allocators/full_att_kv_allocator.py @@ -51,9 +51,9 @@ def hold_token_indices(self) -> tuple[int, ...]: def alloc(self, need_size: int) -> torch.Tensor: assert need_size >= 0 and need_size % self.page_size == 0, "allocation must contain complete KV pages" - assert need_size <= self.mark_end - self.mark_start, ( - f"need {need_size} KV tokens, only {self.can_use_mem_size} available" - ) + assert ( + need_size <= self.mark_end - self.mark_start + ), f"need {need_size} KV tokens, only {self.can_use_mem_size} available" start = self.mark_start end = start + need_size self.mark_start = end diff --git a/lightllm/common/attention_state/bounded_att/copy_kernel.py b/lightllm/common/attention_state/bounded_att/copy_kernel.py index b4ec296fda..b66d3eb11b 100644 --- a/lightllm/common/attention_state/bounded_att/copy_kernel.py +++ b/lightllm/common/attention_state/bounded_att/copy_kernel.py @@ -261,16 +261,8 @@ def _copy_bounded_att_3d( axis1 = remaining_indices % SHAPE[1] remaining_indices = remaining_indices // SHAPE[1] axis0 = remaining_indices - source_offsets = ( - axis0 * src_stride_0 - + axis1 * src_stride_1 - + axis2 * src_stride_2 - ) - destination_offsets = ( - axis0 * dst_stride_0 - + axis1 * dst_stride_1 - + axis2 * dst_stride_2 - ) + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + axis2 * src_stride_2 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + axis2 * dst_stride_2 # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 @@ -328,18 +320,8 @@ def _copy_bounded_att_4d( axis1 = remaining_indices % SHAPE[1] remaining_indices = remaining_indices // SHAPE[1] axis0 = remaining_indices - source_offsets = ( - axis0 * src_stride_0 - + axis1 * src_stride_1 - + axis2 * src_stride_2 - + axis3 * src_stride_3 - ) - destination_offsets = ( - axis0 * dst_stride_0 - + axis1 * dst_stride_1 - + axis2 * dst_stride_2 - + axis3 * dst_stride_3 - ) + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + axis2 * src_stride_2 + axis3 * src_stride_3 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + axis2 * dst_stride_2 + axis3 * dst_stride_3 # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index 2cb26f0aed..b8a641a73d 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -57,7 +57,13 @@ def __init__( device=device, ) self.ssm_state = torch.zeros( - (storage_layer_count, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), dtype=config.ssm_dtype, device=device, ) @@ -67,10 +73,7 @@ def __init__( def get_state_type(self) -> str: """返回 GDN 算法及 Conv/SSM 存储精度。""" - return ( - f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" - f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - ) + return f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" def initialize_page_pools( self, @@ -134,7 +137,9 @@ def save_bounded_att_runtime_to_pages( page_slot_ids: List[int], ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 @@ -154,7 +159,9 @@ def restore_bounded_att_runtime_from_pages( page_slot_ids: List[int], ) -> None: """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] @@ -267,9 +274,9 @@ def copy_big_pages_from_dp( """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( - "DP copy requires the same algorithm and equivalent TP state configurations" - ) + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" src_pool = remote_dp_bounded_state.big_pages_pool dst_pool = self.big_pages_pool assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" @@ -291,22 +298,26 @@ def copy_runtime_state_from_dp( """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( - "DP copy requires the same algorithm and equivalent TP state configurations" - ) + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append(( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - )) - ssm_copies.append(( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - )) + conv_copies.append( + ( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + ) + ) + ssm_copies.append( + ( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + ) + ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) if self.mtp_state_index is not None: @@ -427,7 +438,8 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: ) regions.append(conv_region) logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", self.state_name, name, local_channel_start, diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py index 0ec20ae4fe..df17447cc4 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py @@ -18,9 +18,9 @@ def __init__( pin_memory: bool = True, ) -> None: assert size >= 0, "pool size must be nonnegative" - assert config.num_k_heads == config.num_v_heads and config.key_dim == config.value_dim, ( - "KDA pages require equal Q/K/V heads and dimensions" - ) + assert ( + config.num_k_heads == config.num_v_heads and config.key_dim == config.value_dim + ), "KDA pages require equal Q/K/V heads and dimensions" self.size = size device = torch.device(device) pin_memory = pin_memory and device.type == "cpu" diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index 9bb105dd1f..e5480a3ab0 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -57,7 +57,13 @@ def __init__( device=device, ) self.ssm_state = torch.zeros( - (storage_layer_count, request_capacity * (mtp_step + 1), config.num_v_heads, config.key_dim, config.value_dim), + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), dtype=config.ssm_dtype, device=device, ) @@ -67,10 +73,7 @@ def __init__( def get_state_type(self) -> str: """返回 KDA 算法及 Conv/SSM 存储精度。""" - return ( - f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" - f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - ) + return f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" def initialize_page_pools( self, @@ -134,7 +137,9 @@ def save_bounded_att_runtime_to_pages( page_slot_ids: List[int], ) -> None: """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 @@ -154,7 +159,9 @@ def restore_bounded_att_runtime_from_pages( page_slot_ids: List[int], ) -> None: """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] @@ -267,9 +274,9 @@ def copy_big_pages_from_dp( """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( - "DP copy requires the same algorithm and equivalent TP state configurations" - ) + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" src_pool = remote_dp_bounded_state.big_pages_pool dst_pool = self.big_pages_pool assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" @@ -291,22 +298,26 @@ def copy_runtime_state_from_dp( """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config, ( - "DP copy requires the same algorithm and equivalent TP state configurations" - ) + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append(( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - )) - ssm_copies.append(( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - )) + conv_copies.append( + ( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + ) + ) + ssm_copies.append( + ( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + ) + ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) if self.mtp_state_index is not None: @@ -427,7 +438,8 @@ def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: ) regions.append(conv_region) logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, channels=%d, replicas=%d, write_enabled=%s", + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", self.state_name, name, local_channel_start, diff --git a/lightllm/common/attention_state/bounded_att/transfer_utils.py b/lightllm/common/attention_state/bounded_att/transfer_utils.py index 112358c06a..c5123600d1 100644 --- a/lightllm/common/attention_state/bounded_att/transfer_utils.py +++ b/lightllm/common/attention_state/bounded_att/transfer_utils.py @@ -11,6 +11,7 @@ # 一、CPU/PD 页交换：按状态字段生成复制视图 # ======================================================================== + def bounded_att_buffer_copies( page_fields: list[Mapping[str, TensorField]], page_buffers: torch.Tensor, @@ -41,7 +42,7 @@ def bounded_att_buffer_copies( # 阶段二：region 的字节区间和 shape 定位外部页中的全局字段，再对齐轴顺序。 page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) - external_tensor = page_bytes[region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + external_tensor = page_bytes[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) assert local_tensor.dtype == external_tensor.dtype, "local and external field dtypes must match" external_tensor = external_tensor.permute(tuple(region.axes.index(axis) for axis in local_field.axes)) @@ -52,7 +53,9 @@ def bounded_att_buffer_copies( for axis_mapping in region.axis_mappings: axis = local_field.axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local_view = local_tensor[tuple(local_slices)] external_view = external_tensor[tuple(external_slices)] assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" diff --git a/lightllm/common/attention_state/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py index d8e4c2fbb7..93f5b2ebe3 100644 --- a/lightllm/common/attention_state/bounded_att/utils.py +++ b/lightllm/common/attention_state/bounded_att/utils.py @@ -23,6 +23,7 @@ # 一、页池分配：普通存储或 DP fetch 所需的共享 pinned 存储 # ======================================================================== + def create_page_pool_tensor( size: tuple[int, ...], dtype: torch.dtype, @@ -81,6 +82,7 @@ def create_page_pool_tensor( # 二、本进程注册：首次创建和跨进程接收复用同一入口 # ======================================================================== + def register_page_pool_tensor(tensor: torch.Tensor) -> None: """将共享 CPU 存储注册到当前进程，已 pinned、空存储或 GPU Tensor 无需处理。 diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index 9fe226d456..cea0a49af2 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -114,8 +114,7 @@ def get_state_type(self) -> str: """用途、窗口、步长和精度共同标识可复用的压缩器方案。""" precision = {torch.float32: "fp32", torch.float16: "fp16", torch.bfloat16: "bf16"}[self.config.dtype] return ( - f"compressor_{self.config.target}_w{self.config.window_size}" - f"_s{self.config.compress_stride}_{precision}" + f"compressor_{self.config.target}_w{self.config.window_size}" f"_s{self.config.compress_stride}_{precision}" ) def initialize_page_pools( @@ -225,14 +224,14 @@ def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: """校验缓存页对齐产出步长；无重叠历史的窗口不占 CPU 页，也不创建大小页池。""" args = get_env_start_args() # 大页是小页的整数倍；小页对齐产出步长后，大小页终点均不会落在未完成的组内。 - assert args.linear_att_hash_page_size % self.config.compress_stride == 0, ( - "state cache pages must align to the compressor stride" - ) + assert ( + args.linear_att_hash_page_size % self.config.compress_stride == 0 + ), "state cache pages must align to the compressor stride" # CPU 服务只查询 meta 布局，也要在忽略无重叠窗口之前拒绝不安全的恢复边界。 if args.cpu_cache_token_page_size is not None: - assert args.cpu_cache_token_page_size % self.config.compress_stride == 0, ( - "CPU cache pages must align to the compressor stride" - ) + assert ( + args.cpu_cache_token_page_size % self.config.compress_stride == 0 + ), "CPU cache pages must align to the compressor stride" return self._get_transfer_regions(runtime=False) def write_state_pages_to_cpu( diff --git a/lightllm/common/attention_state/cpu_cache_operator.py b/lightllm/common/attention_state/cpu_cache_operator.py index 9f07ec0356..814d340481 100644 --- a/lightllm/common/attention_state/cpu_cache_operator.py +++ b/lightllm/common/attention_state/cpu_cache_operator.py @@ -83,17 +83,13 @@ def write_req( big_page_slot_ids = [] if big_page_slot_ids is None else big_page_slot_ids small_page_slot_ids = [] if small_page_slot_ids is None else small_page_slot_ids big_page_count = len(big_page_slot_ids) - assert big_page_count + len(small_page_slot_ids) == len(cpu_page_indices), ( - "one state slot is required per CPU page" - ) + assert big_page_count + len(small_page_slot_ids) == len( + cpu_page_indices + ), "one state slot is required per CPU page" big_cpu_page_indices = cpu_page_indices[:big_page_count] small_cpu_page_indices = cpu_page_indices[big_page_count:] - big_page_readies = ( - cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None - ) - small_page_readies = ( - cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None - ) + big_page_readies = cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None + small_page_readies = cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None for state in full_states: state.write_to_cpu( @@ -148,9 +144,9 @@ def read_req( page_tokens = get_env_start_args().cpu_cache_token_page_size has_partial_page = token_indices.numel() % page_tokens != 0 assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert len(big_page_slot_ids) + int(has_partial_page) == len(cpu_page_indices), ( - "slots must cover complete CPU pages only" - ) + assert len(big_page_slot_ids) + int(has_partial_page) == len( + cpu_page_indices + ), "slots must cover complete CPU pages only" load_big_page_slot_ids = big_page_slot_ids.copy() if has_partial_page: load_big_page_slot_ids.append(self.owner.bounded_att_cpu_load_big_page_slot) diff --git a/lightllm/common/attention_state/dp_shared_operator.py b/lightllm/common/attention_state/dp_shared_operator.py index 8aa90330cd..5f0eb7fa0d 100644 --- a/lightllm/common/attention_state/dp_shared_operator.py +++ b/lightllm/common/attention_state/dp_shared_operator.py @@ -22,9 +22,9 @@ def create_ipc_copy(self) -> AttentionStateManager: if pool is not None: for tensor in vars(pool).values(): if isinstance(tensor, torch.Tensor): - assert tensor.device.type == "cpu" and tensor.is_shared(), ( - "DP IPC requires shared CPU page pools" - ) + assert ( + tensor.device.type == "cpu" and tensor.is_shared() + ), "DP IPC requires shared CPU page pools" ipc_copy = copy(self.owner) ipc_copy.full_att_allocator = None ipc_copy.bounded_att_big_page_allocator = None diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py index 4c57c49c1e..848fe47a51 100644 --- a/lightllm/common/attention_state/full_att/__init__.py +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -40,7 +40,11 @@ "PackedMLALayerView", ] -from .compressed_att.compressed_att_state import CompressedAttentionConfig, CompressedAttentionState, CompressedAttentionLayerView +from .compressed_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, + CompressedAttentionLayerView, +) from .compressed_att.create_compressed_att_state import create_compressed_att_state __all__ += [ diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index a9d389a91d..64cc0e384b 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -54,7 +54,9 @@ def __post_init__(self) -> None: isinstance(group, tuple) and group for group in self.logical_layer_groups ), "layer groups must be a tuple of nonempty tuples" layers = [layer for group in self.logical_layer_groups for layer in group] - assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size assert ( @@ -118,10 +120,7 @@ def allocate_storage(self, token_capacity: int) -> None: def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" config = self.config - return ( - len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim - * self.storage_dtype.itemsize - ) + return len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 @@ -150,18 +149,23 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv(K=kv[:, :self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index]) - destindex_copy_kv(K=kv[:, self.config.head_num:], DestLoc=token_indices, Out=self.v_storage[storage_layer_index]) + destindex_copy_kv( + K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] + ) + destindex_copy_kv( + K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] + ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" - assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( - "token copies require equal one-dimensional index tensors" - ) - assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( - "token indices must be integer tensors" - ) + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" assert self.k_storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return @@ -248,9 +252,9 @@ def copy_tokens_from_dp( remote_dp_full_att_state: GQAFullAttentionState, ) -> None: """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" - assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( - "DP copy requires the same encoding and equivalent TP state configurations" - ) + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" copy_gqa_tokens( diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py index 994cd3904a..516c947864 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py @@ -86,7 +86,7 @@ def transfer_gqa_pages( else: raise AssertionError(f"unknown transfer field {region.field}") # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 - external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) external = external.view(page_buffers.shape[0], *region.shape) # 外部统一为 [page, layer, page_token, head, dim]；+1 是因为多了 page 轴。 external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) @@ -97,13 +97,17 @@ def transfer_gqa_pages( # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 axis = axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 assert local.dtype == external.dtype assert (local.shape[0], local.shape[2], local.shape[3]) == ( - external.shape[1], external.shape[3], external.shape[4] + external.shape[1], + external.shape[3], + external.shape[4], ), "local and external copy views must have the same shape" copy_views.append((local, external)) diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py index a35a1aad9b..28b84c9551 100644 --- a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py @@ -58,7 +58,9 @@ def __post_init__(self) -> None: isinstance(group, tuple) and group for group in self.logical_layer_groups ), "layer groups must be a tuple of nonempty tuples" layers = [layer for group in self.logical_layer_groups for layer in group] - assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size assert ( @@ -66,18 +68,19 @@ def __post_init__(self) -> None: or self.global_head_num % self.tp_world_size == 0 or self.tp_world_size % self.global_head_num == 0 ), "KV heads require complete TP shards or replicas" - assert self.bits in (4, 8) and self.group_size >= 2 and not ( - self.group_size & (self.group_size - 1) + assert ( + self.bits in (4, 8) and self.group_size >= 2 and not (self.group_size & (self.group_size - 1)) ), "INT4/INT8 quantization requires a power-of-two group size" assert self.head_dim % self.group_size == 0 and self.scale_dtype in ( - torch.float16, torch.bfloat16, torch.float32 + torch.float16, + torch.bfloat16, + torch.float32, ), "group scales require a floating dtype and complete groups" class GroupQuantizedFullAttentionState(FullAttentionState): """管理 INT4/INT8 KV 与逐 token 分组 scale，复制和传输始终覆盖二者。""" - # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -108,7 +111,9 @@ def __init__( def get_state_type(self) -> str: """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" - return f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.scale_dtype)}_scale" + return ( + f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.scale_dtype)}_scale" + ) def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" @@ -117,8 +122,10 @@ def allocate_storage(self, token_capacity: int) -> None: self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) shape = ( - len(self.config.logical_layer_groups), token_capacity + self.page_size, - self.config.head_num, self.config.head_dim // self.config.group_size, + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + self.config.head_num, + self.config.head_dim // self.config.group_size, ) self.k_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) self.v_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) @@ -142,10 +149,15 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] - def get_layer_view(self, logical_layer_id: int) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + def get_layer_view( + self, logical_layer_id: int + ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """返回该层 K/V 编码数据及各自的分组 scale 视图。""" storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) - return (self.k_storage[storage_layer_index], self.k_scales[storage_layer_index]), (self.v_storage[storage_layer_index], self.v_scales[storage_layer_index]) + return (self.k_storage[storage_layer_index], self.k_scales[storage_layer_index]), ( + self.v_storage[storage_layer_index], + self.v_scales[storage_layer_index], + ) def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: """量化输入浮点 K/V，并将编码数据和分组 scale 写入相同 token 位置。""" @@ -160,7 +172,7 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, (0, self.k_storage, self.k_scales), (self.config.head_num, self.v_storage, self.v_scales), ): - values = kv[:, offset:offset + self.config.head_num] + values = kv[:, offset : offset + self.config.head_num] if self.config.bits == 4: from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv @@ -172,7 +184,9 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, quant_group_size=self.config.group_size, ) else: - from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import destindex_copy_quantize_kv + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( + destindex_copy_quantize_kv, + ) destindex_copy_quantize_kv( K=values, @@ -185,12 +199,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" - assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( - "token copies require equal one-dimensional index tensors" - ) - assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( - "token indices must be integer tensors" - ) + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" assert self.k_storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return @@ -294,12 +309,10 @@ def copy_tokens_from_dp( assert ( remote_dp_full_att_state.config.bits, remote_dp_full_att_state.config.group_size, - ) == ( - self.config.bits, self.config.group_size - ), "incompatible group quantization" - assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( - "DP copy requires the same encoding and equivalent TP state configurations" - ) + ) == (self.config.bits, self.config.group_size), "incompatible group quantization" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" copy_group_quantized_tokens( @@ -458,6 +471,8 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: def _storage_shape(self) -> tuple[int, ...]: return ( - len(self.config.logical_layer_groups), self.capacity + self.page_size, - self.config.head_num, self.config.head_dim * self.config.bits // 8, + len(self.config.logical_layer_groups), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim * self.config.bits // 8, ) diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py index 1662896ff0..72eea0ba6b 100644 --- a/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py @@ -96,7 +96,7 @@ def transfer_group_quantized_pages( else: raise AssertionError(f"unknown transfer field {region.field}") # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 - external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) external = external.view(page_buffers.shape[0], *region.shape) # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) @@ -107,13 +107,17 @@ def transfer_group_quantized_pages( # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 axis = axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 assert local.dtype == external.dtype assert (local.shape[0], local.shape[2], local.shape[3]) == ( - external.shape[1], external.shape[3], external.shape[4] + external.shape[1], + external.shape[3], + external.shape[4], ), "local and external copy views must have the same shape" copy_views.append((local, external)) diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py index d907d7813d..20a79bd917 100644 --- a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py @@ -56,7 +56,9 @@ def __post_init__(self) -> None: isinstance(group, tuple) and group for group in self.logical_layer_groups ), "layer groups must be a tuple of nonempty tuples" layers = [layer for group in self.logical_layer_groups for layer in group] - assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size assert ( @@ -70,7 +72,6 @@ def __post_init__(self) -> None: class StaticFP8FullAttentionState(FullAttentionState): """使用静态标定参数的 FP8 KV 存储，固定 scale 随 State 初始化，不逐 token 搬运。""" - # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -112,7 +113,9 @@ def bind_calibration( 无文件时仅 per_tensor 可使用单位 scale；per_head 必须有 K/V 标定，Q 标定可选。 未提供 Q 标定时由 attention backend 动态量化 Q。 """ - assert self.k_storage is None and self.scales is None, "calibration must be bound once before storage allocation" + assert ( + self.k_storage is None and self.scales is None + ), "calibration must be bound once before storage allocation" config = self.config storage_layer_count = len(config.logical_layer_groups) if calibration is None: @@ -123,28 +126,30 @@ def bind_calibration( # 阶段一：检查文件编码和物理行覆盖范围。 calibrated_layer_count = calibration["num_layers"] storage_layer_end = storage_layer_start + storage_layer_count - assert 0 <= storage_layer_start < storage_layer_end <= calibrated_layer_count, ( - "calibration does not cover this State's physical layers" - ) - assert calibration.get("quant_type", config.scale_mode) == config.scale_mode, ( - "calibration quantization mode does not match the State" - ) + assert ( + 0 <= storage_layer_start < storage_layer_end <= calibrated_layer_count + ), "calibration does not cover this State's physical layers" + assert ( + calibration.get("quant_type", config.scale_mode) == config.scale_mode + ), "calibration quantization mode does not match the State" limits = torch.finfo(torch.float8_e4m3fn) - assert calibration.get("qmin", limits.min) == limits.min and calibration.get("qmax", limits.max) == limits.max, ( - "calibration range does not match float8_e4m3fn" - ) + assert ( + calibration.get("qmin", limits.min) == limits.min and calibration.get("qmax", limits.max) == limits.max + ), "calibration range does not match float8_e4m3fn" def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]) -> torch.Tensor: tensor = torch.tensor(data=values["scales"], dtype=torch.float32) assert tuple(values["scales_shape"]) == tensor.shape == shape, "invalid calibration scale shape" - assert bool((torch.isfinite(tensor) & (tensor > 0)).all()), "calibration scales must be finite and positive" + assert bool( + (torch.isfinite(tensor) & (tensor > 0)).all() + ), "calibration scales must be finite and positive" return tensor[storage_layer_start:storage_layer_end] def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: heads = config.head_num - assert calibrated_heads > 0 and (heads * config.tp_world_size) % calibrated_heads == 0, ( - "calibration heads do not cover the model's KV heads" - ) + assert ( + calibrated_heads > 0 and (heads * config.tp_world_size) % calibrated_heads == 0 + ), "calibration heads do not cover the model's KV heads" expanded = tensor.repeat_interleave(repeats=heads * config.tp_world_size // calibrated_heads, dim=-1) return expanded[:, config.tp_rank * heads : (config.tp_rank + 1) * heads].contiguous() @@ -158,7 +163,8 @@ def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: q_scales = None if q_config is not None: q_scales = shard( - scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), q_config["num_head"], + scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), + q_config["num_head"], ) else: scales = scale_tensor(calibration, (calibrated_layer_count, 2)) @@ -190,10 +196,7 @@ def allocate_storage(self, token_capacity: int) -> None: def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" config = self.config - return ( - len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim - * self.storage_dtype.itemsize - ) + return len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize # ======================================================================== # 二、Full attention：计算视图、KV 写入与完整 token 复制 @@ -233,12 +236,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" - assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( - "token copies require equal one-dimensional index tensors" - ) - assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( - "token indices must be integer tensors" - ) + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" assert self.k_storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return @@ -325,9 +329,9 @@ def copy_tokens_from_dp( remote_dp_full_att_state: StaticFP8FullAttentionState, ) -> None: """复制相同配置的 FP8 KV 编码；各 DP 使用相同的静态标定参数。""" - assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( - "DP copy requires the same encoding and equivalent TP state configurations" - ) + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" copy_gqa_tokens( diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py index 9c70016acd..95db4d79d4 100644 --- a/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py @@ -88,7 +88,7 @@ def transfer_static_fp8_pages( else: raise AssertionError(f"unknown transfer field {region.field}") # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 - external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) external = external.view(page_buffers.shape[0], *region.shape) # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) @@ -99,13 +99,17 @@ def transfer_static_fp8_pages( # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 axis = axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 assert local.dtype == external.dtype assert (local.shape[0], local.shape[2], local.shape[3]) == ( - external.shape[1], external.shape[3], external.shape[4] + external.shape[1], + external.shape[3], + external.shape[4], ), "local and external copy views must have the same shape" copy_views.append((local, external)) diff --git a/lightllm/common/attention_state/full_att/mla/mla_state.py b/lightllm/common/attention_state/full_att/mla/mla_state.py index e10dd953a1..a818f5af3a 100644 --- a/lightllm/common/attention_state/full_att/mla/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla/mla_state.py @@ -41,7 +41,9 @@ def __post_init__(self) -> None: isinstance(group, tuple) and group for group in self.logical_layer_groups ), "layer groups must be a tuple of nonempty tuples" layers = [layer for group in self.logical_layer_groups for layer in group] - assert all(isinstance(layer, int) and layer >= 0 for layer in layers), "logical layer IDs must be nonnegative integers" + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" assert self.latent_dim > 0 and 0 <= self.tp_rank < self.tp_world_size @@ -126,12 +128,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" - assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( - "token copies require equal one-dimensional index tensors" - ) - assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( - "token indices must be integer tensors" - ) + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" assert self.storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return @@ -209,9 +212,9 @@ def copy_tokens_from_dp( remote_dp_full_att_state: MLAFullAttentionState, ) -> None: """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" - assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( - "DP copy requires the same encoding and equivalent TP state configurations" - ) + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" copy_packed_tokens( diff --git a/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py index afb237eb7a..f41646f119 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py @@ -17,9 +17,7 @@ def create_packed_mla_state( """Create a packed MLA State; packed formats always include an indexer field.""" assert type(config) is PackedMLALogicalConfig, "packed MLA creation requires PackedMLALogicalConfig" if config.format == "packed_group64_e8m0": - assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, ( - "FP8/E8M0 MLA requires a 512-dimensional latent" - ) + assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" return FP8E8M0PackedMLAState( layout_config=FP8E8M0PackedMLAConfig( logical_layer_groups=config.logical_layer_groups, diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py index fb6281fd7c..99c920fe26 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py @@ -50,7 +50,10 @@ def __post_init__(self) -> None: and 0 <= self.tp_rank < self.tp_world_size ) assert self.format in ( - "packed_fp16", "packed_bf16", "packed_group128_fp8", "packed_group64_e8m0" + "packed_fp16", + "packed_bf16", + "packed_group128_fp8", + "packed_group64_e8m0", ), "unknown packed MLA logical format" @@ -143,7 +146,6 @@ class PackedMLALayerView: class PackedMLAFullAttentionState(FullAttentionState): """在同一物理 token 行中存储 MLA latent、量化参数及稀疏 indexer key/scale。""" - # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -172,7 +174,8 @@ def get_state_type(self) -> str: """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" encoding = ( f"fp8_group{self.config.quant_group_size}" - if self.config.kv_encoding == "fp8_group128" else dtype_name(dtype=self.config.dtype) + if self.config.kv_encoding == "fp8_group128" + else dtype_name(dtype=self.config.dtype) ) return "mla_packed_" + encoding @@ -214,7 +217,8 @@ def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: raw = row.view(torch.uint8) kv = ( raw[..., : self.config.kv_bytes] - if self.config.kv_encoding == "fp8_group128" else row[..., : self.config.latent_dim] + if self.config.kv_encoding == "fp8_group128" + else row[..., : self.config.latent_dim] ) indexer = raw[..., self.config.indexer_offset : self.config.indexer_offset + self.config.indexer_dim + 4] return PackedMLALayerView(kv=kv, indexer_k=indexer) @@ -248,17 +252,20 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, else: from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.latent_dim]) + destindex_copy_kv( + K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.latent_dim] + ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" - assert src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel(), ( - "token copies require equal one-dimensional index tensors" - ) - assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in (torch.int32, torch.int64), ( - "token indices must be integer tensors" - ) + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" assert self.storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return @@ -351,9 +358,9 @@ def copy_tokens_from_dp( self.config.latent_dim, self.config.indexer_dim, ), "incompatible packed MLA format" - assert type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config, ( - "DP copy requires the same encoding and equivalent TP state configurations" - ) + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" copy_packed_tokens( @@ -441,10 +448,12 @@ def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: ] else: field_specs = [("kv", config.dtype, config.latent_dim, "dim")] - field_specs.extend(( - ("indexer_k", torch.uint8, config.indexer_dim, "dim"), - ("indexer_scale", torch.float32, 1, "scale"), - )) + field_specs.extend( + ( + ("indexer_k", torch.uint8, config.indexer_dim, "dim"), + ("indexer_scale", torch.float32, 1, "scale"), + ) + ) for field_name, field_dtype, field_width, last_axis in field_specs: axes = ("layer", "token", "head", last_axis) local_shape = (len(config.logical_layer_groups), page_tokens, 1, field_width) diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py index 5b5bb4534f..86a0dc4145 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py @@ -100,12 +100,12 @@ def transfer_packed_mla_pages( local = raw[..., indexer_offset:indexer_end] axes = ("layer", "token", "head", "dim") elif region.field == "indexer_scale": - local = raw[..., indexer_end:indexer_end + 4].view(torch.float32) + local = raw[..., indexer_end : indexer_end + 4].view(torch.float32) axes = ("layer", "token", "head", "scale") else: raise AssertionError(f"unknown transfer field {region.field}") # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 - external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) external = external.view(page_buffers.shape[0], *region.shape) # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) @@ -116,13 +116,17 @@ def transfer_packed_mla_pages( # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 axis = axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 assert local.dtype == external.dtype assert (local.shape[0], local.shape[2], local.shape[3]) == ( - external.shape[1], external.shape[3], external.shape[4] + external.shape[1], + external.shape[3], + external.shape[4], ), "local and external copy views must have the same shape" copy_views.append((local, external)) diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py index bd04051a98..03d00746bd 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py @@ -83,7 +83,7 @@ def transfer_packed_pages( else: raise AssertionError(f"unknown transfer field {region.field}") # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 - external = page_bytes[:, region.byte_start:region.byte_end].view(region.dtype) + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) external = external.view(page_buffers.shape[0], *region.shape) # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) @@ -94,13 +94,17 @@ def transfer_packed_pages( # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 axis = axes.index(axis_mapping.axis) local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) - external_slices[axis + 1] = slice(axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) local = local[tuple(local_slices)] external = external[tuple(external_slices)] # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 assert local.dtype == external.dtype assert (local.shape[0], local.shape[2], local.shape[3]) == ( - external.shape[1], external.shape[3], external.shape[4] + external.shape[1], + external.shape[3], + external.shape[4], ), "local and external copy views must have the same shape" copy_views.append((local, external)) diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index ff55d7c34d..6fb18c7ea1 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -77,9 +77,7 @@ def get_pd_operator(self) -> PDOperator: """返回当前资源管理器唯一的 PD 操作对象。""" return self._pd_operator - def register_state( - self, model_role: str, state: FullAttentionState | BoundedAttentionState - ) -> None: + def register_state(self, model_role: str, state: FullAttentionState | BoundedAttentionState) -> None: """按模型角色和具体方案注册 State，保留原对象及其存储。 成员名和 state_name 均为 {model_role}_{state.get_state_type()}_state。 @@ -93,7 +91,9 @@ def register_state( assert state_type.isidentifier() and state_type.islower(), f"invalid State type {state_type}" name = f"{model_role}_{state_type}_state" assert name not in self.__dict__, f"State {name} is already bound" - assert all(state is not value for value in self.__dict__.values()), "a State cannot belong to multiple model roles" + assert all( + state is not value for value in self.__dict__.values() + ), "a State cannot belong to multiple model roles" # config 只描述布局，资源身份仅保存在 State；绑定时保留原配置对象。 state.state_name = name setattr(self, name, state) @@ -155,9 +155,9 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: if any(state.config.scale_mode == "per_tensor" for state in calibration_states): from lightllm.utils.config_utils import get_model_architectures - assert calibration["architectures"] == get_model_architectures(args.model_dir), ( - "calibration architecture does not match the model" - ) + assert calibration["architectures"] == get_model_architectures( + args.model_dir + ), "calibration architecture does not match the model" # 阶段二：显式使用模型角色顺序，避免注册先后及 draft10 的字符串排序影响标定行。 storage_layer_start = 0 @@ -165,7 +165,8 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: states = sorted( self.get_full_attention_states(role), key=lambda state: ( - min(layer for group in state.config.logical_layer_groups for layer in group), state.state_name, + min(layer for group in state.config.logical_layer_groups for layer in group), + state.state_name, ), ) for state in states: @@ -175,7 +176,9 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: storage_layer_start=storage_layer_start, ) storage_layer_start += len(state.config.logical_layer_groups) - logger.info("Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales") + logger.info( + "Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales" + ) def initialize_storage( self, @@ -281,7 +284,8 @@ def allocate_storage( if not state.get_cpu_transfer_regions(): continue state.initialize_page_pools( - big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, + big_page_capacity=big_page_capacity, + small_page_capacity=small_page_capacity, pin_memory=torch.device(device).type == "cuda", ) self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 @@ -343,11 +347,15 @@ def save_bounded_att_runtime_to_pages( ) -> None: """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" bounded_states = self.get_bounded_attention_states() - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(page_slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" + assert ( + len(page_slot_ids) == len(req_indices) == len(state_token_lens) + ), "one length and slot are required per request" for state in bounded_states: state.save_bounded_att_runtime_to_pages( req_indices=req_indices, @@ -365,11 +373,15 @@ def restore_bounded_att_runtime_from_pages( ) -> None: """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" bounded_states = self.get_bounded_attention_states() - assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "request and slot indices must be lists" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator assert allocator is not None, f"{page_kind} pages are disabled" - assert len(page_slot_ids) == len(req_indices) == len(state_token_lens), "one length and slot are required per request" + assert ( + len(page_slot_ids) == len(req_indices) == len(state_token_lens) + ), "one length and slot are required per request" assert not page_slot_ids or bounded_states, "state restore requires participating bounded states" for state in bounded_states: state.restore_bounded_att_runtime_from_pages( diff --git a/lightllm/common/attention_state/pd_operator.py b/lightllm/common/attention_state/pd_operator.py index 576ce7e9ff..ec1c35c190 100644 --- a/lightllm/common/attention_state/pd_operator.py +++ b/lightllm/common/attention_state/pd_operator.py @@ -119,9 +119,9 @@ def write_page( ) case "att_state": assert req_idx is not None, "state PD pages require a request index" - assert state_token_len is not None and state_token_len >= 0, ( - "PD state requires its prefill prefix length" - ) + assert ( + state_token_len is not None and state_token_len >= 0 + ), "PD state requires its prefill prefix length" for state in self.owner.get_bounded_attention_states(): state.write_to_pd_page( req_idx=req_idx, @@ -155,9 +155,9 @@ def read_page( ) case "att_state": assert req_idx is not None, "state PD pages require a request index" - assert state_token_len is not None and state_token_len >= 0, ( - "PD state requires its prefill prefix length" - ) + assert ( + state_token_len is not None and state_token_len >= 0 + ), "PD state requires its prefill prefix length" for state in self.owner.get_bounded_attention_states(): state.read_from_pd_page( pd_page_buffers=pd_page_buffers, diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py index c2f331f5bf..5de4267819 100644 --- a/lightllm/common/attention_state/utils.py +++ b/lightllm/common/attention_state/utils.py @@ -20,6 +20,7 @@ # 一、State 命名与返回值副本 # ======================================================================== + def dtype_name(dtype: torch.dtype): return {torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32"}[dtype] @@ -70,9 +71,9 @@ def merge_contiguous_tail_dims( if start_dim < 0: start_dim += ndim tail_dims = ndim - start_dim - assert len(destination_shape) >= tail_dims and source_shape[start_dim:] == destination_shape[-tail_dims:], ( - "copy tail shapes must match" - ) + assert ( + len(destination_shape) >= tail_dims and source_shape[start_dim:] == destination_shape[-tail_dims:] + ), "copy tail shapes must match" # 阶段一：从末轴向前检查；两侧的前缀容量和轴数独立，按距末轴的距离对齐。 first_tail_dim = ndim - 1 diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 81a740ffa9..986214568f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -141,8 +141,10 @@ def prefill_att( backend: LinearAttBackend = self.backend state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, state_type_filter="gdn", - logical_layer_id=layer_num, state_class=GDNState, + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 @@ -220,7 +222,8 @@ def init_state(self): # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 (state,) = ( - state for state in self.infer_state.att_state_manager.get_bounded_attention_states( + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states( self.infer_state.att_state_role ) if isinstance(state, GDNState) @@ -297,8 +300,10 @@ def decode_att( mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, state_type_filter="gdn", - logical_layer_id=layer_num, state_class=GDNState, + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=GDNState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 0b7f7eaa1d..cd2f8afb39 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -85,8 +85,10 @@ def prefill_att( backend: KDALinearAttBackend = self.backend state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, state_type_filter="kda", - logical_layer_id=layer_num, state_class=KDAState, + model_role=self.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=KDAState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] @@ -144,7 +146,8 @@ def init_state(self): # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 (state,) = ( - state for state in self.infer_state.att_state_manager.get_bounded_attention_states( + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states( self.infer_state.att_state_role ) if isinstance(state, KDAState) @@ -220,8 +223,10 @@ def decode_att( backend: KDALinearAttBackend = self.backend state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, state_type_filter="kda", - logical_layer_id=layer_num, state_class=KDAState, + model_role=self.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=KDAState, ) conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py index 1fd992aa37..93560de07a 100644 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py @@ -81,7 +81,8 @@ def get_att_input_params(self, infer_state, logical_layer_id: int): compressed_kv = infer_state.get_compressed_kv_view(logical_layer_id=logical_layer_id).kv if ratio else None if infer_state.is_prefill: swa = infer_state.prefill_att_state.swa_buffer[config.logical_layer_ids.index(logical_layer_id)].view( - -1, config.get_packed_kv_bytes(config.page_token_size), + -1, + config.get_packed_kv_bytes(config.page_token_size), ) return swa, compressed_kv diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e0c439cd8a..78af939391 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -147,7 +147,9 @@ def prepare_attention_states(self) -> None: if self.is_mtp_draft_model: self.att_state_manager = self.main_model.att_state_manager assert not self.att_state_manager._state_registration_finished, "State registration is already finished" - assert self.max_req_num + 1 == self.att_state_manager._request_capacity, "models must share the request domain" + assert ( + self.max_req_num + 1 == self.att_state_manager._request_capacity + ), "models must share the request domain" with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): if not self.is_mtp_draft_model: self.att_state_manager = AttentionStateManager(request_capacity=self.max_req_num + 1) @@ -158,12 +160,19 @@ def prepare_attention_states(self) -> None: # 当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；各工厂自行选择参与层及共享分组。 logical_layer_ids = tuple(layer.layer_num_ for layer in self.layers_infer) full_states = self.create_full_att_states( - model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, - tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, + model_config=self.config, + args=self.args, + model_logical_layer_ids=logical_layer_ids, + tp_rank=self.tp_rank_, + tp_world_size=self.tp_world_size_, ) bounded_states = self.create_bounded_att_states( - model_config=self.config, args=self.args, model_logical_layer_ids=logical_layer_ids, - request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, + model_config=self.config, + args=self.args, + model_logical_layer_ids=logical_layer_ids, + request_capacity=self.max_req_num + 1, + tp_rank=self.tp_rank_, + tp_world_size=self.tp_world_size_, device="cuda", ) for state in (*full_states, *bounded_states): @@ -254,8 +263,13 @@ def load_weights(self, weight_dict: dict): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """返回本模型的 full-attention 定义列表；默认没有此类资源，不分配 KV 容量。 @@ -266,9 +280,15 @@ def create_full_att_states( @classmethod def create_bounded_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], request_capacity: int = 1, - tp_rank: int = 0, tp_world_size: int = 1, device: str = "meta", + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", ) -> list["BoundedAttentionState"]: """返回本模型的 bounded 运行态列表；meta 用于布局查询，大小页后续分配。 diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py index 23be8798e5..b94126fcaa 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py @@ -35,7 +35,12 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( if SPLIT_KV: # 两块独立分配的 K/V 具有相同形状和 stride，输入仍按 K heads、V heads 排列。 k_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] - v_ptrs = Out_v + dest_index * stride_o_bs + stride_o_h * (offs_h[:, None] - head_num // 2) + stride_o_d * offs_d[None, :] + v_ptrs = ( + Out_v + + dest_index * stride_o_bs + + stride_o_h * (offs_h[:, None] - head_num // 2) + + stride_o_d * offs_d[None, :] + ) o_ptrs = tl.where(offs_h[:, None] < head_num // 2, k_ptrs, v_ptrs) else: o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 5cf512e10c..8e963d7f84 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -33,8 +33,13 @@ def __init__(self, kvargs): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -50,9 +55,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 2521054b4a..e06c18ea16 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -65,8 +65,13 @@ def _verify_params(self): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" from lightllm.common.attention_state.full_att import MLAConfig, create_mla_state @@ -81,9 +86,11 @@ def create_full_att_states( logical_layer_groups=logical_layer_groups, latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], latent_dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - page_size=args.page_size, device="meta", + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py index a086349983..795f75459f 100644 --- a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py +++ b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py @@ -24,7 +24,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( softmax_scale, out=None, alloc_tensor_func=torch.empty, - **run_config + **run_config, ): batch_size = infer_state.batch_size max_kv_seq_len = infer_state.max_kv_seq_len @@ -73,7 +73,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, softmax_scale, get_sm_count=True, - **run_config + **run_config, ) if not hasattr(infer_state, "decode_att_block_seq"): @@ -122,7 +122,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, softmax_scale, get_sm_count=False, - **run_config + **run_config, ) flash_decode_stage2( @@ -132,7 +132,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), - **run_config + **run_config, ) return o_tensor diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index ca050e911b..b691d986cc 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -22,8 +22,13 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state @@ -37,11 +42,13 @@ def create_full_att_states( config=PackedMLALogicalConfig( logical_layer_groups=logical_layer_groups, latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, format="packed_group128_fp8" if args.llm_kv_type == "fp8kv_dsa" else "packed_bf16", indexer_dim=config.get("indexer_dim", 128), ), - page_size=args.page_size, device="meta", + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/deepseek_v4/layer_infer/indexer.py b/lightllm/models/deepseek_v4/layer_infer/indexer.py index 6c0d8b14aa..177c5529d6 100644 --- a/lightllm/models/deepseek_v4/layer_infer/indexer.py +++ b/lightllm/models/deepseek_v4/layer_infer/indexer.py @@ -75,7 +75,8 @@ def write_indexer_k( infer_state.mem_index.reshape(-1), infer_state.position_ids, infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k, - infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size // self.compress_ratio, + infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size + // self.compress_ratio, ) def build_metadata( @@ -171,7 +172,10 @@ def _c4_indices(self, infer_state: DeepseekV4InferStateInfo, idx_q_fp8, weights) return slots.unsqueeze(1), lengths paged_meta = infer_state.get_c4_indexer_metadata(c4_cap) - page_size = infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size // self.compress_ratio + page_size = ( + infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size + // self.compress_ratio + ) packed = infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k indexer_k_cache = packed.view( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index 9a296328ac..76758da0cc 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -105,7 +105,9 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list[FullAttentionState]: """分别创建 C4/C128 定义；逻辑索引保留原始 token 域，不提前分配存储。""" - from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import create_compressed_att_state + from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import ( + create_compressed_att_state, + ) from .config import create_attention_config assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" @@ -119,11 +121,13 @@ def create_full_att_states( for ratio in (4, 128): if ratio not in config.compress_ratios: continue - states.append(create_compressed_att_state( - config=config.create_compressed_config(compress_ratio=ratio), - page_size=args.page_size, - device="meta", - )) + states.append( + create_compressed_att_state( + config=config.create_compressed_config(compress_ratio=ratio), + page_size=args.page_size, + device="meta", + ) + ) return states @classmethod @@ -151,39 +155,45 @@ def create_bounded_att_states( tp_rank=tp_rank, tp_world_size=tp_world_size, ) - states = [create_sliding_window_state( - config=config.create_swa_config(), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - )] + states = [ + create_sliding_window_state( + config=config.create_swa_config(), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ] # 模型决定窗口组合：CSA 使用 W8/S4，HCA 使用 W128/S128；State 本身不识别 C4/C128。 for ratio, window_size in ((4, 8), (128, 128)): layers = tuple( - layer for layer, compress_ratio in zip(model_logical_layer_ids, config.compress_ratios) + layer + for layer, compress_ratio in zip(model_logical_layer_ids, config.compress_ratios) if compress_ratio == ratio ) if not layers: continue targets = ( (("kv", config.head_dim), ("indexer", config.indexer_head_dim)) - if ratio == 4 else (("kv", config.head_dim),) + if ratio == 4 + else (("kv", config.head_dim),) ) for target, head_dim in targets: - states.append(create_window_compressor_state( - config=WindowCompressorStateConfig( - logical_layer_ids=layers, - window_size=window_size, - compress_stride=ratio, - candidate_dim=head_dim * (2 if window_size > ratio else 1), - target=target, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - )) + states.append( + create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=layers, + window_size=window_size, + compress_stride=ratio, + candidate_dim=head_dim * (2 if window_size > ratio else 1), + target=target, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) return states def prepare_attention_states(self) -> None: diff --git a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py index 0f9365b1b1..63d936afc0 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py +++ b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py @@ -33,9 +33,7 @@ def get_swa_token_indices(swa_state, req_idx, position_ids, *, prefill_state=Non hold_slot = swa_state.HOLD_SWA_SLOT if prefill_state is not None: ready = prefill_state.prefill_token_ranges[req_idx, 0] - first = ( - (ready - config.window_size + 1).clamp_min(0) // config.page_token_size * config.page_token_size - ) + first = (ready - config.window_size + 1).clamp_min(0) // config.page_token_size * config.page_token_size workspace = prefill_state.swa_prefill_workspace_offsets[req_idx] swa_token_indices = workspace + position_ids - first hold_slot = prefill_state.HOLD_SWA_SLOT @@ -104,7 +102,11 @@ def _copy_swa_prefix_to_workspace_kernel( return source_slot = req_idx * SWA_CAPACITY_TOKENS + position % SWA_CAPACITY_TOKENS source = runtime_swa_buffer + layer * RUNTIME_LAYER_STRIDE + source_slot // SWA_PAGE_TOKEN_NUM * PACKED_BYTES - target = workspace_swa_buffer + layer * WORKSPACE_LAYER_STRIDE + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES + target = ( + workspace_swa_buffer + + layer * WORKSPACE_LAYER_STRIDE + + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES + ) lanes = tl.arange(0, BLOCK) for chunk in range(0, PACKED_BYTES, BLOCK): offset = chunk + lanes diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index 15dc1f37f0..1267979882 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -145,8 +145,13 @@ def _init_custom(self): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -161,9 +166,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=256, dtype=torch.bfloat16, - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index dd87a1a505..4f998ea25e 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -210,12 +210,16 @@ def _init_to_get_rotary_gemma4(self): sliding_partial = sliding_params.get("partial_rotary_factor", 1.0) sliding_rot_dim = int(sliding_head_dim * sliding_partial) inv_freq_sliding = 1.0 / ( - sliding_theta ** (torch.arange( - start=0, - end=sliding_rot_dim, - step=2, - dtype=torch.float32, - ) / sliding_rot_dim) + sliding_theta + ** ( + torch.arange( + start=0, + end=sliding_rot_dim, + step=2, + dtype=torch.float32, + ) + / sliding_rot_dim + ) ) freqs_s = torch.outer(t, inv_freq_sliding) self._cos_cached_sliding = torch.cos(freqs_s).to(self.data_type).cuda() @@ -235,12 +239,16 @@ def _init_to_get_rotary_gemma4(self): if rope_type == "proportional": rope_angles = int(full_partial * full_head_dim // 2) inv_freq_rot = 1.0 / ( - full_theta ** (torch.arange( - start=0, - end=2 * rope_angles, - step=2, - dtype=torch.float32, - ) / full_head_dim) + full_theta + ** ( + torch.arange( + start=0, + end=2 * rope_angles, + step=2, + dtype=torch.float32, + ) + / full_head_dim + ) ) nope_angles = full_head_dim // 2 - rope_angles if nope_angles > 0: @@ -249,12 +257,18 @@ def _init_to_get_rotary_gemma4(self): inv_freq_full = inv_freq_rot else: full_rot_dim = int(full_head_dim * full_partial) - inv_freq_full = 1.0 / (full_theta ** (torch.arange( - start=0, - end=full_rot_dim, - step=2, - dtype=torch.float32, - ) / full_rot_dim)) + inv_freq_full = 1.0 / ( + full_theta + ** ( + torch.arange( + start=0, + end=full_rot_dim, + step=2, + dtype=torch.float32, + ) + / full_rot_dim + ) + ) freqs_f = torch.outer(t, inv_freq_full) self._cos_cached_full = torch.cos(freqs_f).to(self.data_type).cuda() diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index 3abaf00466..e16d08cfce 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -43,8 +43,13 @@ def _verify_params(self): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -60,9 +65,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config["hidden_size"] // config["num_attention_heads"], dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, replicated=True, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 67ade58148..5b5235c76e 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -53,46 +53,56 @@ def create_bounded_att_states( from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) - assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" + assert not ( + args.mtp_step and args.enable_prefill_decode_mixed + ), "MTP linear attention does not support mixed prefill/decode" cfg = Glm5NextResourceConfig.from_model_config(config, args) states = [] - linear_layers = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval) + linear_layers = tuple( + layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval + ) if linear_layers: - states.append(create_kda_state( - config=ConvSSMStateConfig( - logical_layer_ids=linear_layers, - global_k_heads=cfg.global_linear_k_heads, - global_v_heads=cfg.global_linear_v_heads, - key_dim=cfg.head_linear_k_dim, - value_dim=cfg.head_linear_v_dim, - conv_width=cfg.conv_kernel_size - 1, - conv_dtype=cfg.conv_state_dtype, - ssm_dtype=cfg.ssm_state_dtype, - tp_rank=tp_rank, - tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - )) + states.append( + create_kda_state( + config=ConvSSMStateConfig( + logical_layer_ids=linear_layers, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) if full_layers: - assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" - states.append(create_window_compressor_state( - config=WindowCompressorStateConfig( - logical_layer_ids=full_layers, - window_size=config["index_kpool"], - compress_stride=config["index_kpool"], - candidate_dim=config["index_head_dim"], - target="indexer", - dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, - tp_world_size=tp_world_size, - ), - request_capacity=request_capacity, - mtp_step=args.mtp_step, - device=device, - )) + assert ( + args.linear_att_hash_page_size % config["index_kpool"] == 0 + ), "GLM cache pages must align to index_kpool" + states.append( + create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=full_layers, + window_size=config["index_kpool"], + compress_stride=config["index_kpool"], + candidate_dim=config["index_head_dim"], + target="indexer", + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) return states @classmethod @@ -159,9 +169,9 @@ def create_full_att_states( latent_dim=cfg.full_att_head_dim, tp_rank=tp_rank, tp_world_size=tp_world_size, - format="packed_group64_e8m0" if args.llm_kv_type == "fp8kv_dsa" else ( - "packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16" - ), + format="packed_group64_e8m0" + if args.llm_kv_type == "fp8kv_dsa" + else ("packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16"), ), page_size=args.page_size, device="meta", @@ -171,7 +181,9 @@ def create_full_att_states( def _init_att_backend(self): from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend - backend_class = Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + backend_class = ( + Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + ) self.prefill_att_backend = backend_class(model=self) self.decode_att_backend = self.prefill_att_backend diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 37663e4fb9..a59c6f1ab1 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -95,9 +95,9 @@ def create_full_att_states( latent_dim=config["kv_lora_rank"], tp_rank=tp_rank, tp_world_size=tp_world_size, - format="packed_group64_e8m0" if args.llm_kv_type == "fp8kv_dsa" else ( - "packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16" - ), + format="packed_group64_e8m0" + if args.llm_kv_type == "fp8kv_dsa" + else ("packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16"), ), page_size=args.page_size, device="meta", @@ -116,7 +116,10 @@ def create_bounded_att_states( device: str = "meta", ) -> list["BoundedAttentionState"]: """直接选择本模型的 K-pool 窗口层，不依赖 full State 创建过程。""" - from lightllm.common.attention_state.bounded_att import create_window_compressor_state, WindowCompressorStateConfig + from lightllm.common.attention_state.bounded_att import ( + create_window_compressor_state, + WindowCompressorStateConfig, + ) from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 2bec35f191..0b3bb2e867 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -79,7 +79,11 @@ def _token_attention_kernel( _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( - q=_q, k=_k, v=_v, att_control=AttControl(logical_layer_id=self.layer_num_), alloc_func=self.alloc_tensor, + q=_q, + k=_k, + v=_v, + att_control=AttControl(logical_layer_id=self.layer_num_), + alloc_func=self.alloc_tensor, ) return o_tensor.view(q.shape) diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index c89567b3f1..13036b0b83 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -41,8 +41,13 @@ def __init__(self, kvargs): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -58,9 +63,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 26e8fd1ffe..f906270d5c 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -35,8 +35,13 @@ def __init__(self, kvargs): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -52,9 +57,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index e2c298967c..b749698baa 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -43,8 +43,13 @@ def _init_config(self): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -60,9 +65,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 0604d063e5..f1a3a65748 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -55,10 +55,14 @@ def create_bounded_att_states( config = cls.normalize_model_config(model_config) assert config["num_attention_heads"] % tp_world_size == 0 - assert not (args.mtp_step and args.enable_prefill_decode_mixed), "MTP linear attention does not support mixed prefill/decode" + assert not ( + args.mtp_step and args.enable_prefill_decode_mixed + ), "MTP linear attention does not support mixed prefill/decode" cfg = Qwen3NextResourceConfig.from_model_config(config, args) full_layer_indices = set(cfg.get_full_att_layer_indices()) - logical_layer_ids = tuple(layer for i, layer in enumerate(model_logical_layer_ids) if i not in full_layer_indices) + logical_layer_ids = tuple( + layer for i, layer in enumerate(model_logical_layer_ids) if i not in full_layer_indices + ) if not logical_layer_ids: return [] state = create_gdn_state( diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index 5ca086b03e..39a90415bc 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -31,8 +31,13 @@ def __init__(self, kvargs): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -48,10 +53,13 @@ def create_full_att_states( global_head_num=1, head_dim=config["hidden_size"] // config["num_attention_heads"], dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, replicated=True, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 66d5ebb4a5..50fc740e04 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -34,8 +34,13 @@ def __init__(self, kvargs): @classmethod def create_full_att_states( - cls, model_config: dict[str, object], args: "StartArgs | EasyDict", *, - model_logical_layer_ids: tuple[int, ...], tp_rank: int = 0, tp_world_size: int = 1, + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, ) -> list["FullAttentionState"]: """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -51,9 +56,12 @@ def create_full_att_states( global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), dtype=get_torch_dtype(args.data_type), - tp_rank=tp_rank, tp_world_size=tp_world_size, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ), - kv_type=args.llm_kv_type, page_size=args.page_size, device="meta", + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) return [state] diff --git a/lightllm/server/router/dynamic_prompt/radix_cache.py b/lightllm/server/router/dynamic_prompt/radix_cache.py index afc7df82af..41a292a8a0 100644 --- a/lightllm/server/router/dynamic_prompt/radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/radix_cache.py @@ -119,9 +119,7 @@ def __init__(self, total_token_num, rank_in_node, att_state_manager=None, page_s raise ValueError(f"page_size must be >= 1, got {page_size}") allocator = None if att_state_manager is None else att_state_manager.full_att_allocator if allocator is not None and page_size != allocator.page_size: - raise ValueError( - f"RadixCache page_size {page_size} must match allocator page_size {allocator.page_size}" - ) + raise ValueError(f"RadixCache page_size {page_size} must match allocator page_size {allocator.page_size}") self.page_size = page_size self.root_node = TreeNode(page_size=page_size) diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index b2de4fa93b..8513fc030a 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -235,8 +235,7 @@ async def loop_for_fwd( for dp_index in range(self.dp_size_in_node): token_ratio1 = self.get_used_tokens(dp_index) / self.max_total_token_num token_ratio2 = ( - self.max_total_token_num - - self.kv_capacity_reader.get_available_tokens(dp_index) + self.max_total_token_num - self.kv_capacity_reader.get_available_tokens(dp_index) ) / self.max_total_token_num d_i = dp_index estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(d_i) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 0c911b3799..056c36f596 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -180,9 +180,7 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer req.tail_small_page_buffer_id = None # 释放请求执行期间申请、但不再插入 radix cache 的大页状态。 if req.hybrid_len_to_big_page_id: - self.att_state_manager.bounded_att_big_page_allocator.free( - list(req.hybrid_len_to_big_page_id.values()) - ) + self.att_state_manager.bounded_att_big_page_allocator.free(list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() # 解除请求对已命中 GPU radix cache 前缀节点的引用。 @@ -296,9 +294,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): # big page 模式下请求在 prefill 跨过 big page 边界后、到达末尾前被 pause / abort。 # 若不释放，会泄漏 big page state slot，并触发 free_a_req_mem 中 dict 为空的断言。 if req.hybrid_len_to_big_page_id: - self.att_state_manager.bounded_att_big_page_allocator.free( - list(req.hybrid_len_to_big_page_id.values()) - ) + self.att_state_manager.bounded_att_big_page_allocator.free(list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() req.cur_kv_len = shared_kv_len diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 9351a7739f..7997f49934 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -154,7 +154,9 @@ def init_model(self, kvargs): models = (self.model, *self.draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" - assert all(item.max_req_num == self.model.max_req_num for item in models), "models must share the request domain" + assert all( + item.max_req_num == self.model.max_req_num for item in models + ), "models must share the request domain" # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): @@ -240,13 +242,17 @@ def init_model(self, kvargs): # PD 和 DP 使用独立导出，避免 PD 读者收到大小页池或 DP 读者丢失页池。 if self.args.run_mode in ("prefill", "decode"): publish_att_resources( - self.model.att_state_manager, self.model.req_manager.req_to_token_indexs, - kind="pd", reader_count=self.node_world_size, + self.model.att_state_manager, + self.model.req_manager.req_to_token_indexs, + kind="pd", + reader_count=self.node_world_size, ) if self.args.enable_dp_prompt_cache_fetch: publish_att_resources( - self.model.att_state_manager, self.model.req_manager.req_to_token_indexs, - kind="dp", reader_count=self.dp_size_in_node - 1, + self.model.att_state_manager, + self.model.req_manager.req_to_token_indexs, + kind="dp", + reader_count=self.dp_size_in_node - 1, ) dist.barrier(group=self.node_nccl_group) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 2fadb6b3b5..d88f5af2eb 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -23,7 +23,8 @@ def __init__(self, max_req_num: int, dp_size_in_node: int, backend): self.big_page_tokens = ( backend.args.linear_att_hash_page_size * backend.args.linear_att_page_block_num - if backend.is_hybrid_att_model else 0 + if backend.is_hybrid_att_model + else 0 ) max_big_pages = backend.args.max_req_total_len // self.big_page_tokens if self.big_page_tokens else 0 # 原共享信息只保存 KV 长度和请求槽位，大页列表使用独立共享数组。 @@ -60,7 +61,7 @@ def fill_reqs_info(self, reqs: List[InferReq]) -> None: # radix 页按路径从根到叶排列；请求持有的页按其前缀长度放到对应位置。 if req.shared_kv_node is not None: shared_pages = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - page_ids[:len(shared_pages)] = shared_pages + page_ids[: len(shared_pages)] = shared_pages for token_len, page_id in req.hybrid_len_to_big_page_id.items(): if token_len <= req.cur_kv_len: page_ids[token_len // self.big_page_tokens - 1] = page_id @@ -107,7 +108,9 @@ def build_shared_kv_trans_tasks( src_big_page_ids = self.shared_req_big_page_ids.arr[ req_index, src_dp_rank, first_page:end_page ].tolist() - assert all(page_id >= 0 for page_id in src_big_page_ids), "DP source is missing prefix big pages" + assert all( + page_id >= 0 for page_id in src_big_page_ids + ), "DP source is missing prefix big pages" allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" @@ -116,7 +119,7 @@ def build_shared_kv_trans_tasks( # 只搬运有效 KV；页尾预留槽位已写入请求表，不参与本次传输。 dst_token_indices = dst_token_indices[:trans_size] src_token_indices = self.backend.dp_req_to_token_indices[src_dp_rank][ - src_req_idx, req.cur_kv_len:target_kv_len + src_req_idx, req.cur_kv_len : target_kv_len ] trans_tasks.append( TransTask( diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 67ff390a63..ed7e5bb167 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -98,7 +98,9 @@ def _init_reqs(self, reqs: List[Tuple]): infer_reqs = g_infer_context.add_reqs(reqs, init_prefix_cache=True) req_dp_ranks = [req[3] for req in reqs] self.dp_kv_shared_module.fill_reqs_info(reqs=infer_reqs) - transfer_tasks = self.dp_kv_shared_module.build_shared_kv_trans_tasks(reqs=infer_reqs, req_dp_ranks=req_dp_ranks) + transfer_tasks = self.dp_kv_shared_module.build_shared_kv_trans_tasks( + reqs=infer_reqs, req_dp_ranks=req_dp_ranks + ) self.dp_kv_shared_module.kv_trans(trans_tasks=transfer_tasks) # other_dp_reqs 只是为本 DP 做完 prefix cache / KV 拉取后的临时本地对象， diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index e7a1045074..261f090b0b 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -160,10 +160,12 @@ def read_req_from_cpu( # FA3 的资源搬运需等待已有计算流完成。 torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) assert 0 <= start_token_len <= req.cur_kv_len, "invalid CPU page start" - token_indices = torch.cat([ - req_manager.req_to_token_indexs[req.req_idx, start_token_len:req.cur_kv_len], - mem_indexes.cuda(non_blocking=True), - ]) + token_indices = torch.cat( + [ + req_manager.req_to_token_indexs[req.req_idx, start_token_len : req.cur_kv_len], + mem_indexes.cuda(non_blocking=True), + ] + ) assert token_indices.numel() == page_lengths[-1] - start_token_len # 阶段二：只为完整 bounded 页分配普通大页，碎页 LOAD 槽位由资源层管理。 @@ -171,9 +173,9 @@ def read_req_from_cpu( if self.backend.is_hybrid_att_model: page_tokens = self.args.cpu_cache_token_page_size complete_page_lengths = [length for length in page_lengths if length % page_tokens == 0] - assert len(complete_page_lengths) == len(page_lengths) - int(page_lengths[-1] % page_tokens != 0), ( - "only the final CPU page may be partial" - ) + assert len(complete_page_lengths) == len(page_lengths) - int( + page_lengths[-1] % page_tokens != 0 + ), "only the final CPU page may be partial" allocator = resources.bounded_att_big_page_allocator big_page_slot_ids = allocator.alloc(len(complete_page_lengths)) assert big_page_slot_ids is not None, "not enough bounded state slots for CPU cache load" @@ -248,9 +250,7 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L return true_finished_reqs - def write_req_to_cpu( - self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream - ) -> Optional["TransTask"]: + def write_req_to_cpu(self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream) -> Optional["TransTask"]: """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers @@ -324,7 +324,9 @@ def write_req_to_cpu( if self.backend.is_hybrid_att_model: page_tokens = self.args.cpu_cache_token_page_size radix_cache = g_infer_context.radix_cache - shared_slots = radix_cache.get_big_page_ids_by_node(req.shared_kv_node) if radix_cache is not None else [] + shared_slots = ( + radix_cache.get_big_page_ids_by_node(req.shared_kv_node) if radix_cache is not None else [] + ) big_slots = {(i + 1) * page_tokens: slot for i, slot in enumerate(shared_slots)} big_slots.update(req.hybrid_len_to_big_page_id) big_page_slot_ids = [ diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py index 54fd9bc119..516744de25 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -53,9 +53,9 @@ def copy_page( state_layouts = pd_operator.get_pd_bounded_att_page_layouts() case _: assert False, f"unknown PD task kind {task_kind}" - assert state_layouts and pd_operator.get_page_bytes() == pd_page_buffers.shape[1], ( - "PD page layout does not match this group" - ) + assert ( + state_layouts and pd_operator.get_page_bytes() == pd_page_buffers.shape[1] + ), "PD page layout does not match this group" tp_transfer_layouts.append(state_layouts) for att_manager, state_layouts in zip(tp_att_managers, tp_transfer_layouts): diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py index 1d4504015a..46a07e46f9 100644 --- a/lightllm/server/router/model_infer/shared_att_resources_utils.py +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -43,7 +43,9 @@ def publish_att_resources( payloads = [ForkingPickler.dumps((ipc_resources, request_mapping)).tobytes() for _ in range(count)] size = max(len(payload) for payload in payloads) shm = create_or_link_shm( - name=f"rank_storage_{kind}_{get_current_rank_in_node()}", expected_size=8 + count * size, force_mode="create" + name=f"rank_storage_{kind}_{get_current_rank_in_node()}", + expected_size=8 + count * size, + force_mode="create", ) shm.buf[:4] = count.to_bytes(4, "little") shm.buf[4:8] = size.to_bytes(4, "little") @@ -54,7 +56,8 @@ def publish_att_resources( def load_att_resources( - rank_in_node: int, kind: Literal["pd", "dp"], + rank_in_node: int, + kind: Literal["pd", "dp"], ) -> tuple[AttentionStateManager, torch.Tensor]: """读取指定用途的资源；DP 接收方随后通过 manager 统一注册共享页池。""" assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py index abcdbd0650..05c90e17f6 100644 --- a/lightllm/utils/cpu_cache_layout.py +++ b/lightllm/utils/cpu_cache_layout.py @@ -38,7 +38,9 @@ def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: count = config["num_hidden_layers"] logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) full_states = model_class.create_full_att_states( - model_config=config, args=args, model_logical_layer_ids=logical_layer_ids, + model_config=config, + args=args, + model_logical_layer_ids=logical_layer_ids, ) bounded_states = model_class.create_bounded_att_states( model_config=config, diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index ec89a652ba..f7720e64fb 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -64,9 +64,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": page_bytes=page_bytes, ) - cpu_cache_page_num = int( - (args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.page_bytes) - ) + cpu_cache_page_num = int((args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.page_bytes)) cpu_cache_meta.page_num = cpu_cache_page_num logger.info(f"cpu kv cache page num: {cpu_cache_meta.page_num}") diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index d8beb27faf..4e4317f3fa 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -105,7 +105,9 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP next_token = output.logits.argmax(dim=-1).cpu() source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() - cpu = torch.zeros((len(pages), resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + cpu = torch.zeros( + (len(pages), resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) resources.get_cpu_cache_operator().write_req( token_indices=source_indices, cpu_page_buffers=cpu, @@ -119,7 +121,9 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP runtime_layouts = resources.get_pd_operator().get_pd_bounded_att_page_layouts() pd = {"kv": torch.empty((1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda")} if runtime_layouts: - pd["att_state"] = torch.empty((1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + pd["att_state"] = torch.empty( + (1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda" + ) resources.get_pd_operator().write_page( pd_page_buffers=pd["kv"], pd_page_index=0, @@ -148,8 +152,10 @@ def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpP resources.copy_full_att_tokens(source_indices, indices) if hybrid: resources.restore_bounded_att_runtime_from_pages( - req_indices=request.tolist(), state_token_lens=[prefix_len], - page_kind="big", page_slot_ids=[source_slots[-1]], + req_indices=request.tolist(), + state_token_lens=[prefix_len], + page_kind="big", + page_slot_ids=[source_slots[-1]], ) elif kind == "cpu_pages": slots = big.alloc(len(pages)) if hybrid else None diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 5912d2bbda..59e1b66ade 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -63,7 +63,9 @@ def main(): resources = AttentionStateManager(request_capacity=1) bounded_state = create_gdn_state( config=ConvSSMStateConfig( - logical_layer_ids=tuple(layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval), + logical_layer_ids=tuple( + layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval + ), global_k_heads=cfg.global_linear_k_heads, global_v_heads=cfg.global_linear_v_heads, key_dim=cfg.head_linear_k_dim, @@ -88,7 +90,9 @@ def main(): resources.finish_state_registration() pages = list(range(args.pages)) slots = resources.bounded_att_big_page_allocator.alloc(args.pages) - cpu_page_buffers = torch.zeros((args.pages, resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + cpu_page_buffers = torch.zeros( + (args.pages, resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) # Populate every TP region through the production API before measuring restores. full.k_storage.zero_() full.v_storage.zero_() @@ -128,11 +132,16 @@ def restore(): + full.v_storage[:, :tokens].numel() * full.v_storage.element_size() ) pool = ( - resources.main_kda_conv_bf16_ssm_fp32_state if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") else resources.main_gdn_conv_bf16_ssm_fp32_state + resources.main_kda_conv_bf16_ssm_fp32_state + if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") + else resources.main_gdn_conv_bf16_ssm_fp32_state ).big_pages_pool state_bytes = args.pages * sum(t[0].numel() * t.element_size() for t in (pool.conv_state, pool.ssm_state)) median = statistics.median(times) - print(f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; CPU bytes/page: {resources.get_cpu_cache_operator().get_page_bytes():,}") + print( + f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; " + f"CPU bytes/page: {resources.get_cpu_cache_operator().get_page_bytes():,}" + ) print(f"Mean: {statistics.mean(times) * 1000:.3f} ms; median: {median * 1000:.3f} ms") print(f"KV + state payload throughput: {(token_bytes + state_bytes) / median / 1024**3:.3f} GiB/s") diff --git a/unit_tests/common/att_resources/conftest.py b/unit_tests/common/att_resources/conftest.py index 829c82ac3f..c074e861b3 100644 --- a/unit_tests/common/att_resources/conftest.py +++ b/unit_tests/common/att_resources/conftest.py @@ -36,7 +36,9 @@ def pd_args(monkeypatch: pytest.MonkeyPatch): from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(pd_kv_page_size=8, cpu_cache_token_page_size=8)))) + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(pd_kv_page_size=8, cpu_cache_token_page_size=8))) + ) get_env_start_args.cache_clear() yield get_env_start_args() get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index 36ff514ab6..ef44a54dff 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -34,7 +34,8 @@ def __init__(self, device: str | torch.device | int = "cpu"): latent_dim=self.main_gqa_fp32_state.config.head_dim, latent_dtype=self.main_gqa_fp32_state.config.dtype, ), - page_size=self.main_gqa_fp32_state.page_size, device=device, + page_size=self.main_gqa_fp32_state.page_size, + device=device, ) prepared_full_state.allocate_storage(self.main_gqa_fp32_state.capacity) self.register_state("draft0", prepared_full_state) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index b3966eb62c..10c52c8cea 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -11,10 +11,19 @@ from lightllm.common.attention_state.types import TensorField from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState from lightllm.common.attention_state.full_att import ( - FullAttentionState, GQAAttentionConfig, MLAConfig, - GroupQuantizedAttentionConfig, StaticFP8AttentionConfig, PackedMLAAttentionConfig, - CompressedAttentionConfig, FP8E8M0PackedMLAConfig, StaticFP8FullAttentionState, - GQAFullAttentionState, GroupQuantizedFullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, + FullAttentionState, + GQAAttentionConfig, + MLAConfig, + GroupQuantizedAttentionConfig, + StaticFP8AttentionConfig, + PackedMLAAttentionConfig, + CompressedAttentionConfig, + FP8E8M0PackedMLAConfig, + StaticFP8FullAttentionState, + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, ) @@ -46,7 +55,9 @@ def bind_test_fp8_calibration( q_scales = q_scales.repeat(1, config.tp_world_size) calibration["num_head"] = heads * config.tp_world_size calibration["q_calibration"] = { - "num_head": q_scales.shape[1], "scales_shape": list(q_scales.shape), "scales": q_scales.tolist(), + "num_head": q_scales.shape[1], + "scales_shape": list(q_scales.shape), + "scales": q_scales.tolist(), } calibration["scales_shape"] = list(scales.shape) calibration["scales"] = scales.tolist() @@ -56,9 +67,14 @@ def bind_test_fp8_calibration( def rebuild_full_state( template: FullAttentionState, config: ( - GQAAttentionConfig | MLAConfig | GroupQuantizedAttentionConfig - | StaticFP8AttentionConfig | PackedMLAAttentionConfig - | CompressedAttentionConfig | FP8E8M0PackedMLAConfig | None + GQAAttentionConfig + | MLAConfig + | GroupQuantizedAttentionConfig + | StaticFP8AttentionConfig + | PackedMLAAttentionConfig + | CompressedAttentionConfig + | FP8E8M0PackedMLAConfig + | None ) = None, token_capacity: int | None = None, page_size: int | None = None, @@ -73,7 +89,9 @@ def rebuild_full_state( device = template.device if device is None else device if isinstance(template, StaticFP8FullAttentionState): prepared_full_state = type(template)( - config, page_size=page_size, device=device, + config, + page_size=page_size, + device=device, ) bind_test_fp8_calibration( state=prepared_full_state, @@ -83,7 +101,9 @@ def rebuild_full_state( prepared_full_state.allocate_storage(token_capacity) return prepared_full_state prepared_full_state = type(template)( - config, page_size=page_size, device=device, + config, + page_size=page_size, + device=device, ) prepared_full_state.allocate_storage(token_capacity) return prepared_full_state @@ -128,15 +148,16 @@ def prepare_model_states(model: TpPartBaseModel, logical_layer_ids: Sequence[int def create_full_attention_roles( - template: FullAttentionState, main_layer_count: int, + template: FullAttentionState, + main_layer_count: int, ) -> tuple[FullAttentionState, FullAttentionState | None]: """仅用于测试：按配置层号创建各自拥有独立存储的 main 和单个 draft。""" layers = template.config.logical_layer_groups boundary = sum(all(layer < main_layer_count for layer in group) for group in layers) assert boundary and all(layer < main_layer_count for group in layers[:boundary] for layer in group) - assert all(layer >= main_layer_count for group in layers[boundary:] for layer in group), ( - "draft layers must follow main layers and cannot share main storage" - ) + assert all( + layer >= main_layer_count for group in layers[boundary:] for layer in group + ), "draft layers must follow main layers and cannot share main storage" states = [] for start, end in ((0, boundary), (boundary, len(layers))): if start == end: @@ -178,8 +199,10 @@ def full_state_test_buffers(state: FullAttentionState) -> dict[str, torch.Tensor """仅供多方案参数化测试枚举实际存储，检查容量、独立分配及字节一致性。""" if isinstance(state, GroupQuantizedFullAttentionState): return { - "k": state.k_storage, "v": state.v_storage, - "k_scale": state.k_scales, "v_scale": state.v_scales, + "k": state.k_storage, + "v": state.v_storage, + "k_scale": state.k_scales, + "v_scale": state.v_scales, } if isinstance(state, (GQAFullAttentionState, StaticFP8FullAttentionState)): return {"k": state.k_storage, "v": state.v_storage} @@ -188,7 +211,8 @@ def full_state_test_buffers(state: FullAttentionState) -> dict[str, torch.Tensor def full_state_test_fields( - state: FullAttentionState, buffers: dict[str, torch.Tensor] | None = None, + state: FullAttentionState, + buffers: dict[str, torch.Tensor] | None = None, ) -> dict[str, TensorField]: """测试中独立描述实际数据，用于校验不同存储方案的传输内容。""" buffers = full_state_test_buffers(state) if buffers is None else buffers @@ -199,15 +223,19 @@ def full_state_test_fields( if state.config.dtype == torch.float8_e4m3fn: fields = { "kv_nope": TensorField(tensor=raw[..., :512], axes=axes), - "kv_scale": TensorField(tensor=raw[..., 512:528].view(torch.float32), axes=("layer", "token", "head", "group")), + "kv_scale": TensorField( + tensor=raw[..., 512:528].view(torch.float32), axes=("layer", "token", "head", "group") + ), "kv_rope": TensorField(tensor=raw[..., 528:656].view(torch.bfloat16), axes=axes), } else: - fields = {"kv": TensorField(tensor=storage[..., :state.config.latent_dim], axes=axes)} + fields = {"kv": TensorField(tensor=storage[..., : state.config.latent_dim], axes=axes)} start = state.config.indexer_offset end = start + state.config.indexer_dim fields["indexer_k"] = TensorField(tensor=raw[..., start:end], axes=axes) - fields["indexer_scale"] = TensorField(tensor=raw[..., end:end + 4].view(torch.float32), axes=("layer", "token", "head", "scale")) + fields["indexer_scale"] = TensorField( + tensor=raw[..., end : end + 4].view(torch.float32), axes=("layer", "token", "head", "scale") + ) return fields return { name: TensorField(tensor=tensor, axes=("layer", "token", "head", "group") if name.endswith("_scale") else axes) @@ -256,7 +284,9 @@ def build_hybrid_resources( device="cuda", ) bounded_state.initialize_page_pools( - big_page_capacity=big_page_capacity, small_page_capacity=small_page_capacity, pin_memory=True, + big_page_capacity=big_page_capacity, + small_page_capacity=small_page_capacity, + pin_memory=True, ) manager.register_state("main", bounded_state) for role, state in (("main", full_state), ("draft0", draft_full_state)): @@ -293,11 +323,15 @@ def window_states( manager: AttentionStateManager, role: str | None = None, ) -> tuple[WindowCompressorState, ...]: - return tuple(state for state in manager.get_bounded_attention_states(role) if isinstance(state, WindowCompressorState)) + return tuple( + state for state in manager.get_bounded_attention_states(role) if isinstance(state, WindowCompressorState) + ) def linear_states( manager: AttentionStateManager, role: str | None = None, ) -> tuple[GDNState | KDAState, ...]: - return tuple(state for state in manager.get_bounded_attention_states(role) if isinstance(state, (GDNState, KDAState))) + return tuple( + state for state in manager.get_bounded_attention_states(role) if isinstance(state, (GDNState, KDAState)) + ) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index d5d5ace566..e3a09adb8a 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -153,12 +153,15 @@ def mark_backends(model: TpPartBaseModel): logical_layer_ids = {layer_id for group in state.config.logical_layer_groups for layer_id in group} for layer in item.layers_infer: if layer.layer_num_ in logical_layer_ids: - assert item.att_state_manager.get_state( - model_role=item.att_state_role, - state_type_filter="", - logical_layer_id=layer.layer_num_, - state_class=type(state), - ) is state + assert ( + item.att_state_manager.get_state( + model_role=item.att_state_role, + state_type_filter="", + logical_layer_id=layer.layer_num_, + state_class=type(state), + ) + is state + ) model.prefill_att_backend = model.decode_att_backend = object() model.prefill_att_backend1 = model.decode_att_backend1 = None model.backends_bound = True @@ -255,7 +258,9 @@ def start_runtime(): ) ) else: - assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" + assert all( + getattr(item, "att_state_manager", None) is None for item in models + ), "resources already initialized" req_manager = create_req_manager( max_request_num=models[0].max_req_num, max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), @@ -276,9 +281,7 @@ def start_runtime(): for index, item in enumerate(models): item.init_runtime() assert all(not getattr(other, "backends_bound", False) for other in models[index + 1 :]) - assert warmed == padded == captured == prefill_captured == { - id(other) for other in models[: index + 1] - } + assert warmed == padded == captured == prefill_captured == {id(other) for other in models[: index + 1]} start_runtime() assert request_allocations == [(2, 32)] @@ -294,7 +297,8 @@ def start_runtime(): assert len({state.conv_state.data_ptr() for state in draft_states}) == draft_count manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), role="main", + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), + role="main", ) assert all(not state.mtp_state_index.any() for state in draft_states) if main_linear: @@ -303,7 +307,8 @@ def start_runtime(): assert len(state.config.logical_layer_ids) == 1 manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([i + 2], dtype=torch.int32, device="cuda"), role=f"draft{i}", + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), + role=f"draft{i}", ) assert state.mtp_state_index[0] == i + 1 assert all(not other.mtp_state_index.any() for other in draft_states[i + 1 :]) @@ -330,7 +335,9 @@ def start_runtime(): assert torch.equal(state.q_scales[:, 0].cpu(), scales[start:end] * 2) layouts = build_cpu_cache_layouts() - assert (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 == manager.get_cpu_cache_operator().get_page_bytes() + assert ( + max(layout.byte_end for layout in layouts.values()) + 15 + ) // 16 * 16 == manager.get_cpu_cache_operator().get_page_bytes() assert layouts.keys() == manager.get_cpu_cache_operator().get_page_layouts().keys() for state_name, layout in layouts.items(): runtime_layout = manager.get_cpu_cache_operator().get_page_layouts()[state_name] diff --git a/unit_tests/common/att_resources/test_bounded_copy_kernel.py b/unit_tests/common/att_resources/test_bounded_copy_kernel.py index 243fb1fa00..22d7ab470a 100644 --- a/unit_tests/common/att_resources/test_bounded_copy_kernel.py +++ b/unit_tests/common/att_resources/test_bounded_copy_kernel.py @@ -37,10 +37,10 @@ def test_aggregate_same_dtype_and_strides(direction: str, grid_num: int, dtype: device=dst_device, pin_memory=dst_device == "cpu", ) - source_view = source[..., :shape[3] - 1] - target_view = target.permute(1, 0, 2, 3)[:, :, 1:shape[2] + 1, 2:shape[3] + 1] + source_view = source[..., : shape[3] - 1] + target_view = target.permute(1, 0, 2, 3)[:, :, 1 : shape[2] + 1, 2 : shape[3] + 1] expected = target.cpu().clone() - expected.permute(1, 0, 2, 3)[:, :, 1:shape[2] + 1, 2:shape[3] + 1].copy_(source_view.cpu()) + expected.permute(1, 0, 2, 3)[:, :, 1 : shape[2] + 1, 2 : shape[3] + 1].copy_(source_view.cpu()) sources.append(source) targets.append(target) copies.append((source_view, target_view)) @@ -72,12 +72,15 @@ def test_registered_shared_pages_and_stream_dependency() -> None: assert torch.equal(target, values.cpu()) -@pytest.mark.parametrize("float_dtype,int_dtype,patterns", [ - (torch.float16, torch.int16, [0x7E01, -32768, 0x7C00, 1]), - (torch.bfloat16, torch.int16, [0x7FC1, -32768, 0x7F80, 1]), - (torch.float32, torch.int32, [0x7FC00001, -2147483648, 0x7F800000, 1]), - (torch.float64, torch.int64, [0x7FF8000000000001, -9223372036854775808, 0x7FF0000000000000, 1]), -]) +@pytest.mark.parametrize( + "float_dtype,int_dtype,patterns", + [ + (torch.float16, torch.int16, [0x7E01, -32768, 0x7C00, 1]), + (torch.bfloat16, torch.int16, [0x7FC1, -32768, 0x7F80, 1]), + (torch.float32, torch.int32, [0x7FC00001, -2147483648, 0x7F800000, 1]), + (torch.float64, torch.int64, [0x7FF8000000000001, -9223372036854775808, 0x7FF0000000000000, 1]), + ], +) def test_exact_float_bit_patterns(float_dtype: torch.dtype, int_dtype: torch.dtype, patterns: list[int]) -> None: # 包含 NaN payload、负零和无穷，等宽整数复制不能触发精度转换或规范化。 bits = torch.tensor(data=patterns, dtype=int_dtype, device="cuda") @@ -152,7 +155,9 @@ def __getitem__(self, grid: tuple[int]): for _ in range(3): shape = (2,) * (ndim - 1) + (517,) storage_shape = (4,) * (ndim - 1) + (1034,) - source = torch.arange(torch.Size(storage_shape).numel(), dtype=torch.float32, device="cuda").reshape(storage_shape) + source = torch.arange(torch.Size(storage_shape).numel(), dtype=torch.float32, device="cuda").reshape( + storage_shape + ) # 每个轴都保留间隙，保证降维后仍需调用对应的 1～6 维算子。 source_view = source[(slice(None, None, 2),) * ndim] target = torch.full(size=(*shape[:-1], 1034), fill_value=-17, dtype=source.dtype, pin_memory=True) @@ -170,7 +175,21 @@ def __getitem__(self, grid: tuple[int]): assert torch.equal(target, expected) -@pytest.mark.parametrize("layout", ["ssm", "conv_mtp", "dense", "width4", "middle_unaligned", "prefix", "singletons", "wide_singleton", "wide_singleton_prefix", "all_singletons"]) +@pytest.mark.parametrize( + "layout", + [ + "ssm", + "conv_mtp", + "dense", + "width4", + "middle_unaligned", + "prefix", + "singletons", + "wide_singleton", + "wide_singleton_prefix", + "all_singletons", + ], +) def test_transfer_views_preserve_gaps_and_alignment(layout: str, monkeypatch: pytest.MonkeyPatch) -> None: from lightllm.common.attention_state.bounded_att import copy_kernel @@ -185,6 +204,7 @@ def __getitem__(self, grid: tuple[int]): def launch(**kwargs): launches.append((self.ndim, kwargs["SHAPE"], kwargs["DTYPE"])) return self.kernel[grid](**kwargs) + return launch for ndim in range(1, 7): @@ -246,12 +266,14 @@ def launch(**kwargs): copies = [] # 首尾地址都满足 8 字节对齐，中间请求偏移不对齐；整批必须保留 2 字节搬运。 for src_offset, dst_offset in ((0, 0), (9, 16), (24, 32)): - copies.append((source[src_offset:src_offset + 8], target[dst_offset:dst_offset + 8])) - expected[dst_offset:dst_offset + 8].copy_(source[src_offset:src_offset + 8].cpu()) + copies.append((source[src_offset : src_offset + 8], target[dst_offset : dst_offset + 8])) + expected[dst_offset : dst_offset + 8].copy_(source[src_offset : src_offset + 8].cpu()) expected_launch = (1, (8,), tl.uint16) else: shape = (4, 2, 3, 128) if layout == "dense" else (4, 6) - source = torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").to(torch.bfloat16).view(shape) + source = ( + torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").to(torch.bfloat16).view(shape) + ) target = torch.full(size=shape, fill_value=-1, dtype=source.dtype, pin_memory=True) copies = [(source[index], target[index]) for index in (0, 2)] expected = target.clone() diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index 9322b32867..ffd59db4aa 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -15,7 +15,9 @@ ConvSSMStateConfig, ) from lightllm.common.attention_state.types import ( - AxisMapping, StateTransferLayout, TransferRegion, + AxisMapping, + StateTransferLayout, + TransferRegion, ) if TYPE_CHECKING: @@ -24,6 +26,7 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for aggregated copies") + def make_state( cls: type[KDAState] | type[GDNState] = KDAState, device: str | torch.device | int = "cpu", @@ -369,13 +372,15 @@ def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp target.big_pages_pool.conv_state.fill_(91) target.big_pages_pool.ssm_state.fill_(92) for req in (0, 2): - source.conv_state[:, req, :, :source.config.conv_width].fill_(10 + req) + source.conv_state[:, req, :, : source.config.conv_width].fill_(10 + req) source.ssm_state[:, req * (src_mtp_step + 1)].fill_(20 + req) src_conv = source.conv_state.clone() src_ssm = source.ssm_state.clone() target.copy_runtime_state_from_dp( - src_req_indices=[2, 0], dst_req_indices=[0, 1], remote_dp_bounded_state=source, + src_req_indices=[2, 0], + dst_req_indices=[0, 1], + remote_dp_bounded_state=source, ) for src_req, dst_req in ((2, 0), (0, 1)): @@ -383,10 +388,10 @@ def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp assert (target.conv_state[:, dst_req, :, 2:] == -1).all() dst_slot = dst_req * (dst_mtp_step + 1) assert (target.ssm_state[:, dst_slot] == 20 + src_req).all() - assert (target.ssm_state[:, dst_slot + 1:dst_slot + dst_mtp_step + 1] == -2).all() + assert (target.ssm_state[:, dst_slot + 1 : dst_slot + dst_mtp_step + 1] == -2).all() assert target.mtp_state_index[dst_req] == 0 assert (target.conv_state[:, 2] == -1).all() - assert (target.ssm_state[:, 2 * (dst_mtp_step + 1):] == -2).all() + assert (target.ssm_state[:, 2 * (dst_mtp_step + 1) :] == -2).all() assert target.mtp_state_index[2] == 1 assert (target.big_pages_pool.conv_state == 91).all() assert (target.big_pages_pool.ssm_state == 92).all() @@ -412,7 +417,9 @@ def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): src_ssm = source.big_pages_pool.ssm_state.clone() target.copy_big_pages_from_dp( - src_page_ids=[2, 0], dst_page_ids=[0, 1], remote_dp_bounded_state=source, + src_page_ids=[2, 0], + dst_page_ids=[0, 1], + remote_dp_bounded_state=source, ) torch.cuda.synchronize() diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index 5d4425663c..1784d06cb8 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -62,8 +62,10 @@ def definition(kind: str): device="meta", ) kv_type = { - "int4": "int4kv", "int8": "int8kv", - "static_head": "fp8kv_sph", "static_tensor": "fp8kv_spt", + "int4": "int4kv", + "int8": "int8kv", + "static_head": "fp8kv_sph", + "static_tensor": "fp8kv_spt", }.get(kind, "None") state = create_gqa_state( config=GQAAttentionConfig( @@ -79,7 +81,8 @@ def definition(kind: str): if kind.startswith("static"): heads = 2 if kind == "static_head" else 1 bind_test_fp8_calibration( - state=state, scales=torch.ones((2, heads * 2), dtype=torch.float32), + state=state, + scales=torch.ones((2, heads * 2), dtype=torch.float32), q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, ) return state @@ -221,7 +224,9 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( tail = WindowCompressorState( config=WindowCompressorStateConfig( logical_layer_ids=(1,), - window_size=4, compress_stride=4, target="indexer", + window_size=4, + compress_stride=4, + target="indexer", candidate_dim=3, dtype=torch.bfloat16, ), @@ -238,7 +243,13 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( conv = reference.conv_state[:, 0, :, : config.conv_width] ssm = reference.ssm_state[:, 0] bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() - tail_bytes = len(tail.config.logical_layer_ids) * tail.config.window_size * 2 * tail.config.candidate_dim * tail.config.dtype.itemsize + tail_bytes = ( + len(tail.config.logical_layer_ids) + * tail.config.window_size + * 2 + * tail.config.candidate_dim + * tail.config.dtype.itemsize + ) manager = AttentionStateManager(request_capacity=3) manager.register_state("main", reference) manager.register_state("main", tail) @@ -279,7 +290,9 @@ def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( assert not pages[0].any() for state in (reference, tail): monkeypatch.setattr( - state, "get_pd_transfer_regions", lambda *args, **kwargs: pytest.fail("cached page size must not rebuild regions") + state, + "get_pd_transfer_regions", + lambda *args, **kwargs: pytest.fail("cached page size must not rebuild regions"), ) assert manager.get_pd_operator().get_page_bytes() == pages.shape[1] @@ -298,11 +311,16 @@ def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( ): if state_cls is MLAFullAttentionState: config = MLAConfig( - logical_layer_groups=((0,),), latent_dim=dim, latent_dtype=dtype, + logical_layer_groups=((0,),), + latent_dim=dim, + latent_dtype=dtype, ) else: config = GQAAttentionConfig( - logical_layer_groups=((0,),), global_head_num=1, head_dim=dim, dtype=dtype, + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=dim, + dtype=dtype, ) manager.register_state(role, state_cls(config, page_size=1, device="meta")) manager.finish_state_registration() @@ -359,12 +377,15 @@ def test_capacity_uses_quantized_field_padding_and_preserves_page_boundaries( ): if not torch.cuda.is_available(): pytest.skip("CUDA required") + class AlignedFieldsState(GroupQuantizedFullAttentionState): def _get_transfer_regions(self, page_tokens: int): regions, offset = [], 0 for region in super()._get_transfer_regions(page_tokens): offset = (offset + 7) // 8 * 8 - regions.append(replace(region, byte_start=offset, byte_end=offset + region.byte_end - region.byte_start)) + regions.append( + replace(region, byte_start=offset, byte_end=offset + region.byte_end - region.byte_start) + ) offset += prod(region.shape) * region.dtype.itemsize return tuple(regions) @@ -390,8 +411,16 @@ def _get_transfer_regions(self, page_tokens: int): for state in manager.get_full_attention_states(): assert state.get_token_bytes() * 3 == 792 monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("capacity must use field layouts")) - get_layouts = manager.get_cpu_cache_operator().get_page_layouts if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts - get_page_bytes = manager.get_cpu_cache_operator().get_page_bytes if transfer == "cpu" else manager.get_pd_operator().get_page_bytes + get_layouts = ( + manager.get_cpu_cache_operator().get_page_layouts + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts + ) + get_page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes + ) if layout_first: with monkeypatch.context() as patch: operator = manager.get_cpu_cache_operator() if transfer == "cpu" else manager.get_pd_operator() @@ -448,8 +477,8 @@ def _get_transfer_regions(self, page_tokens: int): torch.testing.assert_close(field.tensor[:, 4:7], field.tensor[:, :3]) assert (pages[[0, 2]] == 91).all() for offset in (0, 800): - assert (pages[1, offset + 780:offset + 784] == 91).all() - assert (pages[1, offset + 796:offset + 800] == 91).all() + assert (pages[1, offset + 780 : offset + 784] == 91).all() + assert (pages[1, offset + 796 : offset + 800] == 91).all() @pytest.mark.parametrize("full_page_tokens", [None, 8, 10, 16]) @@ -464,7 +493,9 @@ class PaddedGDNState(GDNState): @lru_cache(maxsize=None) def get_pd_transfer_regions(self): return tuple( - replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) if region.field == "ssm" else region + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + if region.field == "ssm" + else region for region in super().get_pd_transfer_regions() ) @@ -489,7 +520,8 @@ def get_pd_transfer_regions(self): conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ), - request_capacity=1, device="cuda", + request_capacity=1, + device="cuda", ) manager.register_state("main", state) manager.finish_state_registration() @@ -574,7 +606,11 @@ def _get_transfer_regions(self, page_tokens: int): ) manager.register_state("main", PaddedState(config, page_size=1, device="meta")) manager.finish_state_registration() - get_page_bytes = manager.get_cpu_cache_operator().get_page_bytes if transfer == "cpu" else manager.get_pd_operator().get_page_bytes + get_page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes + ) assert get_page_bytes() == 48 monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) @@ -586,7 +622,11 @@ def _get_transfer_regions(self, page_tokens: int): device="cpu", ) assert get_page_bytes() == (80 if transfer == "cpu" else 72) - layouts = manager.get_cpu_cache_operator().get_page_layouts() if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts() + layouts = ( + manager.get_cpu_cache_operator().get_page_layouts() + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts() + ) assert layouts["main_gqa_bf16_state"].byte_end == 72 @@ -623,7 +663,9 @@ def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( state = WindowCompressorState( config=WindowCompressorStateConfig( logical_layer_ids=(0,), - window_size=4, compress_stride=4, target="indexer", + window_size=4, + compress_stride=4, + target="indexer", candidate_dim=1, dtype=torch.float32, ), @@ -672,12 +714,15 @@ def test_cpu_page_bytes_derive_from_cached_transfer_layouts( q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None manager = AttentionStateManager(request_capacity=3) for role in ("main", "draft0"): - manager.register_state(role, rebuild_full_state( - template=template, - config=config, - scales=scales, - q_scales=q_scales, - )) + manager.register_state( + role, + rebuild_full_state( + template=template, + config=config, + scales=scales, + q_scales=q_scales, + ), + ) for state_cls, v_heads, value_dim in ((KDAState, 1, 1), (GDNState, 2, 3)): # 奇数长度的 BF16 Conv 与 FP32 SSM 之间需要补齐，MTP 候选区不应计入页大小。 bounded_config = ConvSSMStateConfig( @@ -725,7 +770,8 @@ class PaddedGDNState(GDNState): def get_cpu_transfer_regions(self): return tuple( replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) - if region.field == "ssm" else region + if region.field == "ssm" + else region for region in super().get_cpu_transfer_regions() ) @@ -741,7 +787,8 @@ def get_cpu_transfer_regions(self): conv_dtype=torch.bfloat16, ssm_dtype=torch.float32, ), - request_capacity=1, device="meta", + request_capacity=1, + device="meta", ) manager.register_state("main", state) manager.finish_state_registration() @@ -898,19 +945,23 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation dtype=torch.bfloat16, ) manager.register_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) - manager.register_state("main", GDNState( - ConvSSMStateConfig( - logical_layer_ids=(0,), - global_k_heads=1, - global_v_heads=1, - key_dim=1, - value_dim=1, - conv_width=1, - conv_dtype=torch.bfloat16, - ssm_dtype=torch.bfloat16, + manager.register_state( + "main", + GDNState( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.bfloat16, + ), + request_capacity=1, + device="cpu", ), - request_capacity=1, device="cpu", - )) + ) manager.finish_state_registration() if transfer == "cpu": operator = manager.get_cpu_cache_operator() @@ -940,7 +991,11 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation assert after_layout is not before_layout full_state = manager.main_gqa_bf16_state assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta - page_bytes = manager.get_cpu_cache_operator().get_page_bytes() if transfer == "cpu" else manager.get_pd_operator().get_page_bytes() + page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes() + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes() + ) if transfer == "pd_runtime": assert after_layout.byte_end == 8 and page_bytes == 40 elif transfer == "cpu": @@ -948,15 +1003,17 @@ def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation assert (layout_end + 15) // 16 * 16 == page_bytes else: assert after_layout.byte_end == page_bytes - assert not { - "cpu_page_tokens", "cpu_page_bytes", "cpu_transfer_layouts", "pd_page_bytes", "pd_transfer_layouts" - } & manager.__dict__.keys() + assert ( + not {"cpu_page_tokens", "cpu_page_bytes", "cpu_transfer_layouts", "pd_page_bytes", "pd_transfer_layouts"} + & manager.__dict__.keys() + ) @pytest.mark.parametrize("transfer", ["cpu", "pd"]) def test_manager_uses_state_regions_and_adds_each_resource_offset(transfer: str, pd_args: EasyDict): if not torch.cuda.is_available(): pytest.skip("CUDA required") + class ReversedCpuRegionsState(GQAFullAttentionState): @deepcopy_result @lru_cache(maxsize=None) @@ -964,7 +1021,8 @@ def get_cpu_transfer_regions(self): regions = super().get_cpu_transfer_regions() # 本实现的 CPU 页先保存 V 再保存 K，PD 仍沿用 K/V 顺序。 return tuple( - replace(region, byte_start=other.byte_start, byte_end=other.byte_end) for region, other in zip(regions, reversed(regions)) + replace(region, byte_start=other.byte_start, byte_end=other.byte_end) + for region, other in zip(regions, reversed(regions)) ) pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 @@ -983,7 +1041,11 @@ def get_cpu_transfer_regions(self): manager.register_state("draft0", prepared_full_state) manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) manager.finish_state_registration() - layouts = manager.get_cpu_cache_operator().get_page_layouts() if transfer == "cpu" else manager.get_pd_operator().get_pd_full_att_page_layouts() + layouts = ( + manager.get_cpu_cache_operator().get_page_layouts() + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts() + ) for index, state in enumerate(manager.get_full_attention_states()): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.zero_() @@ -997,7 +1059,11 @@ def get_cpu_transfer_regions(self): assert k_region.byte_start == index * state_bytes + k_offset assert v_region.byte_start == index * state_bytes + state_bytes // 2 - k_offset - page_bytes = manager.get_cpu_cache_operator().get_page_bytes() if transfer == "cpu" else manager.get_pd_operator().get_page_bytes() + page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes() + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes() + ) pages = torch.zeros( size=(1, page_bytes), dtype=torch.uint8, diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py index 355704a32b..badc2adf0c 100644 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -243,7 +243,11 @@ def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): dspark_markov_rank=4, ) args = SimpleNamespace( - llm_kv_type="None", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", run_mode="decode", + llm_kv_type="None", + page_size=256, + mtp_step=3, + mtp_mode="vanilla_with_att", + run_mode="decode", enable_prefill_microbatch_overlap=False, ) main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py index 84fa1f19d0..ae2ba27179 100644 --- a/unit_tests/common/att_resources/test_dsv4_state.py +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -412,9 +412,7 @@ def test_atomic_page_cpu_and_dp_roundtrip(kind: str) -> None: layout=layout, cpu_page_ready_flags=[True, False], ) - target.read_state_pages_from_cpu( - cpu_page_buffers=pages, cpu_page_indices=[2], big_page_slot_ids=[0], layout=layout - ) + target.read_state_pages_from_cpu(cpu_page_buffers=pages, cpu_page_indices=[2], big_page_slot_ids=[0], layout=layout) target.copy_big_pages_from_dp(src_page_ids=[0], dst_page_ids=[2], remote_dp_bounded_state=source) target.restore_bounded_att_runtime_from_pages( req_indices=[0], state_token_lens=[512], page_kind="big", page_slot_ids=[0] @@ -442,9 +440,7 @@ def test_manager_joint_cpu_page_roundtrip(monkeypatch: pytest.MonkeyPatch) -> No monkeypatch.setattr(args, "pd_kv_page_size", 256) manager = AttentionStateManager(request_capacity=3) full_states = [make_full(ratio=ratio) for ratio in (4, 128)] - bounded = [ - make_atomic_state(kind, mtp_step=2, device="cuda") for kind in ("swa", "kv_w8", "indexer_w8", "kv_w128") - ] + bounded = [make_atomic_state(kind, mtp_step=2, device="cuda") for kind in ("swa", "kv_w8", "indexer_w8", "kv_w128")] for state in (*full_states, *bounded): manager.register_state(model_role="main", state=state) manager.finish_state_registration() diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index 6bb3f32774..7d2a8019de 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -45,14 +45,16 @@ def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( scales = template.scales + value scales[1] += 10 q_scales = None if kind == "static_tensor" else template.q_scales + value - states.append(rebuild_full_state( - template=template, - config=replace_logical_layer_groups(template.config, ((3,), (8,))), - token_capacity=8, - device="cuda", - scales=scales, - q_scales=q_scales, - )) + states.append( + rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((3,), (8,))), + token_capacity=8, + device="cuda", + scales=scales, + q_scales=q_scales, + ) + ) manager = AttentionStateManager(request_capacity=2) manager.register_state("main", states[0]) manager.register_state("draft0", states[1]) @@ -141,12 +143,15 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 scale_count = 2 if scale_mode == "per_tensor" else 2 * head_num - scales = torch.arange( - start=1, - end=scale_count + 1, - dtype=torch.float32, - device="cuda", - ).unsqueeze(0) / 4 + scales = ( + torch.arange( + start=1, + end=scale_count + 1, + dtype=torch.float32, + device="cuda", + ).unsqueeze(0) + / 4 + ) state = StaticFP8FullAttentionState( layout_config=StaticFP8AttentionConfig( logical_layer_groups=((5,),), @@ -154,10 +159,12 @@ def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_nu head_dim=128, scale_mode=scale_mode, ), - page_size=4, device="cuda", + page_size=4, + device="cuda", ) bind_test_fp8_calibration( - state=state, scales=scales, + state=state, + scales=scales, q_scales=torch.ones((1, head_num), dtype=torch.float32, device="cuda") if scale_mode == "per_head" else None, ) state.allocate_storage(8) diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py index db871d24c0..8e2ee61214 100644 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -62,7 +62,8 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic with pytest.raises(AssertionError, match="latent 576"): create_packed_mla_state( config=PackedMLALogicalConfig( - logical_layer_groups=((0,),), latent_dim=512, + logical_layer_groups=((0,),), + latent_dim=512, format="packed_group128_fp8", ), device="meta", @@ -70,7 +71,8 @@ def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDic with pytest.raises(AssertionError, match="512-dimensional"): create_packed_mla_state( config=PackedMLALogicalConfig( - logical_layer_groups=((0,),), latent_dim=576, + logical_layer_groups=((0,),), + latent_dim=576, format="packed_group64_e8m0", ), device="meta", @@ -119,9 +121,7 @@ def test_fixed_quantizer_preserves_indexer_and_other_rows() -> None: indices = torch.tensor([1, 7, 12], dtype=torch.int32, device="cuda") kv = torch.randn((3, 1, 512), dtype=torch.bfloat16, device="cuda") expected = torch.empty((3, 1, 584), dtype=torch.uint8, device="cuda") - write_fp8_e8m0_packed_kv( - kv=kv, token_indices=torch.arange(3, dtype=torch.int32, device="cuda"), packed_kv=expected - ) + write_fp8_e8m0_packed_kv(kv=kv, token_indices=torch.arange(3, dtype=torch.int32, device="cuda"), packed_kv=expected) state.write_full_att_kv(logical_layer_id=4, token_indices=indices, kv=kv) view = state.get_layer_view(logical_layer_id=0) torch.testing.assert_close(view.kv[indices], expected, rtol=0, atol=0) @@ -187,9 +187,7 @@ def test_full_payload_roundtrip_skips_padding( source.write_to_pd_page(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout) target.read_from_pd_page(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, layout=layout) stream.synchronize() - torch.testing.assert_close( - target.storage[:, dst_ids, :, :716], source.storage[:, src_ids, :, :716], rtol=0, atol=0 - ) + torch.testing.assert_close(target.storage[:, dst_ids, :, :716], source.storage[:, src_ids, :, :716], rtol=0, atol=0) assert target.storage[..., 716:].eq(173).all() assert target.storage[:, 16:].eq(173).all() assert pages[:, :16].eq(99).all() and pages[:, -16:].eq(99).all() diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index ffddf2f375..0197b17daa 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -56,6 +56,7 @@ def __getitem__(self, grid: tuple[int]): def launch(**kwargs): launches.append((kwargs["HEADS"], kwargs["WIDTH"], kwargs["ELEMENT_BYTES"])) return kernel[grid](**kwargs) + return launch monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) @@ -95,9 +96,20 @@ def launch(**kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kind", [ - "gqa", "mla", "int4", "int8", "static_head", "static_tensor", "packed", "packed_fp8", "packed_glm", -]) +@pytest.mark.parametrize( + "kind", + [ + "gqa", + "mla", + "int4", + "int8", + "static_head", + "static_tensor", + "packed", + "packed_fp8", + "packed_glm", + ], +) @pytest.mark.parametrize("token_major", [False, True]) def test_direct_cpu_pages_preserve_bytes_without_gather(kind: str, token_major: bool, monkeypatch: pytest.MonkeyPatch): source = allocated_state(kind, device="cuda") @@ -172,7 +184,7 @@ def reject_index_contiguous(tensor: torch.Tensor, *args: object, **kwargs: objec assert (pages[:, :32] == 91).all() and (pages[:, -32:] == 91).all() assert (pages[[1, 3]] == 91).all() for region in layout.regions: - external = page_view[0, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + external = page_view[0, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) # 最后一页只有一个有效 token，其余槽位仍保留填充值。 unused = [slice(None)] * len(region.shape) unused[region.axes.index("token")] = slice(1, None) @@ -193,7 +205,9 @@ def test_direct_cpu_pages_only_update_selected_head_slice_and_skip_replica_write local_start=1, external_start=0, length=1, - ) if item.axis == "head" else item + ) + if item.axis == "head" + else item for item in region.axis_mappings ) regions.append(replace(region, axis_mappings=slices)) @@ -340,9 +354,20 @@ def test_gqa_single_cpu_page_and_empty_transfer(token_count: int) -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kind", [ - "gqa", "mla", "int4", "int8", "static_head", "static_tensor", "packed", "packed_fp8", "packed_glm", -]) +@pytest.mark.parametrize( + "kind", + [ + "gqa", + "mla", + "int4", + "int8", + "static_head", + "static_tensor", + "packed", + "packed_fp8", + "packed_glm", + ], +) def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: state = allocated_state(kind, device="cuda") layout = transfer_layout(state) diff --git a/unit_tests/common/att_resources/test_full_att_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py index 2a63449def..78bc924405 100644 --- a/unit_tests/common/att_resources/test_full_att_dp_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -39,6 +39,7 @@ def __getitem__(self, grid: tuple[int]): def launch(**kwargs): launches.append((kwargs["SHAPES"], kwargs["ELEMENT_BYTES"])) return kernel[grid](**kwargs) + return launch monkeypatch.setattr(dp_transfer, "_copy_indexed_token_storages", CaptureLaunch()) @@ -73,7 +74,8 @@ def test_dp_copy_preserves_complete_rows_without_temporary_payload( pytest.skip("CUDA peer access required") # 仅测试直接跨卡 Tensor 的场景需要预先开 peer；生产通过 IPC 装配，不在复制时建映射。 from cuda.bindings import runtime - error, = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + + (error,) = runtime.cudaDeviceEnablePeerAccess(source_device, 0) assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) template = definition(kind) @@ -142,10 +144,10 @@ def test_dp_copy_uses_real_storage_strides(column_stride: int, offset: int) -> N target_base = torch.empty(size=(2, 23, 3, 48), dtype=torch.bfloat16, device="cuda") source_base.view(torch.uint8).random_(0, 256) target_base.view(torch.uint8).fill_(173) - source = source_base[:, ::2, :, offset:offset + 8 * column_stride:column_stride] - target = target_base[:, ::3, :, offset:offset + 8 * column_stride:column_stride] + source = source_base[:, ::2, :, offset : offset + 8 * column_stride : column_stride] + target = target_base[:, ::3, :, offset : offset + 8 * column_stride : column_stride] expected = target_base.clone() - expected[:, ::3, :, offset:offset + 8 * column_stride:column_stride][:, [4, 0, 6]] = source[:, [7, 2, 7]] + expected[:, ::3, :, offset : offset + 8 * column_stride : column_stride][:, [4, 0, 6]] = source[:, [7, 2, 7]] copy_gqa_tokens( src_k_storage=source, src_v_storage=source, diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py index 059f67fdda..8c11087762 100644 --- a/unit_tests/common/att_resources/test_full_att_transfer.py +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -12,7 +12,11 @@ from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState from lightllm.common.attention_state.types import StateTransferLayout -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, full_state_test_fields, rebuild_full_state +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + rebuild_full_state, +) from unit_tests.common.att_resources.test_capacity import definition if TYPE_CHECKING: @@ -66,11 +70,14 @@ def test_full_att_pages_copy_all_fields_with_bounded_grid( regions = source.get_cpu_transfer_regions() if transfer == "cpu" else source.get_pd_transfer_regions() if grid_num == 1: # 调换层/token 轴，验证寻址使用真实 stride；region 顺序也不要求与存储成员相同。 - regions = tuple(replace( - region, - axes=("token", "layer", *region.axes[2:]), - shape=(region.shape[1], region.shape[0], *region.shape[2:]), - ) for region in reversed(regions)) + regions = tuple( + replace( + region, + axes=("token", "layer", *region.axes[2:]), + shape=(region.shape[1], region.shape[0], *region.shape[2:]), + ) + for region in reversed(regions) + ) layout = StateTransferLayout( state_name=source.state_name, tokens_per_page=page_tokens, @@ -152,7 +159,9 @@ def reject_gather(*args: object, **kwargs: object) -> None: assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() last_count = token_count - (page_count - 1) * page_tokens for region in regions: - last_page = page_view[page_indices[-1], region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + last_page = ( + page_view[page_indices[-1], region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + ) unused = [slice(None)] * 4 unused[region.axes.index("token")] = slice(last_count, None) assert (last_page[tuple(unused)].contiguous().view(torch.uint8) == 91).all() @@ -160,7 +169,8 @@ def reject_gather(*args: object, **kwargs: object) -> None: @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") def test_all_full_att_schemes_transfer_through_peer_ipc( - pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ) -> None: """各编码字段通过生产 IPC peer 映射跨 GPU 搬运，padding 保留目标值。""" from multiprocessing.reduction import ForkingPickler @@ -184,13 +194,15 @@ def test_all_full_att_schemes_transfer_through_peer_ipc( for name, field in source_fields.items(): expected_fields[name].tensor.copy_(other=field.tensor[:, 1:2]) regions = state.get_pd_transfer_regions() - layouts.append(StateTransferLayout( - state_name=state.state_name, - tokens_per_page=4, - byte_start=0, - byte_end=max(region.byte_end for region in regions), - regions=regions, - )) + layouts.append( + StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + ) states.append(state) expected_buffers.append(expected) torch.cuda.synchronize(device=1) @@ -219,7 +231,9 @@ def test_all_full_att_schemes_transfer_through_peer_ipc( def _peer_transfer_worker( - connection: Connection, payload: bytes, layouts: list[StateTransferLayout], + connection: Connection, + payload: bytes, + layouts: list[StateTransferLayout], ) -> None: from multiprocessing.reduction import ForkingPickler import traceback @@ -256,7 +270,9 @@ def _peer_transfer_worker( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kind", ["gqa", *SCHEMES]) -@pytest.mark.parametrize("readies", [[False, True, False], [True, False, True], [True, True, True], [False, False, False]]) +@pytest.mark.parametrize( + "readies", [[False, True, False], [True, False, True], [True, True, True], [False, False, False]] +) @pytest.mark.parametrize("grid_num", [1, 36]) def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( kind: str, diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 97105484be..1339ebb4b0 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -20,18 +20,25 @@ MLAConfig, ) from lightllm.common.attention_state.types import ( - AxisMapping, StateTransferLayout, TransferRegion, + AxisMapping, + StateTransferLayout, + TransferRegion, ) def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): if cls is MLAFullAttentionState: config = MLAConfig( - logical_layer_groups=((2,), (5,)), latent_dim=4, latent_dtype=torch.float32, + logical_layer_groups=((2,), (5,)), + latent_dim=4, + latent_dtype=torch.float32, ) else: config = GQAAttentionConfig( - logical_layer_groups=((2,), (5,)), global_head_num=1, head_dim=4, dtype=torch.float32, + logical_layer_groups=((2,), (5,)), + global_head_num=1, + head_dim=4, + dtype=torch.float32, ) prepared_full_state = cls(config, page_size=4, device=device) prepared_full_state.allocate_storage(12) @@ -103,6 +110,7 @@ def test_full_attention_write_copy_and_pd( head_end = head_start + buffer.shape[2] buffer[layer, ids.long()] = values[:, head_start:head_end] head_start = head_end + def reject_host_list(tensor: torch.Tensor): raise AssertionError("token copies must not convert indices to Python lists") @@ -197,7 +205,8 @@ def test_cpu_partial_last_page(state_offset: int, extra_page_bytes: int): ) assert torch.equal(read_kv_storage(target)[:, [6, 7, 8, 9, 10]], read_kv_storage(state)[:, [0, 1, 2, 3, 7]]) assert (buffer[[1, 3]] == 252).all() - assert (buffer[:, :layout.byte_start] == 252).all() and (buffer[:, layout.byte_end:] == 252).all() + assert (buffer[:, : layout.byte_start] == 252).all() and (buffer[:, layout.byte_end :] == 252).all() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_invalid_late_region_shape_cannot_partially_write_pd_page(): @@ -244,7 +253,8 @@ def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int) if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): target.copy_tokens_from_dp( - src_token_indices=torch.tensor([2, 0]), dst_token_indices=torch.tensor([6, 7]), + src_token_indices=torch.tensor([2, 0]), + dst_token_indices=torch.tensor([6, 7]), remote_dp_full_att_state=source, ) return diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index 29a1fa0265..26d4425ac7 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -26,12 +26,14 @@ from pathlib import Path -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("runtime"), +] @pytest.mark.parametrize("kind,encoding,model_cls", CASES) def test_model_storage_write_and_copy_preserve_all_fields( - runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, kind: str, @@ -75,7 +77,6 @@ def test_model_storage_write_and_copy_preserve_all_fields( @pytest.mark.parametrize("mode", ["normal", "prefill"]) @pytest.mark.parametrize("kind,encoding,model_cls", CASES[:8]) def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( - runtime: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mode: str, @@ -114,4 +115,6 @@ def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( assert state.capacity == 16 assert allocated_bytes(state) == storage_bytes if mode == "prefill": - assert model.att_state_manager.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() + assert ( + model.att_state_manager.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() + ) diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py index 1f64107157..2870ca126e 100644 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -55,20 +55,24 @@ def test_gqa_pd_direct_copy_between_tp_layouts( regions = [] for region in state.get_pd_transfer_regions(): axes = ("token", "layer", "head", "dim") if token_major else region.axes - regions.append(replace( - region, - byte_start=region.byte_start + 16, - byte_end=region.byte_end + 16, - axes=axes, - shape=tuple(region.shape[region.axes.index(axis)] for axis in axes), - )) - layouts.append(StateTransferLayout( - state_name=state.state_name, - tokens_per_page=4, - byte_start=16, - byte_end=max(region.byte_end for region in regions), - regions=tuple(regions), - )) + regions.append( + replace( + region, + byte_start=region.byte_start + 16, + byte_end=region.byte_end + 16, + axes=axes, + shape=tuple(region.shape[region.axes.index(axis)] for axis in axes), + ) + ) + layouts.append( + StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=16, + byte_end=max(region.byte_end for region in regions), + regions=tuple(regions), + ) + ) layout, target_layout = layouts pages = torch.full( size=(3, layout.byte_end + 64), @@ -79,7 +83,9 @@ def test_gqa_pd_direct_copy_between_tp_layouts( page_view = pages[:, 32:-32] # 页池视图的真实跨度大于可见宽度。 src_ids = torch.tensor(data=[6, 9, 1, 9, 7, 9], dtype=torch.int32, device="cuda")[::2][:token_count] dst_ids = torch.tensor(data=[2, 6, 4], dtype=torch.int64, device="cuda")[:token_count] - expected = [tensor[:, src_ids, 2:].contiguous().view(torch.uint8) for tensor in (source.k_storage, source.v_storage)] + expected = [ + tensor[:, src_ids, 2:].contiguous().view(torch.uint8) for tensor in (source.k_storage, source.v_storage) + ] def reject_gather(*args: object, **kwargs: object) -> None: pytest.fail("PD transfer must not gather or scatter intermediate KV tensors") @@ -106,11 +112,16 @@ def reject_gather(*args: object, **kwargs: object) -> None: stream.synchronize() for tensor, expected_bytes in zip((target.k_storage, target.v_storage), expected): assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected_bytes) - assert (tensor[:, [index for index in range(12) if index not in [2, 6, 4][:token_count]]].contiguous().view(torch.uint8) == 173).all() + assert ( + tensor[:, [index for index in range(12) if index not in [2, 6, 4][:token_count]]] + .contiguous() + .view(torch.uint8) + == 173 + ).all() assert (pages[[0, 2]] == 91).all() assert (pages[:, :48] == 91).all() and (pages[:, -32:] == 91).all() for region in layout.regions: - external = page_view[1, region.byte_start:region.byte_end].view(dtype).view(region.shape) + external = page_view[1, region.byte_start : region.byte_end].view(dtype).view(region.shape) unused = [slice(None)] * 4 unused[region.axes.index("token")] = slice(token_count, None) assert (external[tuple(unused)].contiguous().view(torch.uint8) == 91).all() @@ -124,7 +135,7 @@ def reject_gather(*args: object, **kwargs: object) -> None: layout=target_layout, ) for region, expected_bytes in zip(target_layout.regions, expected): - external = page_view[2, region.byte_start:region.byte_end].view(dtype).view(region.shape) + external = page_view[2, region.byte_start : region.byte_end].view(dtype).view(region.shape) external = external.permute(tuple(region.axes.index(axis) for axis in ("layer", "token", "head", "dim"))) assert torch.equal(external[:, :token_count, 2:].contiguous().view(torch.uint8), expected_bytes) assert (external[:, :, :2].contiguous().view(torch.uint8) == 91).all() @@ -161,7 +172,7 @@ def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: ) pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda") for region, value in zip(regions, (3, 7)): - pages[1, region.byte_start:region.byte_end].view(region.dtype).fill_(value) + pages[1, region.byte_start : region.byte_end].view(region.dtype).fill_(value) snapshot = pages.clone() ids = torch.tensor(data=[5, 1], dtype=torch.int32, device="cuda") state.write_to_pd_page( diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py index bb6e213a45..ac19dc39ab 100644 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -76,27 +76,32 @@ def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_a ) torch.cuda.synchronize() for src, dst in ((source.k_storage, target.k_storage), (source.v_storage, target.v_storage)): - assert torch.equal(dst[:, dst_ids].contiguous().view(torch.uint8), src[:, src_ids].contiguous().view(torch.uint8)) + assert torch.equal( + dst[:, dst_ids].contiguous().view(torch.uint8), src[:, src_ids].contiguous().view(torch.uint8) + ) assert (dst[:, dst_order[token_count:]].contiguous().view(torch.uint8) == 173).all() assert (dst[:, 128:].view(torch.uint8) == 173).all() assert (pages[[2, 5]] == 91).all() assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() for region in regions: - last_page = page_view[1, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + last_page = page_view[1, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) assert (last_page[:, 1:].contiguous().view(torch.uint8) == 91).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("page_device", ["cpu", "cuda"]) @pytest.mark.parametrize("num_stages", [1, 3]) -@pytest.mark.parametrize("byte_start, extra_bytes, dim_stride, element_bytes", [ - (8, 32, 1, 8), - (2, 32, 1, 2), - (4, 32, 1, 4), - (8, 34, 1, 2), - (8, 36, 1, 4), - (8, 32, 2, 2), -]) +@pytest.mark.parametrize( + "byte_start, extra_bytes, dim_stride, element_bytes", + [ + (8, 32, 1, 8), + (2, 32, 1, 2), + (4, 32, 1, 4), + (8, 34, 1, 2), + (8, 36, 1, 4), + (8, 32, 2, 2), + ], +) def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( page_device: str, num_stages: int, @@ -143,7 +148,7 @@ def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( device=page_device, pin_memory=page_device == "cpu", ) - page_view = pages[:, byte_start:byte_start + layout.byte_end] + page_view = pages[:, byte_start : byte_start + layout.byte_end] src_ids = torch.tensor(data=[3, 3, 1, 1, 7, 7], dtype=torch.int32, device="cuda")[::2] dst_ids = torch.tensor(data=[6, 2, 4], dtype=torch.int64, device="cuda") expected = {} @@ -160,6 +165,7 @@ def launch(**kwargs): if kwargs["WRITE"]: write_units.append(kwargs["ELEMENT_BYTES"]) return actual_kernel[grid](**kwargs) + return launch monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) @@ -180,7 +186,7 @@ def launch(**kwargs): assert torch.equal(getattr(target, name).view(torch.uint8), expected[name].view(torch.uint8)) assert (pages[1] == 91).all() assert (pages[:, :byte_start] == 91).all() - assert (pages[:, byte_start + layout.byte_end:] == 91).all() + assert (pages[:, byte_start + layout.byte_end :] == 91).all() for region in regions: - tail_page = page_view[0, region.byte_start:region.byte_end].view(region.dtype).view(region.shape) + tail_page = page_view[0, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) assert (tail_page[:, 1:].contiguous().view(torch.uint8) == 91).all() diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index d9ba1052ce..15d29bd85d 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -60,10 +60,15 @@ def test_hybrid_layer_mapping_and_mtp_states( config = Qwen3NextResourceConfig.from_model_config(model_config, args) logical_layer_ids = tuple(range(len(layer_types))) (full_state,) = Qwen3NextTpPartModel.create_full_att_states( - model_config=model_config, args=args, model_logical_layer_ids=logical_layer_ids, + model_config=model_config, + args=args, + model_logical_layer_ids=logical_layer_ids, ) (bounded_state,) = Qwen3NextTpPartModel.create_bounded_att_states( - model_config=model_config, args=args, model_logical_layer_ids=logical_layer_ids, device="meta", + model_config=model_config, + args=args, + model_logical_layer_ids=logical_layer_ids, + device="meta", ) manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py index 6576708961..0f55fe5986 100644 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -77,7 +77,8 @@ def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "E layer_types=["sliding_attention", "full_attention"] * 3, ) (state,) = Gemma4TpPartModel.create_full_att_states( - model_config=config, args=pd_args, + model_config=config, + args=pd_args, model_logical_layer_ids=(10, 11, 12, 13, 14, 15), ) # 共享消费者 14/15 分别使用最后一个同类型生产层 12/13，保持四份物理层存储。 @@ -90,7 +91,9 @@ def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "E @pytest.mark.parametrize("kind", SCHEMES) @pytest.mark.parametrize("transfer", ["cpu", "pd"]) def test_shared_layers_transfer_each_physical_layer_once( - kind: str, transfer: str, pd_args: "EasyDict", + kind: str, + transfer: str, + pd_args: "EasyDict", ) -> None: pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 template = definition(kind) @@ -121,7 +124,9 @@ def test_shared_layers_transfer_each_physical_layer_once( src_ids = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") dst_ids = torch.tensor([1, 4, 7], dtype=torch.int32, device="cuda") if transfer == "cpu": - pages = torch.zeros((2, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + pages = torch.zeros( + (2, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) managers[0].get_cpu_cache_operator().write_req( token_indices=src_ids, cpu_page_buffers=pages, diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 5ad569dd87..c28186809b 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -31,10 +31,13 @@ def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: global_head_num=1, head_dim=128, dtype=torch.bfloat16, - ), device="meta", + ), + device="meta", ) monkeypatch.setattr( - LlamaTpPartModel, "create_full_att_states", classmethod(lambda cls, **kwargs: [state]), + LlamaTpPartModel, + "create_full_att_states", + classmethod(lambda cls, **kwargs: [state]), ) model.prepare_attention_states() @@ -49,14 +52,18 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo manager = AttentionStateManager(request_capacity=1) for model_role in ("main", "draft0"): for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): - manager.register_state(model_role, GQAFullAttentionState( - layout_config=GQAAttentionConfig( - logical_layer_groups=groups, - global_head_num=1, - head_dim=128, - dtype=dtype, - ), device="meta", - )) + manager.register_state( + model_role, + GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=groups, + global_head_num=1, + head_dim=128, + dtype=dtype, + ), + device="meta", + ), + ) model = TpPartBaseModel.__new__(TpPartBaseModel) model.is_mtp_draft_model = role != "main" model.mtp_previous_draft_models = [] @@ -90,7 +97,8 @@ def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> Non plain = definition("gqa") quantized = definition("int8") quantized = type(quantized)( - layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), device="meta", + layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), + device="meta", ) manager.register_state("main", plain) manager.register_state("main", quantized) @@ -152,7 +160,8 @@ def test_mla_selection_excludes_same_layer_tail() -> None: tail = make_window() template = definition("mla") state = type(template)( - layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), device="meta", + layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), + device="meta", ) manager.register_state("main", state) manager.register_state("main", tail) @@ -161,17 +170,19 @@ def test_mla_selection_excludes_same_layer_tail() -> None: assert not tail.contains_logical_layer(0) assert manager.get_state("main", "mla", 7) is state assert manager.get_state("main", "compressor", 7) is tail - assert manager.get_state( - model_role="main", - state_type_filter="", - logical_layer_id=7, - state_class=FullAttentionState, - ) is state + assert ( + manager.get_state( + model_role="main", + state_type_filter="", + logical_layer_id=7, + state_class=FullAttentionState, + ) + is state + ) with pytest.raises(AssertionError, match="expected exactly one"): manager.get_state("main", "", 7) - @pytest.mark.parametrize("role", ["draft1", "draft10"]) def test_query_matches_complete_model_role(role: str) -> None: manager = AttentionStateManager(request_capacity=1) @@ -186,7 +197,9 @@ def test_precision_variants_have_distinct_names_and_can_coexist(algorithm: str) from dataclasses import replace manager = AttentionStateManager(request_capacity=3) - template = make_window() if algorithm == "compressor" else make_state(cls=KDAState if algorithm == "kda" else GDNState) + template = ( + make_window() if algorithm == "compressor" else make_state(cls=KDAState if algorithm == "kda" else GDNState) + ) states = [] for dtype in (torch.bfloat16, torch.float32): config = replace(template.config, **({"dtype": dtype} if algorithm == "compressor" else {"conv_dtype": dtype})) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py index b0e6670ba6..990559a072 100644 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -38,7 +38,10 @@ def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): if kind != "att_state": layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] if kind != "kv": - layouts += [page_layout(state) for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)] + layouts += [ + page_layout(state) + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state) + ] offset = 0 result = {} @@ -61,8 +64,8 @@ def init_source(manager: AttentionStateManager): for i, state in enumerate((manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)): for storage_buffer in full_state_test_buffers(state).values(): storage_buffer.copy_( - torch.arange(storage_buffer.numel(), device=storage_buffer.device).reshape_as(storage_buffer) + i - ) + torch.arange(storage_buffer.numel(), device=storage_buffer.device).reshape_as(storage_buffer) + i + ) for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): state.conv_state[:, 0].fill_(11 + i) state.conv_state[:, 1].fill_(23 + i) @@ -110,8 +113,6 @@ def fail(*args: object, **kwargs: object): assert manager.bounded_att_big_page_allocator.free_count == 3 - - @pytest.mark.parametrize("operation", ["save_bounded_att_runtime_to_pages", "restore_bounded_att_runtime_from_pages"]) def test_page_operations_forward_each_request_length(operation: str, monkeypatch: pytest.MonkeyPatch): manager = composition() @@ -197,10 +198,6 @@ def record_length( assert lengths == [(2, 29)] * len(manager.get_bounded_attention_states()) - - - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("token_count", [2, 4, 6, 8]) def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: int): @@ -225,7 +222,7 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i cpu_page_buffers=pages, cpu_page_indices=page_ids, layouts=layouts, - big_page_slot_ids=source_slots[:token_count // 4], + big_page_slot_ids=source_slots[: token_count // 4], small_page_slot_ids=small_slots, ) torch.cuda.synchronize() @@ -262,19 +259,13 @@ def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: i torch.testing.assert_close(dst.big_pages_pool.conv_state[2], src.big_pages_pool.conv_state[final_slot]) torch.testing.assert_close(dst.big_pages_pool.ssm_state[2], src.big_pages_pool.ssm_state[final_slot]) for src, dst in zip(source.get_full_attention_states(), target.get_full_attention_states()): - torch.testing.assert_close(read_kv_storage(dst)[:, 4:4 + token_count], read_kv_storage(src)[:, :token_count]) - - - - - - - - + torch.testing.assert_close(read_kv_storage(dst)[:, 4 : 4 + token_count], read_kv_storage(src)[:, :token_count]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("readies", [[True, False, False], [False, True, False], [False, False, True], [True, True, True]]) +@pytest.mark.parametrize( + "readies", [[True, False, False], [False, True, False], [False, False, True], [True, True, True]] +) def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bool]): source = composition("cuda") big_slots = init_source(source) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py index 7866525860..9544084bb6 100644 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -131,7 +131,9 @@ def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encod @pytest.mark.parametrize("tp_world_size", [1, 4]) def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, tp_world_size: int, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + tp_world_size: int, ): from builtins import open as original_open from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state @@ -140,14 +142,18 @@ def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( scales = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) q_scales = scales[:, :2] + 20 calibration_path = tmp_path / "calibration.json" - calibration_path.write_text(json.dumps(dict( - num_layers=4, - num_head=2, - quant_type="per_head", - scales_shape=[4, 4], - scales=scales.tolist(), - q_calibration=dict(num_head=2, scales_shape=[4, 2], scales=q_scales.tolist()), - ))) + calibration_path.write_text( + json.dumps( + dict( + num_layers=4, + num_head=2, + quant_type="per_head", + scales_shape=[4, 4], + scales=scales.tolist(), + q_calibration=dict(num_head=2, scales_shape=[4, 2], scales=q_scales.tolist()), + ) + ) + ) args = StartArgs(kv_quant_calibration_config_path=str(calibration_path)) reads = [] @@ -175,7 +181,9 @@ def count_open(file: str, *args: object, **kwargs: object): dtype=torch.bfloat16, tp_rank=tp_rank, tp_world_size=tp_world_size, - ), kv_type=encoding, device="meta", + ), + kv_type=encoding, + device="meta", ) manager.register_state(role, state) if encoding != "None": @@ -187,9 +195,9 @@ def count_open(file: str, *args: object, **kwargs: object): heads = 4 // tp_world_size local_heads = slice(tp_rank * heads, (tp_rank + 1) * heads) for state, row in expected: - expected_k = scales[row:row + 1, :2].repeat_interleave(2, dim=-1)[:, local_heads] - expected_v = scales[row:row + 1, 2:].repeat_interleave(2, dim=-1)[:, local_heads] - expected_q = q_scales[row:row + 1].repeat_interleave(2, dim=-1)[:, local_heads] + expected_k = scales[row : row + 1, :2].repeat_interleave(2, dim=-1)[:, local_heads] + expected_v = scales[row : row + 1, 2:].repeat_interleave(2, dim=-1)[:, local_heads] + expected_q = q_scales[row : row + 1].repeat_interleave(2, dim=-1)[:, local_heads] assert torch.equal(state.scales, torch.cat((expected_k, expected_v), dim=-1)) assert torch.equal(state.q_scales, expected_q) assert state.k_storage is None @@ -234,7 +242,9 @@ def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): global_head_num=2, head_dim=16, dtype=torch.bfloat16, - ), kv_type=encoding, device="cpu", + ), + kv_type=encoding, + device="cpu", ) manager = AttentionStateManager(request_capacity=1) manager.register_state("main", state) diff --git a/unit_tests/common/att_resources/test_model_config_normalization.py b/unit_tests/common/att_resources/test_model_config_normalization.py index f7dfa70a73..e3658f3cbe 100644 --- a/unit_tests/common/att_resources/test_model_config_normalization.py +++ b/unit_tests/common/att_resources/test_model_config_normalization.py @@ -22,7 +22,9 @@ @pytest.mark.parametrize("model_class", [Qwen3DFlash2Model, Qwen3_5DFlashModel, Qwen3_5DSparkModel]) @pytest.mark.parametrize("nested", [False, True]) def test_dflash_runtime_and_cpu_state_use_overridden_config( - model_class: type[TpPartBaseModel], tmp_path: Path, nested: bool, + model_class: type[TpPartBaseModel], + tmp_path: Path, + nested: bool, ) -> None: text_config = { "num_hidden_layers": 32, @@ -63,7 +65,8 @@ def test_dflash_runtime_and_cpu_state_use_overridden_config( @pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel, Qwen3_5MTPModel]) def test_native_mtp_runtime_matches_normalized_config( - model_class: type[TpPartBaseModel], tmp_path: Path, + model_class: type[TpPartBaseModel], + tmp_path: Path, ) -> None: text_config = { "num_hidden_layers": 32, @@ -98,11 +101,16 @@ def test_native_mtp_runtime_matches_normalized_config( @pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel]) def test_native_mtp_weights_keep_checkpoint_layer_start( - model_class: type[TpPartBaseModel], monkeypatch: pytest.MonkeyPatch, + model_class: type[TpPartBaseModel], + monkeypatch: pytest.MonkeyPatch, ) -> None: - config = model_class.normalize_model_config({ - "num_hidden_layers": 32, "num_nextn_predict_layers": 1, "hidden_size": 1024, - }) + config = model_class.normalize_model_config( + { + "num_hidden_layers": 32, + "num_nextn_predict_layers": 1, + "hidden_size": 1024, + } + ) module = importlib.import_module(model_class.pre_and_post_weight_class.__module__) names = [] @@ -118,7 +126,9 @@ def record_weight(**kwargs: object) -> SimpleNamespace: monkeypatch.setattr(module, "RMSNormWeight", record_weight) quant_cfg = SimpleNamespace(get_quant_method=lambda layer, name: None) pre_post_weight = model_class.pre_and_post_weight_class( - data_type=torch.bfloat16, network_config=config, quant_cfg=quant_cfg, + data_type=torch.bfloat16, + network_config=config, + quant_cfg=quant_cfg, ) assert len(names) == 4 assert all(name.startswith("model.layers.32.") for name in names) @@ -127,9 +137,13 @@ def record_weight(**kwargs: object) -> SimpleNamespace: model.config = config model.data_type = torch.bfloat16 model.quant_cfg = quant_cfg - model.main_model = SimpleNamespace(pre_post_weight=SimpleNamespace( - wte_weight_=object(), lm_head_weight_=object(), final_norm_weight_=object(), - )) + model.main_model = SimpleNamespace( + pre_post_weight=SimpleNamespace( + wte_weight_=object(), + lm_head_weight_=object(), + final_norm_weight_=object(), + ) + ) model.pre_and_post_weight_class = lambda *args, **kwargs: pre_post_weight layers = [] diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py index 46c28789cf..a571817fae 100644 --- a/unit_tests/common/att_resources/test_model_resources.py +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -19,7 +19,11 @@ GQAFullAttentionState, GQAAttentionConfig, ) -from lightllm.common.attention_state.full_att import PackedMLAFullAttentionState, PackedMLAAttentionConfig, PackedMLALogicalConfig +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, +) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -66,7 +70,7 @@ def make_resources( latent_dim=512, ) ), - page_size=4, + page_size=4, ) full.allocate_storage(16) else: @@ -79,7 +83,7 @@ def make_resources( tp_rank=rank, tp_world_size=tp, ), - page_size=4, + page_size=4, ) full.allocate_storage(16) main, draft = create_full_attention_roles(full, 4) @@ -106,18 +110,24 @@ def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( mtp: int, cpu_args: EasyDict, ): - sources = [make_resources( - glm=glm, - tp=tp, - rank=rank, - mtp=mtp, - ) for rank in range(tp)] - targets = [make_resources( - glm=glm, - tp=tp, - rank=rank, - mtp=mtp, - ) for rank in range(tp)] + sources = [ + make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) + for rank in range(tp) + ] + targets = [ + make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) + for rank in range(tp) + ] expected = [] for rank, (resources, full) in enumerate(sources): cfg = main_linear_state(resources).config diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index e85b92aad2..ed6b891603 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -84,8 +84,13 @@ def args(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): ), ( Gemma4TpPartModel, - {"num_key_value_heads": 4, "head_dim": 256, "global_head_dim": 256, - "num_kv_shared_layers": 1, "layer_types": ["full_attention", "full_attention"]}, + { + "num_key_value_heads": 4, + "head_dim": 256, + "global_head_dim": 256, + "num_kv_shared_layers": 1, + "layer_types": ["full_attention", "full_attention"], + }, 4, 256, 1, @@ -133,13 +138,18 @@ def test_model_and_cpu_service_use_identical_storage_geometry( if storage_layers == 1: assert state.config.logical_layer_groups == ((0, 1),) assert resources.get_state("main", "gqa", 0) is resources.get_state("main", "gqa", 1) - assert state.get_layer_view(logical_layer_id=0)[0].data_ptr() == state.get_layer_view(logical_layer_id=1)[0].data_ptr() + assert ( + state.get_layer_view(logical_layer_id=0)[0].data_ptr() + == state.get_layer_view(logical_layer_id=1)[0].data_ptr() + ) assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) assert resources.full_att_allocator.size == 16 layout = resources.get_cpu_cache_operator().get_page_layouts()[state.state_name] assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) # 按物理层数计算传输页；共享消费者不重复增加 KV 或传输容量。 - expected_pd_bytes = storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() + expected_pd_bytes = ( + storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() + ) assert resources.get_pd_operator().get_page_bytes() == expected_pd_bytes fields = layout.regions assert all(region.shape[2] == global_heads for region in fields) @@ -202,7 +212,9 @@ def memory_info(device: str | torch.device | int): def test_cpu_layout_uses_all_model_factory_states( - args: StartArgs | EasyDict, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + args: StartArgs | EasyDict, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ): """同一模型返回多种方案和共享层时，CPU 页采用实际 State 列表。""" from lightllm.common.attention_state.full_att import FullAttentionState, GQAFullAttentionState @@ -227,13 +239,21 @@ def create_full_att_states( return [ GQAFullAttentionState( layout_config=dataclasses.replace( - template.config, logical_layer_groups=(model_logical_layer_ids[:2],), dtype=torch.bfloat16, - ), page_size=args.page_size, device="meta", + template.config, + logical_layer_groups=(model_logical_layer_ids[:2],), + dtype=torch.bfloat16, + ), + page_size=args.page_size, + device="meta", ), GQAFullAttentionState( layout_config=dataclasses.replace( - template.config, logical_layer_groups=(model_logical_layer_ids[2:],), dtype=torch.float16, - ), page_size=args.page_size, device="meta", + template.config, + logical_layer_groups=(model_logical_layer_ids[2:],), + dtype=torch.float16, + ), + page_size=args.page_size, + device="meta", ), ] diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index 908167e7cb..f66767b9f4 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -144,7 +144,9 @@ def _snapshot( result["conv"] = state.conv_state[:, req_idx, ..., cursor : cursor + state.config.conv_width].float().cpu().numpy() result["ssm"] = state.ssm_state[:, req_idx * 3 + cursor].cpu().numpy() tail = resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, req_idx] - result["tail"] = torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() + result["tail"] = ( + torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() + ) return result @@ -275,7 +277,9 @@ def _worker(connection: Connection, backend: str, sender: bool): pd_page_index=request.dst_page_index, dp_rank_in_node=0, tp_world_size=1, - token_indices=torch.tensor([8, 12, 10, 15, 11], dtype=torch.long, device=device) if kind == "kv" else torch.tensor([], dtype=torch.long, device=device), + token_indices=torch.tensor([8, 12, 10, 15, 11], dtype=torch.long, device=device) + if kind == "kv" + else torch.tensor([], dtype=torch.long, device=device), task_kind=kind, req_idx=2, state_token_len=request.end_kv_index, diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index b279ea4ce1..032f7674eb 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -51,7 +51,8 @@ def test_page_pool_allocation_strategy( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_shared_registration_lives_until_last_view( - cpu_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ) -> None: cpu_args.enable_dp_prompt_cache_fetch = True unregistered = [] @@ -83,7 +84,8 @@ def test_gpu_page_pool_keeps_device_allocation(cpu_args: EasyDict) -> None: def test_empty_shared_pool_skips_cuda_registration( - cpu_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ) -> None: cpu_args.enable_dp_prompt_cache_fetch = True @@ -99,7 +101,8 @@ def unexpected_cuda_init() -> None: @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") @pytest.mark.parametrize("pool_cls", [KDAPagesPool, GDNPagesPool]) def test_pinned_page_pool_survives_queue_on_another_gpu( - cpu_args: EasyDict, pool_cls: type[KDAPagesPool] | type[GDNPagesPool], + cpu_args: EasyDict, + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], ) -> None: cpu_args.enable_dp_prompt_cache_fetch = True torch.cuda.set_device(0) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index 664db98fbb..af8652dd9e 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -20,11 +20,14 @@ GroupQuantizedFullAttentionState, StaticFP8AttentionConfig, StaticFP8FullAttentionState, - PackedMLAAttentionConfig, PackedMLALogicalConfig, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, PackedMLAFullAttentionState, ) from lightllm.common.attention_state.types import ( - AxisMapping, StateTransferLayout, TransferRegion, + AxisMapping, + StateTransferLayout, + TransferRegion, ) if TYPE_CHECKING: @@ -55,15 +58,19 @@ def make_quant_state(kind: str, device: str | torch.device | int): head_dim=16, scale_mode="per_head" if per_head else "per_tensor", ) - scales = torch.arange( - start=1, - end=9 if per_head else 5, - dtype=torch.float32, - device=device, - ).reshape(2, -1) / 16 + scales = ( + torch.arange( + start=1, + end=9 if per_head else 5, + dtype=torch.float32, + device=device, + ).reshape(2, -1) + / 16 + ) prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device) bind_test_fp8_calibration( - state=prepared_full_state, scales=scales, + state=prepared_full_state, + scales=scales, q_scales=torch.ones((2, config.head_num), dtype=torch.float32) if per_head else None, ) prepared_full_state.allocate_storage(8) @@ -102,7 +109,8 @@ def transfer_layout(state: FullAttentionState): external_start=0, length=dim, ) - for axis, dim in zip(field.axes, shape) if axis not in ("layer", "token") + for axis, dim in zip(field.axes, shape) + if axis not in ("layer", "token") ), ) ) @@ -134,7 +142,9 @@ def assert_token_fields_equal( @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( - kind: str, device: str | torch.device | int, monkeypatch: pytest.MonkeyPatch, + kind: str, + device: str | torch.device | int, + monkeypatch: pytest.MonkeyPatch, ): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") @@ -143,7 +153,9 @@ def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( field.tensor.fill_(i + 1) field.tensor[:, 0].fill_(i + 11) if device == "cuda": - source.copy_full_att_tokens(torch.tensor([0, 2], dtype=torch.long, device=device), torch.tensor([4, 5], dtype=torch.long, device=device)) + source.copy_full_att_tokens( + torch.tensor([0, 2], dtype=torch.long, device=device), torch.tensor([4, 5], dtype=torch.long, device=device) + ) assert_token_fields_equal( source=source, source_ids=[0, 2], @@ -235,7 +247,9 @@ def reject_index_conversion(tensor: torch.Tensor, *args, **kwargs): if device == "cpu": with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): target.copy_tokens_from_dp( - src_token_indices=src_indices, dst_token_indices=dst_indices, remote_dp_full_att_state=source, + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, ) return target.copy_tokens_from_dp( @@ -273,7 +287,9 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): raw = quant.to(torch.uint8) quant = torch.stack(((raw & 15).to(torch.int16) - 7, (raw >> 4).to(torch.int16) - 7), -1).flatten(-2) dequant = (quant.float() * torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids, ..., None]).reshape_as(kv) - assert torch.all((dequant - kv.float()).abs() <= torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids].amax() * 2) + assert torch.all( + (dequant - kv.float()).abs() <= torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids].amax() * 2 + ) elif isinstance(state, StaticFP8FullAttentionState): scales = state.scales[0] if state.config.scale_mode == "per_tensor": @@ -301,7 +317,9 @@ def test_cuda_compute_write_preserves_unrelated_fields(kind: str): def test_deferred_storage_preserves_real_fp8_calibration(kind: str): source = make_quant_state(kind, "cpu") state = StaticFP8FullAttentionState( - source.config, page_size=source.page_size, device="meta", + source.config, + page_size=source.page_size, + device="meta", ) assert state.scales is None and state.q_scales is None bind_test_fp8_calibration(state=state, scales=source.scales, q_scales=source.q_scales) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py index 658c340281..ca97cc7b70 100644 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -19,7 +19,8 @@ def test_quantized_kv_buffers_include_the_full_hold_page(): scale_dtype=torch.float16, bits=bits, ), - page_size=4, device="cpu", + page_size=4, + device="cpu", ) state.allocate_storage(32) allocations.extend(tuple(buffer.shape) for buffer in full_state_test_buffers(state).values()) diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py index f81210c842..88f02eeaf1 100644 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -18,7 +18,9 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("glm", [False, True]) def test_main_and_draft_transfers_copy_pd_and_dp_windows( - glm: bool, pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch, + glm: bool, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, ): pd_args.pd_kv_page_size = 128 source, _ = make_resources( @@ -55,14 +57,22 @@ def test_main_and_draft_transfers_copy_pd_and_dp_windows( pd_page_index=0, kind=kind, layouts=getattr(source.get_pd_operator(), getter)(), - **({"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 0, "state_token_len": 8}), + **( + {"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} + if kind == "kv" + else {"req_idx": 0, "state_token_len": 8} + ), ) target.get_pd_operator().read_page( pd_page_buffers=page.unsqueeze(0), pd_page_index=0, kind=kind, layouts=getattr(target.get_pd_operator(), getter)(), - **({"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} if kind == "kv" else {"req_idx": 1, "state_token_len": 8}), + **( + {"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} + if kind == "kv" + else {"req_idx": 1, "state_token_len": 8} + ), ) src_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") dst_indices = torch.tensor([10, 11], dtype=torch.int64, device="cuda") diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 5c34e35e28..1371cfe47e 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -102,7 +102,9 @@ def reject_host_list(tensor: torch.Tensor) -> list[int]: target_ids=[4, 6], ) - cpu_pages = torch.empty((1, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda") + cpu_pages = torch.empty( + (1, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda" + ) if device == "cpu": with pytest.raises(AssertionError, match="CUDA State storage"): managers[0].get_cpu_cache_operator().write_req( @@ -195,7 +197,7 @@ def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( # 生产拷贝入口只使用已映射的资源，不负责打开设备间访问。 from cuda.bindings import runtime - error, = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + (error,) = runtime.cudaDeviceEnablePeerAccess(source_device, 0) assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) for i, field in enumerate(full_state_test_fields(source_full, full_state_test_buffers(source_full)).values()): if "token" in field.axes: @@ -288,11 +290,14 @@ def barrier(group: str): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("p2p", [False, True]) -@pytest.mark.parametrize("share_pools,ipc_kinds", [ - (False, ("pd",)), - (True, ("pd", "dp")), - (True, ("dp", "pd")), -]) +@pytest.mark.parametrize( + "share_pools,ipc_kinds", + [ + (False, ("pd",)), + (True, ("pd", "dp")), + (True, ("dp", "pd")), + ], +) def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( monkeypatch: pytest.MonkeyPatch, p2p: bool, @@ -359,8 +364,8 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( assert error is None, error assert (read_kv_storage(full)[:, 4] == 9).all() window = resources.main_compressor_indexer_w4_s4_bf16_state - assert not window.buffer[:, 1, :, :window.config.candidate_dim].any() - assert torch.isneginf(window.buffer[:, 1, :, window.config.candidate_dim:]).all() + assert not window.buffer[:, 1, :, : window.config.candidate_dim].any() + assert torch.isneginf(window.buffer[:, 1, :, window.config.candidate_dim :]).all() assert (main_linear_state(resources).conv_state == 11).all() assert (main_linear_state(resources).big_pages_pool.conv_state[slot[0]] == 11).all() assert not main_linear_state(resources).big_pages_pool.ssm_state[slot[0]].any() @@ -386,18 +391,24 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( @pytest.mark.parametrize("source_tp,target_tp", [(1, 4), (4, 1), (2, 8), (8, 2)]) def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, target_tp: int, pd_args: EasyDict): pd_args.pd_kv_page_size = 128 - source = [make_resources( - glm=False, - tp=source_tp, - rank=rank, - mtp=2, - ) for rank in range(source_tp)] - target = [make_resources( - glm=False, - tp=target_tp, - rank=rank, - mtp=2, - ) for rank in range(target_tp)] + source = [ + make_resources( + glm=False, + tp=source_tp, + rank=rank, + mtp=2, + ) + for rank in range(source_tp) + ] + target = [ + make_resources( + glm=False, + tp=target_tp, + rank=rank, + mtp=2, + ) + for rank in range(target_tp) + ] for rank, (manager, full) in enumerate(source): for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): for local_head in range(full.config.head_num): @@ -493,7 +504,10 @@ def shared_memory(name: str, expected_size: int, force_mode: str) -> SimpleNames monkeypatch.setattr(shared_att_resources_utils.ForkingPickler, "dumps", serialize) shared_att_resources_utils.publish_att_resources( - resources=resources, request_mapping=mapping, kind=kind, reader_count=3, + resources=resources, + request_mapping=mapping, + kind=kind, + reader_count=3, ) assert len(payloads) == 3 and len({len(payload) for payload in payloads}) == 3 buffer = storage[f"rank_storage_{kind}_0"].buf @@ -538,7 +552,9 @@ def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: main_linear_state(resources).small_pages_pool.conv_state.flatten()[0] = 37 assert not hasattr(resources, "req_to_token_indexs") assert mapping[1, 0].item() == 7 - resources.copy_full_att_tokens(torch.tensor([0], dtype=torch.long, device="cuda"), torch.tensor([4], dtype=torch.long, device="cuda")) + resources.copy_full_att_tokens( + torch.tensor([0], dtype=torch.long, device="cuda"), torch.tensor([4], dtype=torch.long, device="cuda") + ) for window in window_states(resources): window.reset_bounded_att_state([1]) torch.cuda.synchronize() diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index 3339a71343..dee0d9a836 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -10,12 +10,24 @@ import torch from lightllm.common.attention_state.full_att import ( - GQAAttentionConfig, MLAConfig, FullAttentionState, GQAFullAttentionState, GroupQuantizedFullAttentionState, - StaticFP8FullAttentionState, MLAFullAttentionState, PackedMLAFullAttentionState, - PackedMLALogicalConfig, create_gqa_state, create_mla_state, create_packed_mla_state, + GQAAttentionConfig, + MLAConfig, + FullAttentionState, + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + StaticFP8FullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + PackedMLALogicalConfig, + create_gqa_state, + create_mla_state, + create_packed_mla_state, ) from lightllm.common.attention_state.bounded_att import ( - BoundedAttentionState, ConvSSMStateConfig, create_gdn_state, create_kda_state, + BoundedAttentionState, + ConvSSMStateConfig, + create_gdn_state, + create_kda_state, ) from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state @@ -108,8 +120,11 @@ def test_mla_creation_computes_packed_width_from_logical_dimensions( state = create_packed_mla_state(config=config, page_size=4, device="meta") else: config = MLAConfig( - logical_layer_groups=((2,), (6,)), latent_dim=kv_dim, latent_dtype=torch.bfloat16, - tp_rank=1, tp_world_size=4, + logical_layer_groups=((2,), (6,)), + latent_dim=kv_dim, + latent_dtype=torch.bfloat16, + tp_rank=1, + tp_world_size=4, ) state = create_mla_state(config=config, page_size=4, device="meta") assert state.get_state_type() == expected_type @@ -163,7 +178,9 @@ def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Calla def test_kpool_creation_preserves_pool_and_mtp_dimensions(): config = WindowCompressorStateConfig( logical_layer_ids=(3, 7), - window_size=4, compress_stride=4, target="indexer", + window_size=4, + compress_stride=4, + target="indexer", candidate_dim=128, dtype=torch.bfloat16, ) @@ -188,7 +205,9 @@ def test_create_rejects_encoding_from_another_logical_family(): with pytest.raises(AssertionError, match="unsupported GQA encoding"): create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") packed_config = PackedMLALogicalConfig( - logical_layer_groups=((0,),), latent_dim=576, format="packed_group128_fp8", + logical_layer_groups=((0,),), + latent_dim=576, + format="packed_group128_fp8", ) with pytest.raises(AssertionError, match="normal MLA creation"): create_mla_state(config=packed_config, device="meta") @@ -201,8 +220,13 @@ def test_create_rejects_encoding_from_another_logical_family(): @pytest.mark.parametrize( "state_class", - [GQAFullAttentionState, GroupQuantizedFullAttentionState, StaticFP8FullAttentionState, - MLAFullAttentionState, PackedMLAFullAttentionState], + [ + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + StaticFP8FullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + ], ) def test_full_attention_implementations_are_independent(state_class: type[FullAttentionState]): # 具体方案只共享抽象契约，不继承其他具体方案的存储或操作实现。 diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py index 5b383bba4e..8fffc04927 100644 --- a/unit_tests/common/att_resources/test_state_identity.py +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -92,7 +92,12 @@ def test_numbered_drafts_keep_state_name_order_and_unique_members(): manager.register_state("draft1", definition("gqa")) manager.finish_state_registration() assert [state.state_name.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ - "draft0", "draft10", "draft1", "draft2", "draft35", "main", + "draft0", + "draft10", + "draft1", + "draft2", + "draft35", + "main", ] assert manager.get_full_attention_states("draft1") == (manager.draft1_gqa_bf16_state,) assert manager.get_full_attention_states("draft10") == (manager.draft10_gqa_bf16_state,) @@ -211,7 +216,9 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp source.finish_state_registration() assert getattr(source, main.state_name) is main assert getattr(source, draft.state_name) is draft - source.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) + source.copy_full_att_tokens( + torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda") + ) target_physical = allocated_state(kind, device="cuda") target_main, target_draft = create_full_attention_roles(target_physical, 1) @@ -234,7 +241,9 @@ def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cp layouts=source.get_cpu_cache_operator().get_page_layouts(), ) for src, dst in ((main, target_main), (draft, target_draft)): - fields = {name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items()} + fields = { + name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items() + } for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): if "token" in field.axes: assert torch.equal(field.tensor[:, :2], fields[name][:, 6:8]) @@ -282,7 +291,9 @@ def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(c for state in (full, quant): manager.register_state("main", state) manager.finish_state_registration() - manager.copy_full_att_tokens(torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda")) + manager.copy_full_att_tokens( + torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda") + ) for state in (full, quant): for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): assert (field.tensor[:, 4:6] == 5).all() diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py index 7424d5be90..59d5751f69 100644 --- a/unit_tests/common/att_resources/test_state_queries.py +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -27,12 +27,16 @@ def test_created_state_queries_keep_the_same_objects_through_storage_initializat manager = AttentionStateManager(request_capacity=1) for role, kind in (("main", "gqa"), ("draft0", "int8")): manager.register_state(role, definition(kind)) - definitions.append(next( - state for state in vars(manager).values() - if isinstance(state, FullAttentionState) and state.state_name.startswith(role + "_") - )) + definitions.append( + next( + state + for state in vars(manager).values() + if isinstance(state, FullAttentionState) and state.state_name.startswith(role + "_") + ) + ) for query in ( - manager.get_full_attention_states, manager.get_bounded_attention_states, + manager.get_full_attention_states, + manager.get_bounded_attention_states, ): with pytest.raises(AssertionError, match="State registration"): query() @@ -72,7 +76,8 @@ def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): results = {(query.__name__, role): query(role) for query in queries for role in (None, "main", "draft0")} assert manager._get_states.cache_info().currsize == 6 assert manager.get_full_attention_states() == ( - manager.draft0_mla_fp32_state, manager.main_gqa_fp32_state, + manager.draft0_mla_fp32_state, + manager.main_gqa_fp32_state, ) assert manager.get_bounded_attention_states("main") == (manager.main_kda_conv_fp32_ssm_fp32_state,) assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_conv_fp32_ssm_fp32_state,) @@ -126,7 +131,6 @@ def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): assert getattr(exported, state.state_name) is state - @pytest.mark.parametrize("share_pools", [False, True]) def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share_pools: bool) -> None: cpu_args.enable_dp_prompt_cache_fetch = share_pools @@ -162,7 +166,8 @@ def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): assert imported.bounded_att_small_page_allocator is None misses = source._get_states.cache_info().misses for query in ( - imported.get_full_attention_states, imported.get_bounded_attention_states, + imported.get_full_attention_states, + imported.get_bounded_attention_states, ): for state in query(): name = state.state_name diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py index aa26a9252f..1fd660fa8d 100644 --- a/unit_tests/common/att_resources/test_tp_config.py +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -61,7 +61,9 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( cpu_args.cpu_cache_token_page_size = 4 if state_cls is MLAFullAttentionState: config = MLAConfig( - logical_layer_groups=((0,),), latent_dim=4, latent_dtype=torch.float32, + logical_layer_groups=((0,),), + latent_dim=4, + latent_dtype=torch.float32, ) else: config = GQAAttentionConfig( @@ -75,7 +77,10 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( reference.allocate_storage(4) for storage_buffer in full_state_test_buffers(reference).values(): storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape(storage_buffer.shape)) - fields = {name: field.tensor for name, field in full_state_test_fields(reference, full_state_test_buffers(reference)).items()} + fields = { + name: field.tensor + for name, field in full_state_test_fields(reference, full_state_test_buffers(reference)).items() + } reference_layout = cpu_layout(reference) buffer = torch.zeros((1, reference_layout.byte_end), dtype=torch.uint8, pin_memory=True) @@ -97,7 +102,9 @@ def test_serialized_full_config_restores_cpu_pages_with_different_tp( ) for rank, (start, end) in enumerate(target_ranges): - target = state_cls(replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda") + target = state_cls( + replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda" + ) target.allocate_storage(4) layout = cpu_layout(target) assert layout.byte_end == reference_layout.byte_end @@ -199,14 +206,18 @@ def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( assert page_bytes == page_tokens * 192 assert manager.get_pd_operator().get_page_bytes() == page_bytes getter = ( - manager.get_pd_operator().get_pd_bounded_att_page_layouts if bounded else manager.get_pd_operator().get_pd_full_att_page_layouts + manager.get_pd_operator().get_pd_bounded_att_page_layouts + if bounded + else manager.get_pd_operator().get_pd_full_att_page_layouts ) layouts = getter() assert getter() is layouts assert set(layouts) == {source.state_name, other.state_name} assert all(layout.tokens_per_page == (None if bounded else page_tokens) for layout in layouts.values()) unused_getter = ( - manager.get_pd_operator().get_pd_full_att_page_layouts if bounded else manager.get_pd_operator().get_pd_bounded_att_page_layouts + manager.get_pd_operator().get_pd_full_att_page_layouts + if bounded + else manager.get_pd_operator().get_pd_bounded_att_page_layouts ) assert unused_getter() == {} assert manager.get_pd_operator().get_page_bytes() == page_bytes diff --git a/unit_tests/common/att_resources/test_transfer_views.py b/unit_tests/common/att_resources/test_transfer_views.py index 7626abd0da..e3fa612818 100644 --- a/unit_tests/common/att_resources/test_transfer_views.py +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -18,7 +18,9 @@ ], ) def test_merge_tail_preserves_copy_regions( - layout: str, expected_shape: tuple[int, ...], start_dim: int, + layout: str, + expected_shape: tuple[int, ...], + start_dim: int, ) -> None: source_shape = (2, 5, 3, 8) if layout == "protected_window" else (2, 3, 8) destination_shape = source_shape @@ -121,7 +123,8 @@ def test_merge_tail_singleton_axes_preserve_addressing( source = torch.arange(source_size).as_strided(size=shape, stride=source_strides) destination = torch.arange(destination_size).as_strided(size=shape, stride=destination_strides) source_view, destination_view = merge_contiguous_tail_dims( - copies=[(source, destination)], keepdim=keepdim, + copies=[(source, destination)], + keepdim=keepdim, )[0] assert source_view.shape == destination_view.shape == (kept_shape if keepdim else flat_shape) assert source_view.data_ptr() == source.data_ptr() @@ -153,9 +156,9 @@ def test_max_copy_dtype_preserves_data_and_aliases( source = torch.arange(32).to(dtype=dtype) destination = torch.full(size=(32,), fill_value=-1, dtype=dtype) expected = destination.clone() - source_view = source[src_offset:src_offset + width] - destination_view = destination[dst_offset:dst_offset + width] - expected[dst_offset:dst_offset + width].copy_(source_view) + source_view = source[src_offset : src_offset + width] + destination_view = destination[dst_offset : dst_offset + width] + expected[dst_offset : dst_offset + width].copy_(source_view) copies = [(source_view, destination_view)] converted = view_max_copy_dtype(copies=copies) @@ -175,7 +178,9 @@ def test_max_copy_dtype_preserves_data_and_aliases( [(1, 0, torch.bfloat16), (0, 2, torch.int32), (0, 0, torch.int64)], ) def test_max_copy_dtype_preserves_row_padding( - src_padding: int, dst_padding: int, expected_dtype: torch.dtype, + src_padding: int, + dst_padding: int, + expected_dtype: torch.dtype, ) -> None: source = torch.arange(2 * (8 + src_padding)).to(dtype=torch.bfloat16).view(2, 8 + src_padding) destination = torch.full(size=(2, 8 + dst_padding), fill_value=-1, dtype=torch.bfloat16) @@ -192,7 +197,7 @@ def test_max_copy_dtype_preserves_row_padding( def test_max_copy_dtype_falls_back_for_middle_batch_offset() -> None: source = torch.arange(64).to(dtype=torch.bfloat16) destination = torch.full(size=(64,), fill_value=-1, dtype=torch.bfloat16) - copies = [(source[start:start + 8], destination[end:end + 8]) for start, end in ((0, 0), (9, 16), (24, 32))] + copies = [(source[start : start + 8], destination[end : end + 8]) for start, end in ((0, 0), (9, 16), (24, 32))] assert view_max_copy_dtype(copies=copies) is copies @@ -201,6 +206,8 @@ def test_max_copy_dtype_keeps_strided_last_axis() -> None: destination = torch.empty(size=(16,), dtype=torch.float32) copies = [(source, destination)] assert view_max_copy_dtype(copies=copies) is copies + + def test_deepcopy_result_protects_cached_values_and_keeps_cache_controls() -> None: from functools import lru_cache from lightllm.common.attention_state.utils import deepcopy_result diff --git a/unit_tests/common/att_resources/test_window_compressor_state.py b/unit_tests/common/att_resources/test_window_compressor_state.py index ad907852a6..106743b126 100644 --- a/unit_tests/common/att_resources/test_window_compressor_state.py +++ b/unit_tests/common/att_resources/test_window_compressor_state.py @@ -122,15 +122,18 @@ def test_manager_allocates_page_tensors_only_for_snapshot_states( partial(FullAttKVAllocator, pin_memory=False), ) manager = AttentionStateManager(request_capacity=3) - manager.register_state("main", GQAFullAttentionState( - layout_config=GQAAttentionConfig( - logical_layer_groups=((3,),), - global_head_num=1, - head_dim=4, - dtype=torch.bfloat16, + manager.register_state( + "main", + GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((3,),), + global_head_num=1, + head_dim=4, + dtype=torch.bfloat16, + ), + device="meta", ), - device="meta", - )) + ) manager.register_state("main", state) manager.finish_state_registration() manager.allocate_storage( @@ -189,4 +192,4 @@ def test_pd_window_restores_different_mtp_rings( torch.testing.assert_close(target.buffer[:, 2], expected, rtol=0, atol=0) assert (target.buffer[:, :2] == 37).all() assert (pages[0] == 253).all() and (pages[2] == 253).all() - assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end:] == 253).all() + assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end :] == 253).all() diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 356cb9a112..8ec77c6f4d 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -362,12 +362,20 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st request_capacity = req_to_mtp_state_index.numel() resources = AttentionStateManager(request_capacity=request_capacity) config = ConvSSMStateConfig( - logical_layer_ids=(3,), global_k_heads=1, global_v_heads=1, - key_dim=2, value_dim=2, conv_width=2, conv_dtype=torch.float32, ssm_dtype=torch.float32, + logical_layer_ids=(3,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, ) for model_role in ("main", "draft0"): linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") - linear.mtp_state_index = req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) + linear.mtp_state_index = ( + req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) + ) resources.register_state(model_role, linear) resources.finish_state_registration() infer_state.att_state_manager = resources diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index 6fb3770dfb..ae848f13de 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -65,7 +65,8 @@ def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.Monke att_state_role="main", ) monkeypatch.setattr( - backend_class, "get_gpu_workspace_buffer", + backend_class, + "get_gpu_workspace_buffer", lambda key_name, workspace_size, dtype: torch.empty(workspace_size, dtype=dtype), ) workspace, capacity = backend_class.get_swa_prefill_workspace(model=model) diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index a861c945fb..8a981bc726 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -71,7 +71,8 @@ def attend(**kwargs: object) -> torch.Tensor: state_class = fa3_fp8.Fp8Fa3PrefillAttState if operation == "prefill" else fa3_fp8.Fp8Fa3DecodeAttState else: state_class = ( - flashinfer_fp8.Fp8FlashInferPrefillAttState if operation == "prefill" + flashinfer_fp8.Fp8FlashInferPrefillAttState + if operation == "prefill" else flashinfer_fp8.Fp8FlashInferDecodeAttState ) @@ -84,6 +85,7 @@ def attend(**kwargs: object) -> torch.Tensor: attention.decode_max_q_seq_len = 1 attention.causal = True else: + def run( q: torch.Tensor, paged_kv_cache: tuple[torch.Tensor, torch.Tensor], diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py index a38e67bb5f..35caf326e3 100644 --- a/unit_tests/common/dsv4_test_utils.py +++ b/unit_tests/common/dsv4_test_utils.py @@ -276,7 +276,9 @@ def alloc_paged_kv_move_buffer(self, count: int, page_size: int): self.manager.get_pd_operator().get_pd_full_att_page_layouts.cache_clear() self.manager.get_pd_operator().get_pd_bounded_att_page_layouts.cache_clear() self.manager.get_pd_operator().get_page_bytes.cache_clear() - self.kv_move_buffer = torch.empty((count, self.manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + self.kv_move_buffer = torch.empty( + (count, self.manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda" + ) return self.kv_move_buffer def get_pd_kv_move_buffer_size(self): @@ -333,9 +335,7 @@ def read_page_kv_move_buffer_to_mem( class DeepseekV4ReqManager(ReqManager): - def __init__( - self, count: int, length: int, memory: DeepseekV4MemoryManager | None, config: Any, mtp_step: int = 0 - ): + def __init__(self, count: int, length: int, memory: DeepseekV4MemoryManager | None, config: Any, mtp_step: int = 0): super().__init__(count, length) self.cache_config = config self.states = DeepseekV4TpPartModel.create_bounded_att_states( diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index 983198f835..1278e5e195 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -136,7 +136,9 @@ def test_cache_io_copies_individual_token_slots(cache): buffer.zero_() buffer[source] = torch.arange(1, 5, device="cuda", dtype=torch.uint8)[:, None, None] layouts = manager.manager.get_cpu_cache_operator().get_page_layouts() - pages = torch.empty((1, manager.manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + pages = torch.empty( + (1, manager.manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) for state in manager.full_states.values(): state.write_to_cpu(source, pages, [0], layout=layouts[state.state_name]) state.read_from_cpu(pages, [0], destination, layout=layouts[state.state_name]) @@ -159,16 +161,28 @@ def test_cache_io_copies_individual_token_slots(cache): def test_request_runtime_is_independent_of_compressed_kv_capacity(cache, monkeypatch): original, requests = cache runtime = requests.window - before = tuple(tensor.data_ptr() for tensor in ( - runtime.swa_buffer, runtime.c4_buffer, runtime.indexer_buffer, runtime.c128_buffer, - )) + before = tuple( + tensor.data_ptr() + for tensor in ( + runtime.swa_buffer, + runtime.c4_buffer, + runtime.indexer_buffer, + runtime.c128_buffer, + ) + ) # 扩展压缩 KV 容量不能重新创建或移动按请求数分配的运行态。 for state in original.full_states.values(): state.allocate_storage(token_capacity=16384) assert state.storage.shape[1] == 16384 + state.page_size - assert before == tuple(tensor.data_ptr() for tensor in ( - runtime.swa_buffer, runtime.c4_buffer, runtime.indexer_buffer, runtime.c128_buffer, - )) + assert before == tuple( + tensor.data_ptr() + for tensor in ( + runtime.swa_buffer, + runtime.c4_buffer, + runtime.indexer_buffer, + runtime.c128_buffer, + ) + ) assert runtime.request_capacity == requests.max_request_num + 1 @@ -242,7 +256,10 @@ def _test_backend(requests, token_capacity=8192): "num_attention_heads": 64, "compress_ratios": requests.cache_config.compress_ratios, "index_topk": 512, - "head_dim": 512, "qk_rope_head_dim": 64, "index_head_dim": 128, "sliding_window": requests.sliding_window, + "head_dim": 512, + "qk_rope_head_dim": 64, + "index_head_dim": 128, + "sliding_window": requests.sliding_window, } model.tp_world_size_ = 1 model.max_req_num = requests.max_request_num @@ -268,8 +285,10 @@ def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): state.prepare_swa_workspace( SimpleNamespace( req_manager=requests, - swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, b_req_idx=req_ids, b_ready_cache_len=starts, b_seq_len=ends, @@ -668,7 +687,9 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): for layer in range(requests.cache_config.layer_num): torch.testing.assert_close( _read_swa(requests, layer, destination, positions, prefill=False), - _read_swa(prefill, layer, source, positions, prefill=False), rtol=0, atol=0, + _read_swa(prefill, layer, source, positions, prefill=False), + rtol=0, + atol=0, ) ring_size = requests.swa_capacity_tokens assert requests.req_to_swa.flatten(1, 2)[:, : destination * ring_size].eq(0).all() @@ -679,8 +700,12 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): (prefill.req_to_c128_state, requests.req_to_c128_state, 128), ): positions = torch.arange(max(0, length - width), length, device="cuda") - torch.testing.assert_close(source_buffer[:, source, positions % source_buffer.shape[2]], - target_buffer[:, destination, positions % target_buffer.shape[2]], rtol=0, atol=0) + torch.testing.assert_close( + source_buffer[:, source, positions % source_buffer.shape[2]], + target_buffer[:, destination, positions % target_buffer.shape[2]], + rtol=0, + atol=0, + ) prefill.free([source], src) requests.free([destination], dst) assert manager.allocator.can_use_mem_size == manager.size @@ -746,9 +771,17 @@ def test_pd_transport_completion_covers_peer_gpu_state(cache, length): with torch.cuda.stream(consumer_stream): positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda:1") for layer in range(requests.cache_config.layer_num): - torch.testing.assert_close(_read_swa(requests, layer, destination, positions, prefill=False), - _read_swa(requests, layer, source, positions, prefill=False), rtol=0, atol=0) - for buffer, width in ((requests.req_to_c4_state, 8), (requests.req_to_c4_indexer_state, 8), (requests.req_to_c128_state, 128)): + torch.testing.assert_close( + _read_swa(requests, layer, destination, positions, prefill=False), + _read_swa(requests, layer, source, positions, prefill=False), + rtol=0, + atol=0, + ) + for buffer, width in ( + (requests.req_to_c4_state, 8), + (requests.req_to_c4_indexer_state, 8), + (requests.req_to_c128_state, 128), + ): positions = torch.arange(max(0, length - width), length, device="cuda:1") % buffer.shape[2] torch.testing.assert_close(buffer[:, destination, positions], buffer[:, source, positions], rtol=0, atol=0) for buffer in peer.kv_buffer: @@ -778,8 +811,10 @@ def compress(req_idx, slots, start, end, decode=False): state = SimpleNamespace( mem_manager=manager, req_manager=requests, - swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, mem_index=slots[start:end], position_ids=positions, is_prefill=not decode, @@ -946,8 +981,14 @@ def build(): "Model", (), { - "config": dict(num_attention_heads=64, compress_ratios=[0, 0, 0], head_dim=512, - qk_rope_head_dim=64, index_head_dim=128, sliding_window=128), + "config": dict( + num_attention_heads=64, + compress_ratios=[0, 0, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + ), "layers_infer": [SimpleNamespace(layer_num_=2)], "tp_rank_": 0, "tp_world_size_": 1, @@ -988,8 +1029,10 @@ def forward(): reference_kv = backend.get_att_input_params( SimpleNamespace( req_manager=requests, - swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[source], ), @@ -1362,8 +1405,15 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, monkeypatch.setitem(args, "mtp_mode", "dspark") monkeypatch.setitem(args, "enable_decode_microbatch_overlap", True) monkeypatch.setitem(args, "enable_prefill_microbatch_overlap", False) - config = dict(sliding_window=128, index_topk=512, num_attention_heads=64, compress_ratios=[4, 128, 0], - head_dim=512, qk_rope_head_dim=64, index_head_dim=128) + config = dict( + sliding_window=128, + index_topk=512, + num_attention_heads=64, + compress_ratios=[4, 128, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + ) model = type("TestModel", (), {})() vars(model).update( config=config, @@ -1374,9 +1424,13 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, has_vision=False, args=args, req_manager=requests, - swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, - tp_rank_=0, tp_world_size_=1, att_state_role="main", + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, + tp_rank_=0, + tp_world_size_=1, + att_state_role="main", layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], ) backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) @@ -1551,7 +1605,10 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( disable_linear_att_small_page_cpu_cache=False, ) module.backend = SimpleNamespace( - is_master_in_dp=True, dp_world_size=1, is_hybrid_att_model=True, model=SimpleNamespace(att_state_manager=manager.manager, req_manager=requests) + is_master_in_dp=True, + dp_world_size=1, + is_hybrid_att_model=True, + model=SimpleNamespace(att_state_manager=manager.manager, req_manager=requests), ) module.cpu_cache_client = client module.cpu_cache_handle_queue = deque() @@ -1699,8 +1756,10 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): slots = _swa_slots(requests, req, positions) state = SimpleNamespace( req_manager=requests, - swa_state=requests.swa_state, get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[req], dsv4_swa_write_slots=slots, @@ -1744,16 +1803,22 @@ def test_cpu_cache_preserves_already_published_pages(cache): manager.big_page_buffers.buffer[state_id].fill_(7) checkpoint_ids = [state_id] manager.manager.get_cpu_cache_operator().write_req( - token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), cpu_page_ready_flags=readies, + token_indices=slots, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), + cpu_page_ready_flags=readies, big_page_slot_ids=checkpoint_ids, ) torch.cuda.synchronize() assert pages.eq(221).all() readies[0] = False manager.manager.get_cpu_cache_operator().write_req( - token_indices=slots, cpu_page_buffers=pages, cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), cpu_page_ready_flags=readies, + token_indices=slots, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), + cpu_page_ready_flags=readies, big_page_slot_ids=checkpoint_ids, ) torch.cuda.synchronize() diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 94e4ab5e91..530b95436a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -9,7 +9,8 @@ from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, - PackedMLAAttentionConfig, PackedMLALogicalConfig, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, ) from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -94,7 +95,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( latent_dim=512, ) ), - page_size=args.page_size, + page_size=args.page_size, ) full.allocate_storage(16) req = ReqManager(3, 16) @@ -167,8 +168,8 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( resources.reset_bounded_att_state([2]) assert not linear.conv_state[:, 2].any() assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() - assert not tail.buffer[:, 2, :, :tail.config.candidate_dim].any() - assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim:]).all() + assert not tail.buffer[:, 2, :, : tail.config.candidate_dim].any() + assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim :]).all() assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only # belong to live requests. Bytewise equality checks FP8 scale integrity. diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index 86eebff928..8b4e7aba69 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -1285,9 +1285,13 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.cpu_cache_operator.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.pd_operator.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args) + monkeypatch.setattr( + "lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args + ) config = PackedMLALogicalConfig( - logical_layer_groups=((3,),), latent_dim=512, format="packed_group64_e8m0", + logical_layer_groups=((3,),), + latent_dim=512, + format="packed_group64_e8m0", ) state = create_packed_mla_state(config=config, page_size=4) state.allocate_storage(token_capacity=16) @@ -1313,5 +1317,5 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): torch.cuda.synchronize() for begin in (4, 8): actual = state.get_layer_view(logical_layer_id=3) - torch.testing.assert_close(actual.kv[begin:begin + 4], actual.kv[:4], rtol=0, atol=0) - torch.testing.assert_close(actual.indexer_k[begin:begin + 4], actual.indexer_k[:4], rtol=0, atol=0) + torch.testing.assert_close(actual.kv[begin : begin + 4], actual.kv[:4], rtol=0, atol=0) + torch.testing.assert_close(actual.indexer_k[begin : begin + 4], actual.indexer_k[:4], rtol=0, atol=0) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index b46a87cde6..628547a591 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -242,7 +242,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -299,7 +299,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -391,7 +391,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2**31 // (584 * 2) + 1 + loc = 2 ** 31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -404,7 +404,7 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @@ -481,7 +481,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) out, final = chunk_gated_delta_rule( q, k, @@ -592,7 +592,8 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): att_state_manager=SimpleNamespace( get_state=lambda state_type_filter, **kwargs: ( SimpleNamespace(get_layer_view=lambda logical_layer_id: tail) - if "compressor_indexer" in state_type_filter else full_att_state + if "compressor_indexer" in state_type_filter + else full_att_state ), ), is_draft_model=False, @@ -645,11 +646,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) + scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 1440abd623..8ac9c2d473 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -172,9 +172,7 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} state.att_state_role = "main" state.att_state_manager = SimpleNamespace( - get_state=lambda **kwargs: SimpleNamespace( - get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv) - ) + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv)) ) state.prefill_att_state = Glm5NextSparsePrefillState() state.b1_cu_q_seq_len = torch.arange(4, device="cuda", dtype=torch.int32) @@ -420,6 +418,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): assert all(not hasattr(model, "att_state_manager") for model in (main, *drafts)) assert all(not hasattr(model, "linear_config") for model in (main, *drafts)) from contextlib import nullcontext + for model in (main, *drafts): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) main.prepare_attention_states() diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index fe8e149f3c..82e029c5dd 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -196,8 +196,14 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): ) resources = AttentionStateManager(request_capacity=4) config = ConvSSMStateConfig( - logical_layer_ids=(0,), global_k_heads=heads, global_v_heads=heads, - key_dim=dim, value_dim=dim, conv_width=3, conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, + logical_layer_ids=(0,), + global_k_heads=heads, + global_v_heads=heads, + key_dim=dim, + value_dim=dim, + conv_width=3, + conv_dtype=conv.dtype, + ssm_dtype=ssm.dtype, ) for model_role in ("main", "draft0"): linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 7ed6fb487b..ff93312516 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -12,7 +12,8 @@ from lightllm.common.attention_state.full_att import ( PackedMLAFullAttentionState, - PackedMLAAttentionConfig, PackedMLALogicalConfig, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, ) from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.models.glm5_next.config import Glm5NextResourceConfig @@ -78,7 +79,7 @@ def create(tp, target_mtp_step=None): latent_dim=512, ) ), - page_size=args.page_size, + page_size=args.page_size, ) full.allocate_storage(32) main, draft = create_full_attention_roles(full, config.all_layer_num) @@ -191,7 +192,10 @@ def test_pd_kv_and_runtime_state_roundtrip( source_buffer = allocate_transfer_buffer(source_resources[0], 1, device="cuda") dest_buffer = allocate_transfer_buffer(dest_resources[0], 1, device="cuda") assert source_buffer.shape == dest_buffer.shape - assert source_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() == dest_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() + assert ( + source_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() + == dest_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() + ) for kind in ("kv", "att_state"): copy_page( node_att_managers=source_resources, @@ -284,7 +288,9 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - full_bytes = max(layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_full_att_page_layouts().values()) + full_bytes = max( + layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_full_att_page_layouts().values() + ) runtime_bytes = max( layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_bounded_att_page_layouts().values() ) @@ -389,7 +395,7 @@ def prefill(values, destination, tail): width = target_step + 1 # 第一次验证后回访已接受位置；新增候选历史必须由 D 的本地 ring 正确保留。 for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: - raw[start:start + count] = torch.randn( + raw[start : start + count] = torch.randn( size=(count, 256), device="cuda", dtype=torch.bfloat16, @@ -401,7 +407,7 @@ def prefill(values, destination, tail): dtype=torch.int32, ) compress_pools( - raw=raw[start:start + count], + raw=raw[start : start + count], tail=target.get_layer_view(0), packed_buffer=packed, ape=ape, @@ -415,7 +421,7 @@ def prefill(values, destination, tail): mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) - prefill(raw[:start + count], expected, torch.zeros_like(source.get_layer_view(0))) + prefill(raw[: start + count], expected, torch.zeros_like(source.get_layer_view(0))) closing = torch.arange(start, start + count, device="cuda") closing = closing[(closing + 1) % 4 == 0] assert torch.equal(packed[closing], expected[closing]) diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py index 1e87b85a1b..ea3311e66a 100644 --- a/unit_tests/models/sensenova_loop/test_loop.py +++ b/unit_tests/models/sensenova_loop/test_loop.py @@ -38,6 +38,7 @@ def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch monkeypatch.setattr("lightllm.models.qwen3next.config.get_env_start_args", lambda: args) monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) cache = Qwen3NextResourceConfig.from_model_config(normalized, args) full_layers = cache.get_full_att_layer_indices() @@ -169,11 +170,15 @@ def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_lay normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) raw_cache = Qwen3NextResourceConfig.from_model_config(normalized, args) logical_ids = tuple(range(normalized["num_hidden_layers"])) - full, = SenseNovaLoopTpPartModel.create_full_att_states( - model_config=normalized, args=args, model_logical_layer_ids=logical_ids, + (full,) = SenseNovaLoopTpPartModel.create_full_att_states( + model_config=normalized, + args=args, + model_logical_layer_ids=logical_ids, ) - bounded, = SenseNovaLoopTpPartModel.create_bounded_att_states( - model_config=normalized, args=args, model_logical_layer_ids=logical_ids, + (bounded,) = SenseNovaLoopTpPartModel.create_bounded_att_states( + model_config=normalized, + args=args, + model_logical_layer_ids=logical_ids, ) full_layers = raw_cache.get_full_att_layer_indices() assert full.config.logical_layer_groups == tuple((i,) for i in full_layers) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py index 2b30c127b2..67976b131d 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -76,7 +76,9 @@ def test_fill_orders_radix_and_owned_pages_and_clears_reused_rows(monkeypatch: p @pytest.mark.parametrize("capacity,src_kv_len", [(2, 26), (1, 26), (0, 14)]) def test_build_reserves_only_new_big_pages_before_kv( - monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: int, + monkeypatch: pytest.MonkeyPatch, + capacity: int, + src_kv_len: int, ): module, req, allocator, allocations = _build_fixture(monkeypatch, capacity, src_kv_len) if capacity == 1: @@ -86,7 +88,7 @@ def test_build_reserves_only_new_big_pages_before_kv( assert req.hold_kv_len == 8 and allocator.free_count == 1 return tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) - task, = tasks + (task,) = tasks assert task.dst_req is req and task.src_req_idx == 1 assert task.src_dp_rank == 1 assert task.src_token_indices.tolist() == list(range(40, 32 + src_kv_len)) @@ -144,9 +146,7 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( module.big_page_tokens = 8 if hybrid else 0 module.shared_req_infos = SimpleNamespace(arr=np.zeros((1, 2, 2), dtype=np.int64)) module.shared_req_infos.arr[0, 1 - local_dp_rank] = [remote_len, 9] - module.shared_req_big_page_ids = ( - SimpleNamespace(arr=np.full((1, 2, 4), -1, dtype=np.int64)) if hybrid else None - ) + module.shared_req_big_page_ids = SimpleNamespace(arr=np.full((1, 2, 4), -1, dtype=np.int64)) if hybrid else None owned_pages = {length: 70 + length // 8 for length in range(8, local_len + 1, 8)} if hybrid else {} req = SimpleNamespace( req_idx=7, @@ -158,7 +158,9 @@ def test_local_longest_prefix_skips_copy_but_keeps_barriers( _kv_cache_alloc_need=allocation_need, shm_req=SimpleNamespace(shm_cur_kv_len=local_len), ) - monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=available_tokens)) + monkeypatch.setattr( + dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=available_tokens) + ) monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: barriers.append(group)) module.fill_reqs_info([req]) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py index 54bb5db743..17d564987d 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -132,7 +132,8 @@ def bind_before_listening(*args: object, **kwargs: object): "127.0.0.1", first_port, first_port + 4, - ) for _ in range(2) + ) + for _ in range(2) ] for future in futures: channels.append(future.result(timeout=15)) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py index 6d7018f561..726df8dbee 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -27,7 +27,9 @@ def test_pd_uploads_cpu_indices_on_copy_stream_and_skips_runtime_pages( original_cuda = torch.Tensor.cuda def track_upload(tensor: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: - uploads.append((tensor.device.type, tensor.is_pinned(), kwargs.get("non_blocking"), torch.cuda.current_stream())) + uploads.append( + (tensor.device.type, tensor.is_pinned(), kwargs.get("non_blocking"), torch.cuda.current_stream()) + ) return original_cuda(tensor, *args, **kwargs) def capture_copy(**kwargs: object) -> None: diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index e0c429419e..3c7ab1a6f5 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -328,7 +328,6 @@ def test_cpu_service_layout_and_roundtrip_all_encodings( assert not buffer[:2].any() - @pytest.mark.parametrize("hybrid", [False, True]) def test_production_offload_preserves_ready_pages_without_host_sync( runtime: None, @@ -410,6 +409,7 @@ def capture_write(*args: object, **kwargs: object) -> None: monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) monkeypatch.setattr(common.g_infer_context, "radix_cache", None) + def reject_host_sync(*args: object, **kwargs: object) -> None: pytest.fail("CPU offload must use stream dependencies instead of host synchronization") diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index d4861792c2..07ab123d09 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -71,7 +71,7 @@ def linear_cache(monkeypatch, tmp_path, request): head_dim=8, dtype=torch.bfloat16, ), - page_size=args.page_size, + page_size=args.page_size, ) full.allocate_storage(384) requests = ReqManager(2, 512) @@ -271,7 +271,9 @@ def test_offload_resolves_only_allocated_pages_and_requires_big_slots( full, resources, requests, config = linear_cache slots = resources.bounded_att_big_page_allocator.alloc(allocated_pages) if allocated_pages else [] shared_slots = [slots[0]] if shared_prefix and slots else [] - monkeypatch.setattr(g_infer_context, "radix_cache", SimpleNamespace(get_big_page_ids_by_node=lambda node: shared_slots)) + monkeypatch.setattr( + g_infer_context, "radix_cache", SimpleNamespace(get_big_page_ids_by_node=lambda node: shared_slots) + ) big_slots = {128 * (i + 1): slot for i, slot in enumerate(slots) if not shared_prefix or i > 0} req = SimpleNamespace( req_idx=0, @@ -298,7 +300,9 @@ def write_to_cpu( cpu_page_indices, **kwargs, ): - writes.append((token_indices.numel(), cpu_page_indices, kwargs["big_page_slot_ids"], kwargs["small_page_slot_ids"])) + writes.append( + (token_indices.numel(), cpu_page_indices, kwargs["big_page_slot_ids"], kwargs["small_page_slot_ids"]) + ) module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) module.args = get_env_start_args() @@ -311,7 +315,8 @@ def write_to_cpu( ) module.cpu_cache_client = SimpleNamespace( lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), - allocate_pages=allocate_pages, cpu_kv_cache_tensor=torch.empty((2, 1), dtype=torch.uint8), + allocate_pages=allocate_pages, + cpu_kv_cache_tensor=torch.empty((2, 1), dtype=torch.uint8), ) monkeypatch.setattr(resources.get_cpu_cache_operator(), "write_req", write_to_cpu) broadcasts = [] @@ -325,11 +330,15 @@ def write_to_cpu( assert writes == [] return task.sync_event.synchronize() - assert broadcasts == [[{ - "item_size": allocated_pages, - "page_list": list(range(allocated_pages)), - "ready_list": [False] * allocated_pages, - }]] + assert broadcasts == [ + [ + { + "item_size": allocated_pages, + "page_list": list(range(allocated_pages)), + "ready_list": [False] * allocated_pages, + } + ] + ] # 即使 CPU 只分配一页，也申请完整可缓存范围；未选中的 256 位置不要求提供槽位。 assert requested_hashes == [[11, 22]] assert writes == [(allocated_pages * 128, list(range(allocated_pages)), slots, [])] diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index f24d3dcdae..9779145cf4 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -67,9 +67,7 @@ def test_bounded_cache_miss_resets_runtime_and_preserves_pages(monkeypatch): req.req_idx = 1 req.cur_kv_len = req.cur_output_len = 0 req.sampling_param = SimpleNamespace(disable_prompt_cache=False) - req.shm_req = SimpleNamespace( - input_len=1, hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [], size=0) - ) + req.shm_req = SimpleNamespace(input_len=1, hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [], size=0)) req._hybrid_match_radix_cache() for state in resources.get_bounded_attention_states(): diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py index a263c9bed9..9626b03ea8 100644 --- a/unit_tests/server/router/test_kv_capacity.py +++ b/unit_tests/server/router/test_kv_capacity.py @@ -24,7 +24,9 @@ def test_router_reads_each_local_dp_groups_allocator( selected_ranks: list[int], reader_first: bool, ): - monkeypatch.setattr(full_att_kv_allocator, "get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=dp, nnodes=nnodes)) + monkeypatch.setattr( + full_att_kv_allocator, "get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=dp, nnodes=nnodes) + ) reader = full_att_kv_allocator.KVCapacityReader() if reader_first else None allocators = [] for rank in range(tp // nnodes): diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py index f41b5ad22a..3109fe11b9 100644 --- a/unit_tests/server/test_pd_start_args.py +++ b/unit_tests/server/test_pd_start_args.py @@ -95,9 +95,15 @@ def test_default_attention_factory_uses_gqa(): import torch from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState, create_gqa_state - state = create_gqa_state(config=GQAAttentionConfig( - logical_layer_groups=((0,),), global_head_num=4, head_dim=128, dtype=torch.bfloat16, - ), device="meta") + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=4, + head_dim=128, + dtype=torch.bfloat16, + ), + device="meta", + ) assert isinstance(state, GQAFullAttentionState) diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py index 88e695fac7..aa1c6d6273 100644 --- a/unit_tests/server/test_tokenizer.py +++ b/unit_tests/server/test_tokenizer.py @@ -40,9 +40,14 @@ def test_deepseek_bytelevel_checkpoint_preserves_whitespace(tmp_path: Path, mode config = json.loads(config_path.read_text()) config.update(tokenizer_class="LlamaTokenizerFast", legacy=True) config_path.write_text(json.dumps(config)) - (tmp_path / "config.json").write_text(json.dumps({ - "model_type": model_type, "architectures": ["DeepseekV3ForCausalLM"], - })) + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": model_type, + "architectures": ["DeepseekV3ForCausalLM"], + } + ) + ) tokenizer = get_tokenizer(str(tmp_path)) expected_ids = backend.encode(text).ids From a3ecc4b877e0d88dc43007020ebd698ce73ddd81 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:25:11 +0800 Subject: [PATCH 50/56] fix: handle absent draft attention states --- lightllm/common/basemodel/infer_struct.py | 9 ++++++--- .../att_resources/test_no_attention_draft.py | 3 +++ .../attention/test_state_lifetime.py | 20 +++++++++++++++++++ 3 files changed, 29 insertions(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 82a826fc54..83b312f991 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -135,11 +135,13 @@ def init_some_extra_state(self, model): def init_att_state(self): if self.is_prefill: - self.prefill_att_state.init_state() + if self.prefill_att_state is not None: + self.prefill_att_state.init_state() if self.prefill_att_state1 is not None: self.prefill_att_state1.init_state() else: - self.decode_att_state.init_state() + if self.decode_att_state is not None: + self.decode_att_state.init_state() if self.decode_att_state1 is not None: self.decode_att_state1.init_state() self.check_att_states() @@ -179,7 +181,8 @@ def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): attr_.copy_(attr_value, non_blocking=True) - self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) + if self.decode_att_state is not None: + self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) if self.decode_att_state1 is not None: self.decode_att_state1.copy_for_decode_cuda_graph(new_infer_state.decode_att_state1) return diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index 67a53fb694..2b27dbebab 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -66,6 +66,9 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo assert state.is_draft_model assert state.prefill_att_state is None and state.decode_att_state is None assert state.req_manager is draft.req_manager + state.init_att_state() + if not prefill: + state.copy_for_cuda_graph(draft._create_inferstate(model_input)) @pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py index aa1795058c..9e049d54b4 100644 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -152,3 +152,23 @@ def replay(self): assert runtime_ref() is None graph.graph.clear() assert captured_ref() is None + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_secondary_state_without_primary(is_prefill): + owner = InferStateInfo() + owner.is_prefill = is_prefill + owner.position_ids = torch.arange(4) + slot = "prefill_att_state1" if is_prefill else "decode_att_state1" + state_cls = _PrefillState if is_prefill else _DecodeState + secondary = state_cls(_infer_state_weak_ref=weakref.ref(owner)) + setattr(owner, slot, secondary) + owner.init_att_state() + torch.testing.assert_close(secondary.positions, owner.position_ids) + if not is_prefill: + runtime = InferStateInfo() + runtime.position_ids = owner.position_ids + 7 + runtime.decode_att_state1 = _DecodeState(_infer_state_weak_ref=weakref.ref(runtime)) + runtime.decode_att_state1.init_state() + owner.copy_for_cuda_graph(runtime) + torch.testing.assert_close(secondary.positions, runtime.position_ids) From 9aabee3298b60a7f08fdb2ec74ff7dbe2c1448cb Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:26:29 +0800 Subject: [PATCH 51/56] fix: allow separate INT4 KV storage in prefill --- .../basemodel/attention/triton/int4kv.py | 3 - .../test_quantized_full_state.py | 57 +++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py index e4dd2cee2f..22e723de3b 100644 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ b/lightllm/common/basemodel/attention/triton/int4kv.py @@ -70,9 +70,6 @@ def _groupsize_quant_prefill_att( # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) # batch_size = self.infer_state.b_seq_len.shape[0] - assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr() - assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr() - total_token_num = self.infer_state.total_token_num head_dim = k.shape[2] * 2 # 2个4bit存储为一个int8, 所以维度需要翻倍，才是解量化后的精度 k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index af8652dd9e..de13eb28fc 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -329,3 +329,60 @@ def test_deferred_storage_preserves_real_fp8_calibration(kind: str): state.allocate_storage(source.capacity) assert torch.equal(state.scales, source.scales) assert torch.equal(torch.from_numpy(state.cpu_scales), source.scales) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_int4_prefill_with_separate_storage(device, monkeypatch): + import weakref + from types import SimpleNamespace + from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv + from lightllm.common.basemodel.triton_kernel.att.prefill_att import context_flashattention_nopad + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = GroupQuantizedFullAttentionState( + GroupQuantizedAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=64, + scale_dtype=torch.float32, + bits=4, + ), + device=device, + ) + state.allocate_storage(8) + (k, k_scale), (v, v_scale) = state.get_layer_view(0) + k.view(torch.uint8).fill_(0x88) + v.view(torch.uint8).fill_(0x99) + k_scale.fill_(0.25) + v_scale.fill_(0.5) + assert k.untyped_storage().data_ptr() != v.untyped_storage().data_ptr() + assert k_scale.untyped_storage().data_ptr() != v_scale.untyped_storage().data_ptr() + owner = InferStateInfo() + owner.total_token_num = owner.max_kv_seq_len = owner.max_q_seq_len = 2 + owner.b_seq_len = torch.tensor([2], dtype=torch.int32, device=device) + owner.b_req_idx = torch.tensor([0], dtype=torch.int32, device=device) + owner.b_ready_cache_len = owner.b_q_start_loc = torch.zeros(1, dtype=torch.int32, device=device) + owner.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, device=device).view(1, -1)) + prefill = Int4kvTritonPrefillAttState( + backend=SimpleNamespace(quant_group_size=8), _infer_state_weak_ref=weakref.ref(owner) + ) + prefill.init_state() + if device == "cpu": + + def dequantize(**kwargs): + assert kwargs["k"] is k and kwargs["v"] is v + assert kwargs["k_scale"] is k_scale and kwargs["v_scale"] is v_scale + kwargs["k_out"].fill_(0.25) + kwargs["v_out"].fill_(1) + + def attention(**kwargs): + kwargs["o"].copy_(kwargs["v"].expand(-1, 2, -1)) + + monkeypatch.setattr(ppl_int4kv_copy_kv, "dequantize_int4kv", dequantize) + monkeypatch.setattr(context_flashattention_nopad, "context_attention_fwd_contiguous_kv", attention) + q = torch.randn(2, 2, 64, device=device) + output = prefill.prefill_att(q, (k, k_scale), (v, v_scale)) + torch.testing.assert_close(output, torch.ones_like(q)) From 2c4ba35e708a707dc955c924142fa0a18e21fb37 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:28:13 +0800 Subject: [PATCH 52/56] fix: preserve draft role when committing DFlash KV --- lightllm/models/qwen3_dflash/model.py | 1 + .../mtp_speculative/test_dflash.py | 64 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index dc7a0f685d..f0683dd882 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -92,6 +92,7 @@ def _decode(self, model_input: ModelInput) -> ModelOutput: infer_state.position_cos = torch.index_select(self._cos_cached, 0, position_ids) infer_state.position_sin = torch.index_select(self._sin_cached, 0, position_ids) infer_state.att_state_manager = self.att_state_manager + infer_state.att_state_role = self.att_state_role infer_state.is_draft_model = True infer_state.mem_index = self._select_mem_indexes(model_input) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py index ed42d27088..784848cb13 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import pytest import torch from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer @@ -114,3 +115,66 @@ def forward(model_input): proposal.schedule_scores, flat_draft_token_probs.reshape(2, block_size)[:, :2].float(), ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dflash_hidden_commit_uses_draft_state(device, monkeypatch): + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.full_att import GQAFullAttentionState, GQAAttentionConfig + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl + from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + manager = AttentionStateManager(request_capacity=1) + states = [] + for role, layers in (("main", ((0,), (1,), (2,), (3,))), ("draft0", ((4,),))): + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=layers, global_head_num=1, head_dim=16, dtype=torch.float32), + device=device, + ) + state.allocate_storage(8) + state.k_storage.zero_() + state.v_storage.zero_() + manager.register_state(role, state) + states.append(state) + manager.finish_state_registration() + main, draft = states + ids = torch.tensor([1, 3], dtype=torch.int32, device=device) + if device == "cpu": + + def write(logical_layer_id, token_indices, kv): + assert logical_layer_id == 4 + draft.k_storage[0, token_indices.long()] = kv[:, :1] + draft.v_storage[0, token_indices.long()] = kv[:, 1:] + + monkeypatch.setattr(draft, "write_full_att_kv", write) + model = Qwen3DFlashModel.__new__(Qwen3DFlashModel) + model.infer_state_class = InferStateInfo + model.att_state_manager = manager + model.is_mtp_draft_model = True + model.mtp_previous_draft_models = [] + model._cos_cached = model._sin_cached = torch.zeros(8, 16, device=device) + model._select_mem_indexes = lambda _: ids + model.pre_post_weight = None + model.pre_infer = SimpleNamespace(context_forward=lambda _, state, weight: state.mtp_draft_input_hiddens) + + def commit(hidden, state, weight): + TransformerLayerInferTpl._post_cache_kv(SimpleNamespace(layer_num_=4), hidden, state, weight) + return hidden + + model.layers_infer = [SimpleNamespace(context_forward=commit)] + model.trans_layers_weight = [None] + hidden = torch.arange(64, dtype=torch.float32, device=device).view(2, 2, 16) + output = model._decode( + SimpleNamespace( + mtp_draft_input_hiddens=hidden, + batch_size=2, + b_seq_len=torch.tensor([2, 4], device=device), + ) + ) + assert output.logits.shape == (2, 1) + torch.testing.assert_close(draft.k_storage[0, ids.long()], hidden[:, :1]) + torch.testing.assert_close(draft.v_storage[0, ids.long()], hidden[:, 1:]) + assert not main.k_storage.any() and not main.v_storage.any() From 0878642646988f3f4637c16d1998b24878ddaf1c Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:30:57 +0800 Subject: [PATCH 53/56] fix: evict radix cache before allocating DP state pages --- .../dp_backend/dp_shared_kv_trans.py | 7 ++- .../mode_backend/test_dp_big_pages.py | 59 ++++++++++++++++--- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index d88f5af2eb..03c69be293 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -111,11 +111,12 @@ def build_shared_kv_trans_tasks( assert all( page_id >= 0 for page_id in src_big_page_ids ), "DP source is missing prefix big pages" - allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator - dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) - assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" dst_token_indices = self.backend._alloc_req_kv_mem(req, alloc_token_num) assert dst_token_indices is not None + if src_big_page_ids: + allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator + dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) + assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" # 只搬运有效 KV；页尾预留槽位已写入请求表，不参与本次传输。 dst_token_indices = dst_token_indices[:trans_size] src_token_indices = self.backend.dp_req_to_token_indices[src_dp_rank][ diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py index 67976b131d..0f0fd7ae88 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -74,19 +74,13 @@ def test_fill_orders_radix_and_owned_pages_and_clears_reused_rows(monkeypatch: p assert module.shared_req_infos.arr[0, 0, 0] == 8 -@pytest.mark.parametrize("capacity,src_kv_len", [(2, 26), (1, 26), (0, 14)]) -def test_build_reserves_only_new_big_pages_before_kv( +@pytest.mark.parametrize("capacity,src_kv_len", [(2, 26), (0, 14)]) +def test_build_reserves_only_new_big_pages( monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: int, ): module, req, allocator, allocations = _build_fixture(monkeypatch, capacity, src_kv_len) - if capacity == 1: - with pytest.raises(AssertionError, match="DP fetch big page allocation must succeed"): - module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) - assert not allocations - assert req.hold_kv_len == 8 and allocator.free_count == 1 - return tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) (task,) = tasks assert task.dst_req is req and task.src_req_idx == 1 @@ -211,3 +205,52 @@ def allocate(cur_req: SimpleNamespace, count: int) -> torch.Tensor: monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=lambda: 100)) monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) return module, req, allocator, allocations + + +@pytest.mark.parametrize("prefix_len", [8, 16]) +def test_build_evicts_cached_state_pages(monkeypatch, kv_allocator_environment, prefix_len): + from sortedcontainers import SortedDict + from lightllm.common.attention_state.allocators import FullAttKVAllocator + from lightllm.server.router.dynamic_prompt import hybrid_att_radix_cache + from lightllm.server.router.model_infer.infer_batch import InferenceContext + from lightllm.server.router.model_infer.mode_backend import base_backend + + module, req, allocator, _ = _build_fixture(monkeypatch, 2, prefix_len) + req.cur_kv_len = req.hold_kv_len = 0 + req.hybrid_len_to_big_page_id = {} + module.shared_req_infos.arr[0, 0, 0] = 0 + full = FullAttKVAllocator(16, page_size=4, pin_memory=False) + manager = module.backend.model.att_state_manager + manager.full_att_allocator = full + manager.bounded_att_small_page_allocator = StatePageAllocator(0) + monkeypatch.setattr( + hybrid_att_radix_cache, "SharedArray", lambda name, shape, dtype: SimpleNamespace(arr=np.zeros(shape, dtype)) + ) + radix = hybrid_att_radix_cache.HybridAttPagedRadixCache(16, 0, 4, 2, manager) + slots = allocator.alloc(2) + for i, slot in enumerate(slots): + radix.insert( + key=torch.arange(8) + i * 8, + value=full.alloc(8).clone(), + block_hashs=[i * 2, i * 2 + 1], + block_state_idxs=[None, None], + len_to_big_page_id=SortedDict({8: slot}), + ) + requests = SimpleNamespace(req_to_token_indexs=torch.full((1, 32), -1, dtype=torch.int32)) + context = SimpleNamespace(att_state_manager=manager, radix_cache=radix, req_manager=requests) + context.get_can_alloc_token_num = lambda: InferenceContext.get_can_alloc_token_num(context) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(base_backend, "g_infer_context", context) + module.backend.args = SimpleNamespace(page_size=4) + module.backend._alloc_req_kv_mem = lambda req, count: base_backend.ModeBackend._alloc_req_kv_mem( + module.backend, req, count + ) + assert full.can_use_mem_size == allocator.free_count == 0 + assert context.get_can_alloc_token_num() == 16 + (task,) = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + assert len(task.dst_big_page_ids) == prefix_len // 8 + assert len(task.dst_token_indices) == req.hold_kv_len == prefix_len + assert radix.get_tree_total_tokens_num() == 16 - prefix_len + assert torch.equal(requests.req_to_token_indexs[0, :prefix_len], task.dst_token_indices) + assert allocator.free_count == 0 + assert full.can_use_mem_size == 0 From 245fc550c92ce3f91bee3914ca450c56307a29aa Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:32:00 +0800 Subject: [PATCH 54/56] fix: pass state lengths in CPU cache benchmark --- test/cpu_cache_kernel/test_speed.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 59e1b66ade..322d1b148d 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -96,7 +96,12 @@ def main(): # Populate every TP region through the production API before measuring restores. full.k_storage.zero_() full.v_storage.zero_() - resources.save_bounded_att_runtime_to_pages([0] * args.pages, page_kind="big", page_slot_ids=slots) + resources.save_bounded_att_runtime_to_pages( + [0] * args.pages, + state_token_lens=[(page + 1) * args.tokens_per_page for page in pages], + page_kind="big", + page_slot_ids=slots, + ) torch.cuda.synchronize() token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") resources.get_cpu_cache_operator().write_req( @@ -116,6 +121,7 @@ def restore(): layouts=resources.get_cpu_cache_operator().get_page_layouts(), big_page_slot_ids=slots, req_idx=0, + state_token_len=tokens, ) for _ in range(args.warmup): From aa3bef26ca8c61ef006b2c14a9b75e2765515b90 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:52:12 +0800 Subject: [PATCH 55/56] fix: rebind HOLD page after resuming KV memory --- .../mode_backend/rl_backend_ops.py | 3 +- .../mode_backend/test_rl_backend_ops.py | 44 ++++++++++++++----- 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py index a34adb8c52..839c3b10ad 100644 --- a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py +++ b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py @@ -73,7 +73,8 @@ def _resume_memory_tags(self, tags: Optional[List[MemoryTag]]): for tag in memory_tags: self.backend.model.torch_memory_saver.resume(tag=tag) if MemoryTag.KV_CACHE in memory_tags: - self.backend.model.req_manager.init_hold_request_indexs() + model = self.backend.model + model.req_manager.bind_hold_page(model.att_state_manager.full_att_allocator.hold_token_indices) def release_memory_occupation(self, tags: Optional[List[MemoryTag]]): try: diff --git a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py index dff0e892a9..7a631e398d 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py @@ -1,7 +1,12 @@ from types import SimpleNamespace +from unittest.mock import Mock import pytest +import torch +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.req_manager import ReqManager +from lightllm.common.req_manager.base import _ReqLinkedList from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps from lightllm.utils.torch_memory_saver_utils import MemoryTag @@ -10,24 +15,41 @@ ("tags", "should_init_hold_indexes"), [ ([MemoryTag.WEIGHT], False), + ([MemoryTag.GRAPH], False), + ([], False), ([MemoryTag.KV_CACHE], True), (None, True), ], ) -def test_resume_memory_initializes_hold_indexes_only_for_kv_cache(tags, should_init_hold_indexes): - events = [] +def test_resume_memory_initializes_hold_indexes_only_for_kv_cache( + tags, should_init_hold_indexes, kv_allocator_environment +): + allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) + req_manager = ReqManager.__new__(ReqManager) + req_manager.req_list = _ReqLinkedList(2) + req_manager.HOLD_REQUEST_ID = 2 + req_manager.page_size = allocator.page_size + req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) + resumed = [] + + def resume(tag): + assert req_manager.req_to_token_indexs.eq(-1).all() + resumed.append(tag) + model = SimpleNamespace( - torch_memory_saver=SimpleNamespace(resume=lambda tag: events.append(("resume", tag))), - req_manager=SimpleNamespace(init_hold_request_indexs=lambda: events.append(("init_hold_indexes", None))), + torch_memory_saver=SimpleNamespace(resume=resume), + req_manager=req_manager, + att_state_manager=SimpleNamespace(full_att_allocator=allocator), ) - backend = SimpleNamespace(model=model, logger=None) + backend = SimpleNamespace(model=model, logger=Mock()) ops = RlBackendOps(backend) ops._clear_cuda_cache = lambda: None - ops._resume_memory_tags(tags) + success, message = ops.resume_memory_occupation(tags) - init_events = [event for event in events if event[0] == "init_hold_indexes"] - assert bool(init_events) is should_init_hold_indexes - if should_init_hold_indexes: - last_resume_index = max(index for index, event in enumerate(events) if event[0] == "resume") - assert events.index(("init_hold_indexes", None)) > last_resume_index + assert success, message + assert resumed == list(RlBackendOps.MEMORY_TAG_ORDER if tags is None else tags) + expected = list(allocator.hold_token_indices) * 2 if should_init_hold_indexes else [-1] * 8 + assert req_manager.req_to_token_indexs[req_manager.HOLD_REQUEST_ID].tolist() == expected + assert req_manager.req_to_token_indexs[:2].eq(-1).all() + backend.logger.error.assert_not_called() From 4aa8075e11d137efec0b5a6ed6668ee159ad52eb Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:58:37 +0800 Subject: [PATCH 56/56] docs: move attention resource design into Chinese docs --- docs/{ => CN/source/framework}/cache_manager_design.md | 0 docs/CN/source/index.rst | 1 + 2 files changed, 1 insertion(+) rename docs/{ => CN/source/framework}/cache_manager_design.md (100%) diff --git a/docs/cache_manager_design.md b/docs/CN/source/framework/cache_manager_design.md similarity index 100% rename from docs/cache_manager_design.md rename to docs/CN/source/framework/cache_manager_design.md diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index a0399f1623..f15caff1bf 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -81,6 +81,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran :caption: 架构介绍 架构介绍 + Attention Resources 架构设计 token attention介绍 峰值显存调度器介绍